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) -> 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
201		path_name = data.file_plain[:data.file_plain.rfind('.')]
202		
203		candidate = "{}.json".format(path_name)
204		cnt = 0
205		while os.path.isfile(candidate):
206			cnt += 1
207			candidate = "{}-{}.json".format(path_name, cnt)
208
209		with open(candidate, 'w+') as f:
210			f.write(extracted.model_dump_json())
211
212		return candidate
213
214
215class ExtractTextAgent(ExtractAgent):
216	"""
217		An extraction agent for text, create a subclass for your agent.
218	"""
219	
220	def extract_type() -> RiddleDataType:
221		return RiddleDataType.TEXT
222
223class ExtractAudioAgent(ExtractAgent):
224	"""
225		An extraction agent for audio, create a subclass for your agent.
226	"""
227	
228	def extract_type() -> RiddleDataType:
229		return RiddleDataType.AUDIO
230
231class ExtractImageAgent(ExtractAgent):
232	"""
233		An extraction agent for images, create a subclass for your agent.
234	"""
235	
236	def extract_type() -> RiddleDataType:
237		return RiddleDataType.IMAGE
238
239
240class SolveAgent(BasicAgent):
241	"""
242		A solve agent, create a subclass for your agent.
243	"""
244
245	def agent_capability() -> AgentCapability:
246		return AgentCapability.SOLVE
247	
248	def _process(self):
249		logger.debug(f"Start solve: {self._response.id}")
250		solution = self.handle(self._response.riddle, self._response.data)
251		logger.debug(f"End solve: {self._response.id} ({solution.solution}, {solution.explanation})")
252		
253		if len(solution.solution) == 0 or len(solution.explanation) == 0:
254			logger.info(f"Riddle {self._response.id}: Empty solution/ explanation after handling")
255
256		self._response.solution = solution
257		self._response.status.solve.finished = True
258
259		self._do_response = True
260	
261	@abstractmethod
262	@validate_call
263	def handle(self, riddle: Riddle, data: List[RiddleData]) -> RiddleSolution:
264		"""
265			Solve the `riddle` using `data` and return a solution.
266		"""
267	
268class GatekeeperAgent(BasicAgent):
269	"""
270		A gatekeeper agent, create a subclass for your agent.
271	"""
272
273	def agent_capability() -> AgentCapability:
274		return AgentCapability.GATEKEEPER
275	
276	def _process(self):
277		if self._response.solution is None:
278			self._response.solution = RiddleSolution(solution="", explanation="")
279
280		logger.debug(f"Start validate: {self._response.id}")
281		solution = self.handle(self._response.solution, self._response.riddle)
282		logger.debug(f"End validate: {self._response.id} ({solution.review}, {solution.accepted})")
283		
284		if solution.review is None or len(solution.review) == 0:
285			logger.info(f"Riddle {self._response.id}: Empty review after handling")
286
287		self._response.solution = solution
288		self._response.status.validate.finished = True
289		self._response.status.solved = solution.accepted
290
291		self._do_response = True
292	
293	@abstractmethod
294	@validate_call
295	def handle(self, solution:RiddleSolution, riddle:Riddle) -> RiddleSolution:
296		"""
297			Check the `solution` of `riddle` and return solution with populated `solution.accepted` and `solution.review`.
298		"""