326 lines
9.1 KiB
Python
326 lines
9.1 KiB
Python
# Agenten Plattform
|
|
#
|
|
# (c) 2024 Magnus Bender
|
|
# Institute of Humanities-Centered Artificial Intelligence (CHAI)
|
|
# Universitaet Hamburg
|
|
# https://www.chai.uni-hamburg.de/~bender
|
|
#
|
|
# source code released under the terms of GNU Public License Version 3
|
|
# https://www.gnu.org/licenses/gpl-3.0.txt
|
|
|
|
import random, os, json, time
|
|
|
|
from abc import abstractmethod, ABC
|
|
from enum import Enum
|
|
from typing import List, Callable
|
|
|
|
from pydantic import validate_call
|
|
|
|
from ums.utils import (
|
|
RiddleInformation, AgentMessage, RiddleDataType, RiddleData, Riddle,
|
|
RiddleStatus, RiddleSolution,
|
|
ExtractedData,
|
|
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
|
|
|
|
@validate_call
|
|
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):
|
|
# do a very short sleep
|
|
time.sleep(random.random())
|
|
|
|
# sending
|
|
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}")
|
|
|
|
@validate_call
|
|
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
|
|
|
|
@validate_call
|
|
def message(self) -> AgentMessage:
|
|
"""
|
|
Get the message this agent object is working on.
|
|
"""
|
|
return self._message;
|
|
|
|
@validate_call
|
|
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
|
|
|
|
@validate_call
|
|
def get_extracted(self, data:RiddleData) -> ExtractedData|None:
|
|
"""
|
|
Loads the extracted data from the `data` item (i.e., from the file `data.file_extracted`).
|
|
|
|
Returns None if no extracted data found.
|
|
"""
|
|
|
|
if not data.file_extracted is None:
|
|
return ExtractedData.model_validate(
|
|
json.load(open(data.file_extracted, 'r'))
|
|
)
|
|
|
|
return None
|
|
|
|
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
|
|
@validate_call
|
|
def handle(self, data:RiddleData) -> RiddleData:
|
|
"""
|
|
Process the item `data`, create extraction file and return `data` with populated `data.file_extracted`.
|
|
"""
|
|
|
|
@validate_call
|
|
def store_extracted(self, data:RiddleData, extracted:ExtractedData, allow_overwrite:bool=True) -> str:
|
|
"""
|
|
Stores the newly extracted data (in `extracted`) from `data` (i.e., `data.file_plain`)
|
|
and returns the filename to use in `data.file_extracted`.
|
|
|
|
If there already exists an extracted file for this `data`, the file will be overwritten if `allow_overwrite=True`.
|
|
Generally the system will check, if the contents of the current file are equal to the contents to write.
|
|
File with equal content will not be written again.
|
|
"""
|
|
|
|
# get path and name
|
|
path_name = data.file_plain[:data.file_plain.rfind('.')]
|
|
candidate = "{}.json".format(path_name)
|
|
|
|
# data to write
|
|
data = extracted.model_dump_json()
|
|
|
|
# check for file
|
|
if os.path.isfile(candidate):
|
|
|
|
# get current content
|
|
with open(candidate, 'r') as f:
|
|
content = f.read()
|
|
|
|
# files equal -> no need to rewrite
|
|
if content == data:
|
|
return candidate
|
|
|
|
# not equal and overwrite not allowed
|
|
elif not allow_overwrite:
|
|
# get non-existent file name
|
|
cnt = 0
|
|
while os.path.isfile(candidate):
|
|
cnt += 1
|
|
candidate = "{}-{}.json".format(path_name, cnt)
|
|
|
|
# write file
|
|
with open(candidate, 'w+') as f:
|
|
f.write(data)
|
|
|
|
return candidate
|
|
|
|
|
|
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.append(solution)
|
|
self._response.status.solve.finished = True
|
|
|
|
self._do_response = True
|
|
|
|
@abstractmethod
|
|
@validate_call
|
|
def handle(self, riddle: Riddle, data: List[RiddleData]) -> RiddleSolution:
|
|
"""
|
|
Solve the `riddle` using `data` and return a single solution.
|
|
"""
|
|
|
|
class GatekeeperAgent(BasicAgent):
|
|
"""
|
|
A gatekeeper agent, create a subclass for your agent.
|
|
"""
|
|
|
|
def agent_capability() -> AgentCapability:
|
|
return AgentCapability.GATEKEEPER
|
|
|
|
def _process(self):
|
|
if len(self._response.solution) == 0:
|
|
self._response.solution.append(RiddleSolution(solution="", explanation=""))
|
|
|
|
logger.debug(f"Start validate: {self._response.id}")
|
|
solution = self.handle(self._response.solution, self._response.riddle)
|
|
|
|
for single_solution in solution:
|
|
logger.debug(f"End validate: {self._response.id} ({single_solution.review}, {single_solution.accepted})")
|
|
if single_solution.review is None or len(single_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 = any(single_solution.accepted for single_solution in solution)
|
|
|
|
self._do_response = True
|
|
|
|
@abstractmethod
|
|
@validate_call
|
|
def handle(self, solution:List[RiddleSolution], riddle:Riddle) -> List[RiddleSolution]:
|
|
"""
|
|
Check the `solution` (multiple if multiple solver involved) of `riddle` and return solutions with populated `solution[i].accepted` and `solution[i].review`.
|
|
""" |