ums.management.process

  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 os, re
 12from typing import List
 13
 14import requests
 15from fastapi import BackgroundTasks
 16
 17from ums.management.db import DB
 18from ums.utils import AgentMessage, AgentResponse, logger, RiddleData
 19
 20class MessageProcessor():
 21
 22	SOLUTION_MAX_TRIALS = int(os.environ.get('SOLUTION_MAX_TRIALS', 5))
 23	MESSAGE_MAX_CONTACTS = int(os.environ.get('MESSAGE_MAX_CONTACTS', 100))
 24
 25	REQUIRE_FULL_EXTRACT = os.environ.get('REQUIRE_FULL_EXTRACT', 'false').lower() == 'true'
 26
 27	MANAGEMENT_URL = os.environ.get('MANAGEMENT_URL', 'http://127.0.0.1:80').strip().strip('/')
 28	
 29	AGENTS_PROCESS = tuple(map(
 30		lambda s:s.strip().strip('/'),
 31		os.environ.get('AGENTS_PROCESS', '').split(',')
 32	))
 33	AGENTS_SOLVE = tuple(map(
 34		lambda s:s.strip().strip('/'),
 35		os.environ.get('AGENTS_SOLVE', '').split(',')
 36	))
 37	AGENTS_GATEKEEPER = tuple(map(
 38		lambda s:s.strip().strip('/'),
 39		os.environ.get('AGENTS_GATEKEEPER', '').split(',')
 40	))
 41
 42	def __init__(self, db:DB):
 43		self.db = db
 44		self.management_name = self._get_name(self.MANAGEMENT_URL)
 45
 46		if len(self.AGENTS_PROCESS) == 0:
 47			logger.warning(f"Not Process Agent (AGENTS_PROCESS) found, this may be a problem!")
 48		if len(self.AGENTS_SOLVE) == 0:
 49			logger.warning(f"Not Solve Agent (AGENTS_SOLVE) found, this may be a problem!")
 50		if len(self.AGENTS_GATEKEEPER) == 0:
 51			logger.warning(f"Not Gatekeeper Agent (AGENTS_GATEKEEPER) found, this may be a problem!")
 52
 53	def _get_name(self, url:str) -> str:
 54		m = re.match(r'^https?://([^:]*)(?::(\d+))?$', url)
 55		return "unknown" if m == None else m.group(1)
 56
 57	def new_message(self,
 58		sender:str, receiver:str, message:AgentMessage,
 59		background_tasks: BackgroundTasks
 60	) -> AgentResponse:
 61
 62		try:
 63			db_count = self.db.add_message(sender, receiver, message)
 64			background_tasks.add_task(self._process_message, db_count)
 65
 66			return AgentResponse(
 67				count=db_count,
 68				msg="Added message to queue"
 69			)
 70		except Exception as e:
 71			return AgentResponse(
 72				count=-1,
 73				error=True,
 74				error_msg=str(e)
 75			)
 76
 77	def _process_message(self, count:int, ignore_processed:bool=False):
 78		db_message = self.db.by_count(count)
 79
 80		if db_message.processed and not ignore_processed:
 81			# do not process processed messages again
 82			return
 83		
 84		# now message processed!
 85		self.db.set_processed(count=count, processed=True)
 86		
 87		# increment contacts counter
 88		db_message.message.contacts += 1
 89		if db_message.message.contacts > self.MESSAGE_MAX_CONTACTS:
 90			logger.warning(f"Message reached max number of contacts! {db_message.message.id}, {count}")
 91			return
 92		
 93		# combine different extractions in data items
 94		#	will update items in `db_message.message.data`
 95		fully_extracted = self._add_extractions(db_message.message.id, db_message.message.data)
 96		if (self.REQUIRE_FULL_EXTRACT and not fully_extracted) \
 97			and not (db_message.message.status.extract.required and not db_message.message.status.extract.finished):
 98			logger.warning(f"Postpone message, wait for full extract of items! {db_message.message.id}, {count}")
 99			return
