243 lines
6.7 KiB
Python
243 lines
6.7 KiB
Python
|
|
import random
|
|
|
|
from abc import abstractmethod, ABC
|
|
from enum import Enum
|
|
from typing import List, Callable
|
|
|
|
from ums.utils import (
|
|
RiddleInformation, AgentMessage, RiddleDataType, RiddleData, Riddle,
|
|
RiddleStatus, RiddleSolution,
|
|
logger
|
|
)
|
|
|
|
class AgentCapability(Enum):
|
|
"""
|
|
The three different capabilities an agent can have.
|
|
"""
|
|
|
|
EXTRACT="extract"
|
|
SOLVE="solve"
|
|
GATEKEEPER="gatekeeper"
|
|
|
|
|
|
class BasicAgent(ABC):
|
|
"""
|
|
A basic agent, each agent will be a subclass of this class.
|
|
"""
|
|
|
|
@staticmethod
|
|
@abstractmethod
|
|
def agent_capability() -> AgentCapability:
|
|
"""
|
|
Represents the capabilities of this agent, for messages/ tasks of this capability, the `handle` method will be called.
|
|
"""
|
|
pass
|
|
|
|
def __init__(self, message:AgentMessage, send_message:Callable[[AgentMessage], bool]):
|
|
self._send_message = send_message
|
|
self._sub_cnt = 0
|
|
|
|
self._message = message
|
|
self._response = message.model_copy(deep=True)
|
|
|
|
self._do_response = False
|
|
|
|
self._process()
|
|
self._respond()
|
|
|
|
@abstractmethod
|
|
def _process(self):
|
|
pass
|
|
|
|
def _respond(self):
|
|
send_it = lambda: self._send_message(self._response)
|
|
if self.before_response(self._response, send_it) and self._do_response:
|
|
send_it()
|
|
logger.debug(f"Response sent {self._response.id}")
|
|
else:
|
|
logger.debug(f"Stopped response {self._response.id}")
|
|
|
|
def before_response(self, response:AgentMessage, send_it:Callable[[], None]) -> bool:
|
|
"""
|
|
This method is called before the response is sent.
|
|
If the method returns `False` no response will be sent.
|
|
Thus, by overwriting this method, a response can be prevented.
|
|
|
|
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.
|
|
(Hence, one may stop sending the response and later call `send_it()` to send the response.)
|
|
"""
|
|
return True
|
|
|
|
def message(self) -> AgentMessage:
|
|
"""
|
|
Get the message this agent object is working on.
|
|
"""
|
|
return self._message;
|
|
|
|
def sub_riddle(self,
|
|
riddle:Riddle, data:List[RiddleData]=[], status:RiddleStatus=None
|
|
) -> AgentMessage|bool:
|
|
"""
|
|
Create a new sub-riddle for solving details of the current riddle.
|
|
For the sub-riddle, give a `riddle` and optionally, a selection of `data` items (default none) and a `status` (default `ums.utils.types.RiddleStatus()`).
|
|
By changing a status, different steps can be (de-)selected.
|
|
|
|
Return the message of the sub-riddle or `false` on error.
|
|
"""
|
|
|
|
if status is None:
|
|
status = RiddleStatus()
|
|
|
|
# new sub riddle id
|
|
self._sub_cnt += 1
|
|
new_id = "{}-sub-{}.{}".format(self._message.id, self._sub_cnt, int(random.random()*100))
|
|
|
|
self._message.sub_ids.append(new_id)
|
|
self._response.sub_ids.append(new_id)
|
|
|
|
# create the riddle's message
|
|
sub_msg = AgentMessage(
|
|
id=new_id,
|
|
riddle=riddle,
|
|
data=data,
|
|
status=status
|
|
)
|
|
logger.debug(f"Created sub-riddle {sub_msg.id}")
|
|
|
|
# send it
|
|
if self._send_message(sub_msg):
|
|
return sub_msg
|
|
else:
|
|
return False
|
|
|
|
@abstractmethod
|
|
def handle(self, *args:RiddleInformation) -> RiddleInformation:
|
|
"""
|
|
Handle a single task of the agent, the arguments and return value depends on the actual task (see subclass)!
|
|
|
|
**This is the method to implement!**
|
|
|
|
The full message is available via `message()`, a sub riddle can be created with `sub_riddle()`.
|
|
"""
|
|
pass
|
|
|
|
class ExtractAgent(BasicAgent):
|
|
"""
|
|
An extraction agent.
|
|
"""
|
|
|
|
def agent_capability() -> AgentCapability:
|
|
return AgentCapability.EXTRACT
|
|
|
|
@staticmethod
|
|
@abstractmethod
|
|
def extract_type() -> RiddleDataType:
|
|
"""
|
|
Represents the data this agent can process.
|
|
"""
|
|
pass
|
|
|
|
def _process(self):
|
|
for i, data in enumerate(self._response.data):
|
|
if data.type == self.__class__.extract_type():
|
|
logger.debug(f"Start extraction '{data.file_plain}'")
|
|
result = self.handle(data)
|
|
logger.debug(f"End extraction '{data.file_plain}' ('{result.file_extracted}')")
|
|
|
|
if result.file_extracted is None:
|
|
logger.info(f"Riddle {self._response.id}: 'file_extracted' for data '{data.file_plain}' still empty after handling")
|
|
|
|
self._response.data[i] = result
|
|
self._do_response = True
|
|
|
|
self._response.status.extract.finished = True
|
|
|
|
@abstractmethod
|
|
def handle(self, data:RiddleData) -> RiddleData:
|
|
"""
|
|
Process the item `data`, create extraction file and return `data` with populated `data.file_extracted`.
|
|
"""
|
|
|
|
class ExtractTextAgent(ExtractAgent):
|
|
"""
|
|
An extraction agent for text, create a subclass for your agent.
|
|
"""
|
|
|
|
def extract_type() -> RiddleDataType:
|
|
return RiddleDataType.TEXT
|
|
|
|
class ExtractAudioAgent(ExtractAgent):
|
|
"""
|
|
An extraction agent for audio, create a subclass for your agent.
|
|
"""
|
|
|
|
def extract_type() -> RiddleDataType:
|
|
return RiddleDataType.AUDIO
|
|
|
|
class ExtractImageAgent(ExtractAgent):
|
|
"""
|
|
An extraction agent for images, create a subclass for your agent.
|
|
"""
|
|
|
|
def extract_type() -> RiddleDataType:
|
|
return RiddleDataType.IMAGE
|
|
|
|
|
|
class SolveAgent(BasicAgent):
|
|
"""
|
|
A solve agent, create a subclass for your agent.
|
|
"""
|
|
|
|
def agent_capability() -> AgentCapability:
|
|
return AgentCapability.SOLVE
|
|
|
|
def _process(self):
|
|
logger.debug(f"Start solve: {self._response.id}")
|
|
solution = self.handle(self._response.riddle, self._response.data)
|
|
logger.debug(f"End solve: {self._response.id} ({solution.solution}, {solution.explanation})")
|
|
|
|
if len(solution.solution) == 0 or len(solution.explanation) == 0:
|
|
logger.info(f"Riddle {self._response.id}: Empty solution/ explanation after handling")
|
|
|
|
self._response.solution = solution
|
|
self._response.status.solve.finished = True
|
|
|
|
self._do_response = True
|
|
|
|
@abstractmethod
|
|
def handle(self, riddle:Riddle, data:RiddleData) -> RiddleSolution:
|
|
"""
|
|
Solve the `riddle` using `data` and return a solution.
|
|
"""
|
|
|
|
class GatekeeperAgent(BasicAgent):
|
|
"""
|
|
A gatekeeper agent, create a subclass for your agent.
|
|
"""
|
|
|
|
def agent_capability() -> AgentCapability:
|
|
return AgentCapability.GATEKEEPER
|
|
|
|
def _process(self):
|
|
if self._response.solution is None:
|
|
self._response.solution = RiddleSolution(solution="", explanation="")
|
|
|
|
logger.debug(f"Start validate: {self._response.id}")
|
|
solution = self.handle(self._response.solution, self._response.riddle)
|
|
logger.debug(f"End validate: {self._response.id} ({solution.review}, {solution.accepted})")
|
|
|
|
if solution.review is None or len(solution.review) == 0:
|
|
logger.info(f"Riddle {self._response.id}: Empty review after handling")
|
|
|
|
self._response.solution = solution
|
|
self._response.status.validate.finished = True
|
|
self._response.status.solved = solution.accepted
|
|
|
|
self._do_response = True
|
|
|
|
@abstractmethod
|
|
def handle(self, solution:RiddleSolution, riddle:Riddle) -> RiddleSolution:
|
|
"""
|
|
Check the `solution` of `riddle` and return solution with populated `solution.accepted` and `solution.review`.
|
|
""" |