ums.agent.agent

  1# Agenten Plattform
  2#
  3# (c) 2024 Magnus Bender
  4# 	Institute of Humanities-Centered Artificial Intelligence (CHAI)
  5# 	Universitaet Hamburg
  6# 	https://www.chai.uni-hamburg.de/~bender
  7#  
  8# source code released under the terms of GNU Public License Version 3
  9# https://www.gnu.org/licenses/gpl-3.0.txt
 10
 11import random, os, json, time
 12
 13from abc import abstractmethod, ABC
 14from enum import Enum
 15from typing import List, Callable
 16
 17from pydantic import validate_call
 18
 19from ums.utils import (
 20	RiddleInformation, AgentMessage, RiddleDataType, RiddleData, Riddle,
 21	RiddleStatus, RiddleSolution,
 22	ExtractedData,
 23	logger
 24)
 25
 26class AgentCapability(Enum):
 27	"""
 28		The three different capabilities an agent can have.
 29	"""
 30
 31	EXTRACT="extract"
 32	SOLVE="solve"
 33	GATEKEEPER="gatekeeper"
 34
 35
 36class BasicAgent(ABC):
 37	"""
 38		A basic agent, each agent will be a subclass of this class.
 39	"""
 40
 41	@staticmethod
 42	@abstractmethod
 43	def agent_capability() -> AgentCapability:
 44		"""
 45			Represents the capabilities of this agent, for messages/ tasks of this capability, the `handle` method will be called.
 46		"""
 47		pass
 48
 49	@validate_call
 50	def __init__(self, message:AgentMessage, send_message:Callable[[AgentMessage], bool]):
 51		self._send_message = send_message
 52		self._sub_cnt = 0
 53
 54		self._message = message
 55		self._response = message.model_copy(deep=True)
 56		
 57		self._do_response = False
 58
 59		self._process()
 60		self._respond()
 61
 62	@abstractmethod
 63	def _process(self):
 64		pass
 65
 66	def _respond(self):
 67		# do a very short sleep
 68		time.sleep(random.random())
 69
 70		# sending
 71		send_it = lambda: self._send_message(self._response)
 72		if self.before_response(self._response, send_it) and self._do_response:
 73			send_it()
 74			logger.debug(f"Response sent {self._response.id}")
 75		else:
 76			logger.debug(f"Stopped response {self._response.id}")
 77
 78	@validate_call	
 79	def before_response(self, response:AgentMessage, send_it:Callable[[], None]) -> bool:
 80		"""
 81			This method is called before the response is sent.
 82			If the method returns `False` no response will be sent. 
 83			Thus, by overwriting this method, a response can be prevented.
 84
 85			The response to be sent is in `response` and `send_it` is a callable, which sends the response to the management if it gets called. 
 86			(Hence, one may stop sending the response and later call `send_it()` to send the response.)
 87		"""
 88		return True
 89	
 90	@validate_call
 91	def message(self) -> AgentMessage:
 92		"""
 93			Get the message this agent object is working on. 
 94		"""
 95		return self._message;
 96
 97	@validate_call
 98	def sub_riddle(self,
 99			riddle:Riddle, data:List[RiddleData]=[], status:RiddleStatus=None
100		) -> AgentMessage|bool:
101			"""
102				Create a new sub-riddle for solving details of the current riddle.
103				For the sub-riddle, give a `riddle` and optionally, a selection of `data` items (default none) and a `status` (default `ums.utils.types.RiddleStatus()`).
104				By changing a status, different steps can be (de-)selected.
105
106				Return the message of the sub-riddle or `false` on error.
107			"""
108
109			if status is None:
110				status = RiddleStatus()
111			
112			# new sub riddle id
113			self._sub_cnt += 1
114			new_id = "{}-sub-{}.{}".format(self._message.id, self._sub_cnt, int(random.random()*100))
115
116			self._message.sub_ids.append(new_id)
117			self._response.sub_ids.append(new_id)
118
119			# create the riddle's message
120			sub_msg = AgentMessage(
121				id=new_id,
122				riddle=riddle,
123				data=data,
124				status=status
125			)
126			logger.debug(f"Created sub-riddle {sub_msg.id}")
127
128			# send it
129			if self._send_message(sub_msg):
130				return sub_msg
131			else:
132				return False
133
134	@abstractmethod
135	def handle(self, *args:RiddleInformation) -> RiddleInformation:
136		"""
137			Handle a single task of the agent, the arguments and return value depends on the actual task (see subclass)!
138
139			**This is the method to implement!**
140
141			The full message is available via `message()`, a sub riddle can be created with `sub_riddle()`.
142		"""
143		pass
144
145	@validate_call
146	def get_extracted(self, data:RiddleData) -> ExtractedData|None:
147		"""
148			Loads the extracted data from the `data` item (i.e., from the file `data.file_extracted`).
149
150			Returns None if no extracted data found.
151		"""
152
153		if not data.file_extracted is None:
154			return ExtractedData.model_validate(
155				json.load(open(data.file_extracted, 'r'))
156			)
157			
158		return None
159
160class ExtractAgent(BasicAgent):
161	"""
162		An extraction agent.
163	"""
164
165	def agent_capability() -> AgentCapability:
166		return AgentCapability.EXTRACT
167	
168	@staticmethod
169	@abstractmethod
170	def extract_type() -> RiddleDataType:
171		"""
172			Represents the data this agent can process.
173		"""
174		pass
175
176	def _process(self):
177		for i, data in enumerate(self._response.data):
178			if data.type == self.__class__.extract_type():
179				logger.debug(f"Start extraction '{data.file_plain}'")
180				result = self.handle(data)
181				logger.debug(f"End extraction '{data.file_plain}' ('{result.file_extracted}')")
182
183				if result.file_extracted is None:
184					logger.info(f"Riddle {self._response.id}: 'file_extracted' for data '{data.file_plain}' still empty after handling")
185
186				self._response.data[i] = result
187				self._do_response = True
188
189		self._response.status.extract.finished = True
190
191	@abstractmethod
192	@validate_call
193	def handle(self, data:RiddleData) -> RiddleData:
194		"""
195			Process the item `data`, create extraction file and return `data` with populated `data.file_extracted`.
196		"""
197
198	@validate_call
199	def store_extracted(self, data:RiddleData, extracted:ExtractedData, allow_overwrite:bool=True) -> str:
200		"""
201			Stores the newly extracted data (in `extracted`) from `data` (i.e., `data.file_plain`) 
202			and returns the filename to use in `data.file_extracted`.
203
204			If there already exists an extracted file for this `data`, the file will be overwritten if `allow_overwrite=True`.
205			Generally the system will check, if the contents of the current file are equal to the contents to write. 
206			File with equal content will not be written again.
207		"""
208
209		# get path and name
210		path_name = data.file_plain[:data.file_plain.rfind('.')]
211		candidate = "{}.json".format(path_name)
212
213		# data to write
214		data = extracted.model_dump_json()
215
216		# check for file
217		if os.path.isfile(candidate):
218
219			# get current content
220			with open(candidate, 'r') as f:
221				content = f.read()
222
223			# files equal -> no need to rewrite
224			if content == data:
225				return candidate
226			
227			# not equal and overwrite not allowed
228			elif not allow_overwrite:
229				# get non-existent file name
230				cnt = 0
231				while os.path.isfile(candidate):
232					cnt += 1
233					candidate = "{}-{}.json".format(path_name, cnt)
234
235		# write file
236		with open(candidate, 'w+') as f:
237			f.write(data)
238
239		return candidate
240
241
242class ExtractTextAgent(ExtractAgent):
243	"""
244		An extraction agent for text, create a subclass for your agent.
245	"""
246	
247	def extract_type() -> RiddleDataType:
248		return RiddleDataType.TEXT
249
250class ExtractAudioAgent(ExtractAgent):
251	"""
252		An extraction agent for audio, create a subclass for your agent.
253	"""
254	
255	def extract_type() -> RiddleDataType:
256		return RiddleDataType.AUDIO
257
258class ExtractImageAgent(ExtractAgent):
259	"""
260		An extraction agent for images, create a subclass for your agent.
261	"""
262	
263	def extract_type() -> RiddleDataType:
264		return RiddleDataType.IMAGE
265
266
267class SolveAgent(BasicAgent):
268	"""
269		A solve agent, create a subclass for your agent.
270	"""
271
272	def agent_capability() -> AgentCapability:
273		return AgentCapability.SOLVE
274	
275	def _process(self):
276		logger.debug(f"Start solve: {self._response.id}")
277		solution = self.handle(self._response.riddle, self._response.data)
278		logger.debug(f"End solve: {self._response.id} ({solution.solution}, {solution.explanation})")
279		
280		if len(solution.solution) == 0 or len(solution.explanation) == 0:
281			logger.info(f"Riddle {self._response.id}: Empty solution/ explanation after handling")
282
283		self._response.solution = solution
284		self._response.status.solve.finished = True
285
286		self._do_response = True
287	
288	@abstractmethod
289	@validate_call
290	def handle(self, riddle: Riddle, data: List[RiddleData]) -> RiddleSolution:
291		"""
292			Solve the `riddle` using `data` and return a solution.
293		"""
294	
295class GatekeeperAgent(BasicAgent):
296	"""
297		A gatekeeper agent, create a subclass for your agent.
298	"""
299
300	def agent_capability() -> AgentCapability:
301		return AgentCapability.GATEKEEPER
302	
303	def _process(self):
304		if self._response.solution is None:
305			self._response.solution = RiddleSolution(solution="", explanation="")
306
307		logger.debug(f"Start validate: {self._response.id}")
308		solution = self.handle(self._response.solution, self._response.riddle)
309		logger.debug(f"End validate: {self._response.id} ({solution.review}, {solution.accepted})")
310		
311		if solution.review is None or len(solution.review) == 0:
312			logger.info(f"Riddle {self._response.id}: Empty review after handling")
313
314		self._response.solution = solution
315		self._response.status.validate.finished = True
316		self._response.status.solved = solution.accepted
317
318		self._do_response = True
319	
320	@abstractmethod
321	@validate_call
322	def handle(self, solution:RiddleSolution, riddle:Riddle) -> RiddleSolution:
323		"""
324			Check the `solution` of `riddle` and return solution with populated `solution.accepted` and `solution.review`.
325		"""