Schema Idea

ums/agent/agent.py

@@ -1,13 +1,25 @@
+# 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
+import random, os, json
 
 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
 )
 
@@ -34,6 +46,7 @@ class BasicAgent(ABC):
 		"""
 		pass
 
+	@validate_call
 	def __init__(self, message:AgentMessage, send_message:Callable[[AgentMessage], bool]):
 		self._send_message = send_message
 		self._sub_cnt = 0
@@ -57,7 +70,8 @@ class BasicAgent(ABC):
 			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.
@@ -69,12 +83,14 @@ class BasicAgent(ABC):
 		"""
 		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:
@@ -122,6 +138,21 @@ class BasicAgent(ABC):
 		"""
 		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.
@@ -154,11 +185,33 @@ class ExtractAgent(BasicAgent):
 		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) -> 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`.
+		"""
+
+		path_name = data.file_plain[:data.file_plain.rfind('.')]
+		
+		candidate = "{}.json".format(path_name)
+		cnt = 0
+		while os.path.isfile(candidate):
+			cnt += 1
+			candidate = "{}-{}.json".format(path_name, cnt)
+
+		with open(candidate, 'w+') as f:
+			f.write(extracted.model_dump_json())
+
+		return candidate
+
+
 class ExtractTextAgent(ExtractAgent):
 	"""
 		An extraction agent for text, create a subclass for your agent.
@@ -206,7 +259,8 @@ class SolveAgent(BasicAgent):
 		self._do_response = True
 	
 	@abstractmethod
-	def handle(self, riddle:Riddle, data:RiddleData) -> RiddleSolution:
+	@validate_call
+	def handle(self, riddle: Riddle, data: List[RiddleData]) -> RiddleSolution:
 		"""
 			Solve the `riddle` using `data` and return a solution.
 		"""
@@ -237,6 +291,7 @@ class GatekeeperAgent(BasicAgent):
 		self._do_response = True
 	
 	@abstractmethod
+	@validate_call
 	def handle(self, solution:RiddleSolution, riddle:Riddle) -> RiddleSolution:
 		"""
 			Check the `solution` of `riddle` and return solution with populated `solution.accepted` and `solution.review`.