100
101		# check which step/ state the message requires the management to do
102		if db_message.message.status.extract.required and not db_message.message.status.extract.finished:
103			# send to extract agents
104			self._send_messages(self.AGENTS_PROCESS, db_message.message)
105
106		elif db_message.message.status.solve.required and not db_message.message.status.solve.finished:
107			# send to solve agents
108			self._send_messages(self.AGENTS_SOLVE, db_message.message)
109			
110		elif db_message.message.status.validate.required and not db_message.message.status.validate.finished:
111			# send to solve agents
112			self._send_messages(self.AGENTS_GATEKEEPER, db_message.message)
113
114		else: # all steps "done"
115
116			# validate not required? (then solved will never be set to true, thus set it here)
117			if not db_message.message.status.validate.required:
118				db_message.message.status.solved = True
119
120			if db_message.message.status.solved:
121				# yay, message is solved
122				self.db.set_solution(count=count, solution=True);
123			else:
124				# not solved, but all steps done
125				self.db.set_solution(count=count, solution=False);
126		
127				# try again
128				self._do_again(db_message.message)
129
130	def _hash_data(self, d:RiddleData) -> int:
131		return hash((d.file_plain, d.type, d.prompt))
132
133	def _add_extractions(self, riddle_id:str, data:List[RiddleData]) -> bool:
134		
135		# get all the data items without extraction
136		empty_data = {}
137		for i, d in enumerate(data):
138			if d.file_extracted is None:
139				empty_data[self._hash_data(d)] = i
140
141		# search db for extractions already available
142		for row in self.db.iterate(
143			id=riddle_id,
144			limit=min(self.db.len(id=riddle_id), 250)
145		):
146			# check for required extraction
147			for d in row.message.data:
148				# already extracted ? 
149				#	extraction file exists ?
150				#	one of the items, we do not have extractions for ?
151				#	the same data item ?
152				if not d.file_extracted is None \
153					and not d.file_extracted.startswith("missing:") \
154					and self._hash_data(d) in empty_data:
155						# copy the reference to the extracted data
156						data[empty_data[self._hash_data(d)]].file_extracted = d.file_extracted
157						# remove from items we need extracted data for
158						del empty_data[self._hash_data(d)]
159
160			# break if all extractions found
161			if len(empty_data) == 0:
162				break
163
164		return len(empty_data) == 0 # fully extracted
165
166	def _do_again(self, message:AgentMessage):
167		if message.status.trial < self.SOLUTION_MAX_TRIALS:
168			# try again, recycle message
169
170			# require steps again
171			if message.status.extract.required:
172				message.status.extract.finished = False 
173			if message.status.solve.required:
174				message.status.solve.finished = False 
175			if message.status.validate.required:
176				message.status.validate.finished = False
177
178			# increment trial
179			message.status.trial += 1
180
181			# append current solution als old one
182			if not message.solution is None:
183				message.riddle.solutions_before.append(
184					message.solution
185				)
186				# reset current solution
187				message.solution = None
188
189			# add the riddle as new to management
190			self._send_message(self.MANAGEMENT_URL, message)
191
192		else:
193			logger.info(f"Unsolved riddle after max number of trials: {message.id}")
194
195	def _send_messages(self, recipients:List[str], message:AgentMessage) -> bool:
196		ok  = True
197		for r in recipients:
198			ok = ok and self._send_message(r, message)
199		return ok
200
201	def _send_message(self, recipient:str, message:AgentMessage) -> bool:
202		db_count = self.db.add_message(
203			sender=self.management_name,
204			recipient=self._get_name(recipient),
205			message=message,
206			processed=False
207		)
208
209		r = requests.post(
210			"{}/message".format(recipient),
211			data=message.model_dump_json(),
212			headers={"accept" : "application/json", "content-type" : "application/json"}
213		)
214
215		if r.status_code == 200:
216			self.db.set_processed(db_count, processed=True)
217			return True
218		else:
219			logger.warning(f"Error sending message to: {recipient} {(r.text, r.headers)}")
220			return False
221