ums.agent.agent

  1import random
  2
  3from abc import abstractmethod, ABC
  4from enum import Enum
  5from typing import List, Callable
  6
  7from ums.utils import (
  8	RiddleInformation, AgentMessage, RiddleDataType, RiddleData, Riddle,
  9	RiddleStatus, RiddleSolution,
 10	logger
 11)
 12
 13class AgentCapability(Enum):
 14	"""
 15		The three different capabilities an agent can have.
 16	"""
 17
 18	EXTRACT="extract"
 19	SOLVE="solve"
 20	GATEKEEPER="gatekeeper"
 21
 22
 23class BasicAgent(ABC):
 24	"""
 25		A basic agent, each agent will be a subclass of this class.
 26	"""
 27
 28	@staticmethod
 29	@abstractmethod
 30	def agent_capability() -> AgentCapability:
 31		"""
 32			Represents the capabilities of this agent, for messages/ tasks of this capability, the `handle` method will be called.
 33		"""
 34		pass
 35
 36	def __init__(self, message:AgentMessage, send_message:Callable[[AgentMessage], bool]):
 37		self._send_message = send_message
 38		self._sub_cnt = 0
 39
 40		self._message = message
 41		self._response = message.model_copy(deep=True)
 42		
 43		self._do_response = False
 44
 45		self._process()
 46		self._respond()
 47
 48	@abstractmethod
 49	def _process(self):
 50		pass
 51
 52	def _respond(self):
 53		send_it = lambda: self._send_message(self._response)
 54		if self.before_response(self._response, send_it) and self._do_response:
 55			send_it()
 56			logger.debug(f"Response sent {self._response.id}")
 57		else:
 58			logger.debug(f"Stopped response {self._response.id}")
 59		
 60	def before_response(self, response:AgentMessage, send_it:Callable[[], None]) -> bool:
 61		"""
 62			This method is called before the response is sent.
 63			If the method returns `False` no response will be sent. 
 64			Thus, by overwriting this method, a response can be prevented.
 65
 66			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. 
 67			(Hence, one may stop sending the response and later call `send_it()` to send the response.)
 68		"""
 69		return True
 70	
 71	def message(self) -> AgentMessage:
 72		"""
 73			Get the message this agent object is working on. 
 74		"""
 75		return self._message;
 76
 77	def sub_riddle(self,
 78			riddle:Riddle, data:List[RiddleData]=[], status:RiddleStatus=None
 79		) -> AgentMessage|bool:
 80			"""
 81				Create a new sub-riddle for solving details of the current riddle.
 82				For the sub-riddle, give a `riddle` and optionally, a selection of `data` items (default none) and a `status` (default `ums.utils.types.RiddleStatus()`).
 83				By changing a status, different steps can be (de-)selected.
 84
 85				Return the message of the sub-riddle or `false` on error.
 86			"""
 87
 88			if status is None:
 89				status = RiddleStatus()
 90			
 91			# new sub riddle id
 92			self._sub_cnt += 1
 93			new_id = "{}-sub-{}.{}".format(self._message.id, self._sub_cnt, int(random.random()*100))
 94
 95			self._message.sub_ids.append(new_id)
 96			self._response.sub_ids.append(new_id)
 97
 98			# create the riddle's message
 99			sub_msg = AgentMessage(
100				id=new_id,
101				riddle=riddle,
102				data=data,
103				status=status
104			)
105			logger.debug(f"Created sub-riddle {sub_msg.id}")
106
107			# send it
108			if self._send_message(sub_msg):
109				return sub_msg
110			else:
111				return False
112
113	@abstractmethod
114	def handle(self, *args:RiddleInformation) -> RiddleInformation:
115		"""
116			Handle a single task of the agent, the arguments and return value depends on the actual task (see subclass)!
117
118			**This is the method to implement!**
119
120			The full message is available via `message()`, a sub riddle can be created with `sub_riddle()`.
121		"""
122		pass
123
124class ExtractAgent(BasicAgent):
125	"""
126		An extraction agent.
127	"""
128
129	def agent_capability() -> AgentCapability:
130		return AgentCapability.EXTRACT
131	
132	@staticmethod
133	@abstractmethod
134	def extract_type() -> RiddleDataType:
135		"""
136			Represents the data this agent can process.
137		"""
138		pass
139
140	def _process(self):
141		for i, data in enumerate(self._response.data):
142			if data.type == self.__class__.extract_type():
143				logger.debug(f"Start extraction '{data.file_plain}'")
144				result = self.handle(data)
145				logger.debug(f"End extraction '{data.file_plain}' ('{result.file_extracted}')")
146
147				if result.file_extracted is None:
148					logger.info(f"Riddle {self._response.id}: 'file_extracted' for data '{data.file_plain}' still empty after handling")
149
150				self._response.data[i] = result
151				self._do_response = True
152
153		self._response.status.extract.finished = True
154
155	@abstractmethod
156	def handle(self, data:RiddleData) -> RiddleData:
157		"""
158			Process the item `data`, create extraction file and return `data` with populated `data.file_extracted`.
159		"""
160
161class ExtractTextAgent(ExtractAgent):
162	"""
163		An extraction agent for text, create a subclass for your agent.
164	"""
165	
166	def extract_type() -> RiddleDataType:
167		return RiddleDataType.TEXT
168
169class ExtractAudioAgent(ExtractAgent):
170	"""
171		An extraction agent for audio, create a subclass for your agent.
172	"""
173	
174	def extract_type() -> RiddleDataType:
175		return RiddleDataType.AUDIO
176
177class ExtractImageAgent(ExtractAgent):
178	"""
179		An extraction agent for images, create a subclass for your agent.
180	"""
181	
182	def extract_type() -> RiddleDataType:
183		return RiddleDataType.IMAGE
184
185
186class SolveAgent(BasicAgent):
187	"""
188		A solve agent, create a subclass for your agent.
189	"""
190
191	def agent_capability() -> AgentCapability:
192		return AgentCapability.SOLVE
193	
194	def _process(self):
195		logger.debug(f"Start solve: {self._response.id}")
196		solution = self.handle(self._response.riddle, self._response.data)
197		logger.debug(f"End solve: {self._response.id} ({solution.solution}, {solution.explanation})")
198		
199		if len(solution.solution) == 0 or len(solution.explanation) == 0:
200			logger.info(f"Riddle {self._response.id}: Empty solution/ explanation after handling")
201
202		self._response.solution = solution
203		self._response.status.solve.finished = True
204
205		self._do_response = True
206	
207	@abstractmethod
208	def handle(self, riddle:Riddle, data:RiddleData) -> RiddleSolution:
209		"""
210			Solve the `riddle` using `data` and return a solution.
211		"""
212	
213class GatekeeperAgent(BasicAgent):
214	"""
215		A gatekeeper agent, create a subclass for your agent.
216	"""
217
218	def agent_capability() -> AgentCapability:
219		return AgentCapability.GATEKEEPER
220	
221	def _process(self):
222		if self._response.solution is None:
223			self._response.solution = RiddleSolution(solution="", explanation="")
224
225		logger.debug(f"Start validate: {self._response.id}")
226		solution = self.handle(self._response.solution, self._response.riddle)
227		logger.debug(f"End validate: {self._response.id} ({solution.review}, {solution.accepted})")
228		
229		if solution.review is None or len(solution.review) == 0:
230			logger.info(f"Riddle {self._response.id}: Empty review after handling")
231
232		self._response.solution = solution
233		self._response.status.validate.finished = True
234		self._response.status.solved = solution.accepted
235
236		self._do_response = True
237	
238	@abstractmethod
239	def handle(self, solution:RiddleSolution, riddle:Riddle) -> RiddleSolution:
240		"""
241			Check the `solution` of `riddle` and return solution with populated `solution.accepted` and `solution.review`.
242		"""