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