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
 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		send_it = lambda: self._send_message(self._response)
 68		if self.before_response(self._response, send_it) and self._do_response:
 69			send_it()
 70			logger.debug(f"Response sent {self._response.id}")
 71		else:
 72			logger.debug(f"Stopped response {self._response.id}")
 73
 74	@validate_call	
 75	def before_response(self, response:AgentMessage, send_it:Callable[[], None]) -> bool:
 76		"""
 77			This method is called before the response is sent.
 78			If the method returns `False` no response will be sent. 
 79			Thus, by overwriting this method, a response can be prevented.
 80
 81			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. 
 82			(Hence, one may stop sending the response and later call `send_it()` to send the response.)
 83		"""
 84		return True
 85	
 86	@validate_call
 87	def message(self) -> AgentMessage:
 88		"""
 89			Get the message this agent object is working on. 
 90		"""
 91		return self._message;
 92
 93	@validate_call
 94	def sub_riddle(self,
 95			riddle:Riddle, data:List[RiddleData]=[], status:RiddleStatus=None
 96		) -> AgentMessage|bool:
 97			"""
 98				Create a new sub-riddle for solving details of the current riddle.
 99				For the sub-riddle, give a `riddle` and optionally, a selection of `data` items (default none) and a `status` (default `ums.utils.types.RiddleStatus()`).
100				By changing a status, different steps can be (de-)selected.
101
102				Return the message of the sub-riddle or `false` on error.
103			"""
104
105			if status is None:
106				status = RiddleStatus()
107			
108			# new sub riddle id
109			self._sub_cnt += 1
110			new_id = "{}-sub-{}.{}".format(self._message.id, self._sub_cnt, int(random.random()*100))
111
112			self._message.sub_ids.append(new_id)
113			self._response.sub_ids.append(new_id)
114
115			# create the riddle's message
116			sub_msg = AgentMessage(
117				id=new_id,
118				riddle=riddle,
119				data=data,
120				status=status
121			)
122			logger.debug(f"Created sub-riddle {sub_msg.id}")
123
124			# send it
125			if self._send_message(sub_msg):
126				return sub_msg
127			else:
128				return False
129
130	@abstractmethod
131	def handle(self, *args:RiddleInformation) -> RiddleInformation:
132		"""
133			Handle a single task of the agent, the arguments and return value depends on the actual task (see subclass)!
134
135			**This is the method to implement!**
136
137			The full message is available via `message()`, a sub riddle can be created with `sub_riddle()`.
138		"""
139		pass
140
141	@validate_call
142	def get_extracted(self, data:RiddleData) -> ExtractedData|None:
143		"""
144			Loads the extracted data from the `data` item (i.e., from the file `data.file_extracted`).
145
146			Returns None if no extracted data found.
147		"""
148
149		if not data.file_extracted is None:
150			return ExtractedData.model_validate(
151				json.load(open(data.file_extracted, 'r'))
152			)
153			
154		return None
155
156class ExtractAgent(BasicAgent):
157	"""
158		An extraction agent.
159	"""
160
161	def agent_capability() -> AgentCapability:
162		return AgentCapability.EXTRACT
163	
164	@staticmethod
165	@abstractmethod
166	def extract_type() -> RiddleDataType:
167		"""
168			Represents the data this agent can process.
169		"""
170		pass
171
172	def _process(self):
173		for i, data in enumerate(self._response.data):
174			if data.type == self.__class__.extract_type():
175				logger.debug(f"Start extraction '{data.file_plain}'")
176				result = self.handle(data)
177				logger.debug(f"End extraction '{data.file_plain}' ('{result.file_extracted}')")
178
179				if result.file_extracted is None:
180					logger.info(f"Riddle {self._response.id}: 'file_extracted' for data '{data.file_plain}' still empty after handling")
181
182				self._response.data[i] = result
183				self._do_response = True
184
185		self._response.status.extract.finished = True
186
187	@abstractmethod
188	@validate_call
189	def handle(self, data:RiddleData) -> RiddleData:
190		"""
191			Process the item `data`, create extraction file and return `data` with populated `data.file_extracted`.
192		"""
193
194	@validate_call
195	def store_extracted(self, data:RiddleData, extracted:ExtractedData, allow_overwrite:bool=True) -> str:
196		"""
197			Stores the newly extracted data (in `extracted`) from `data` (i.e., `data.file_plain`) 
198			and returns the filename to use in `data.file_extracted`.
199
200			If there already exists an extracted file for this `data`, the file will be overwritten if `allow_overwrite=True`.
201			Generally the system will check, if the contents of the current file are equal to the contents to write. 
202			File with equal content will not be written again.
203		"""
204
205		# get path and name
206		path_name = data.file_plain[:data.file_plain.rfind('.')]
207		candidate = "{}.json".format(path_name)
208
209		# data to write
210		data = extracted.model_dump_json()
211
212		# check for file
213		if os.path.isfile(candidate):
214
215			# get current content
216			with open(candidate, 'r') as f:
217				content = f.read()
218
219			# files equal -> no need to rewrite
220			if content == data:
221				return candidate
222			
223			# not equal and overwrite not allowed
224			elif not allow_overwrite:
225				# get non-existent file name
226				cnt = 0
227				while os.path.isfile(candidate):
228					cnt += 1
229					candidate = "{}-{}.json".format(path_name, cnt)
230
231		# write file
232		with open(candidate, 'w+') as f:
233			f.write(data)
234
235		return candidate
236
237
238class ExtractTextAgent(ExtractAgent):
239	"""
240		An extraction agent for text, create a subclass for your agent.
241	"""
242	
243	def extract_type() -> RiddleDataType:
244		return RiddleDataType.TEXT
245
246class ExtractAudioAgent(ExtractAgent):
247	"""
248		An extraction agent for audio, create a subclass for your agent.
249	"""
250	
251	def extract_type() -> RiddleDataType:
252		return RiddleDataType.AUDIO
253
254class ExtractImageAgent(ExtractAgent):
255	"""
256		An extraction agent for images, create a subclass for your agent.
257	"""
258	
259	def extract_type() -> RiddleDataType:
260		return RiddleDataType.IMAGE
261
262
263class SolveAgent(BasicAgent):
264	"""
265		A solve agent, create a subclass for your agent.
266	"""
267
268	def agent_capability() -> AgentCapability:
269		return AgentCapability.SOLVE
270	
271	def _process(self):
272		logger.debug(f"Start solve: {self._response.id}")
273		solution = self.handle(self._response.riddle, self._response.data)
274		logger.debug(f"End solve: {self._response.id} ({solution.solution}, {solution.explanation})")
275		
276		if len(solution.solution) == 0 or len(solution.explanation) == 0:
277			logger.info(f"Riddle {self._response.id}: Empty solution/ explanation after handling")
278
279		self._response.solution = solution
280		self._response.status.solve.finished = True
281
282		self._do_response = True
283	
284	@abstractmethod
285	@validate_call
286	def handle(self, riddle: Riddle, data: List[RiddleData]) -> RiddleSolution:
287		"""
288			Solve the `riddle` using `data` and return a solution.
289		"""
290	
291class GatekeeperAgent(BasicAgent):
292	"""
293		A gatekeeper agent, create a subclass for your agent.
294	"""
295
296	def agent_capability() -> AgentCapability:
297		return AgentCapability.GATEKEEPER
298	
299	def _process(self):
300		if self._response.solution is None:
301			self._response.solution = RiddleSolution(solution="", explanation="")
302
303		logger.debug(f"Start validate: {self._response.id}")
304		solution = self.handle(self._response.solution, self._response.riddle)
305		logger.debug(f"End validate: {self._response.id} ({solution.review}, {solution.accepted})")
306		
307		if solution.review is None or len(solution.review) == 0:
308			logger.info(f"Riddle {self._response.id}: Empty review after handling")
309
310		self._response.solution = solution
311		self._response.status.validate.finished = True
312		self._response.status.solved = solution.accepted
313
314		self._do_response = True
315	
316	@abstractmethod
317	@validate_call
318	def handle(self, solution:RiddleSolution, riddle:Riddle) -> RiddleSolution:
319		"""
320			Check the `solution` of `riddle` and return solution with populated `solution.accepted` and `solution.review`.
321		"""