This commit is contained in:
2024-12-21 18:57:39 +01:00
parent 4e98292160
commit 42a4449472
5 changed files with 112 additions and 56 deletions

View File

@ -15,7 +15,7 @@ import requests
from fastapi import BackgroundTasks
from ums.management.db import DB
from ums.utils import AgentMessage, AgentResponse, logger, RiddleData
from ums.utils import AgentMessage, AgentResponse, logger, RiddleData, RiddleSolution
class MessageProcessor():
@ -23,6 +23,7 @@ class MessageProcessor():
MESSAGE_MAX_CONTACTS = int(os.environ.get('MESSAGE_MAX_CONTACTS', 100))
REQUIRE_FULL_EXTRACT = os.environ.get('REQUIRE_FULL_EXTRACT', 'false').lower() == 'true'
REQUIRE_FULL_SOLVE = os.environ.get('REQUIRE_FULL_SOLVE', 'false').lower() == 'true'
MANAGEMENT_URL = os.environ.get('MANAGEMENT_URL', 'http://127.0.0.1:80').strip().strip('/')
@ -90,53 +91,99 @@ class MessageProcessor():
logger.warning(f"Message reached max number of contacts! {db_message.message.id}, {count}")
return
# combine different extractions in data items
# will update items in `db_message.message.data`
fully_extracted = self._add_extractions(db_message.message.id, db_message.message.data)
if (self.REQUIRE_FULL_EXTRACT and not fully_extracted) \
and not (db_message.message.status.extract.required and not db_message.message.status.extract.finished):
logger.warning(f"Postpone message, wait for full extract of items! {db_message.message.id}, {count}")
return
# check which step/ state the message requires the management to do
# -> IF
if db_message.message.status.extract.required and not db_message.message.status.extract.finished:
# send to extract agents
self._send_messages(self.AGENTS_PROCESS, db_message.message)
return
# combine different extractions in data items
# will update items in `db_message.message.data`
fully_extracted = self._add_extractions(db_message.message.id, db_message.message.data)
if self.REQUIRE_FULL_EXTRACT and not fully_extracted:
logger.warning(f"Postpone message, wait for full extract of items! {db_message.message.id}, {count}")
return
elif db_message.message.status.solve.required and not db_message.message.status.solve.finished:
# -> EL IF
if db_message.message.status.solve.required and not db_message.message.status.solve.finished:
# send to solve agents
self._send_messages(self.AGENTS_SOLVE, db_message.message)
elif db_message.message.status.validate.required and not db_message.message.status.validate.finished:
return
# combine different solutions
# will add solutions received before to `db_message.message.solution`
fully_solved = self._add_solutions(db_message.message.id, db_message.message.solution, db_message.message.status.trial)
if self.REQUIRE_FULL_SOLVE and not fully_solved:
logger.warning(f"Postpone message, wait for all solutions of riddle! {db_message.message.id}, {count}")
return
# -> EL IF
if db_message.message.status.validate.required and not db_message.message.status.validate.finished:
# send to solve agents
self._send_messages(self.AGENTS_GATEKEEPER, db_message.message)
return
else: # all steps "done"
# -> ELSE
# all steps "done"
# validate not required? (then solved will never be set to true, thus set it here)
if not db_message.message.status.validate.required:
db_message.message.status.solved = True
# validate not required? (then solved will never be set to true, thus set it here)
if not db_message.message.status.validate.required:
db_message.message.status.solved = True
if db_message.message.status.solved:
# yay, message is solved
self.db.set_solution(count=count, solution=True);
else:
# not solved, but all steps done
self.db.set_solution(count=count, solution=False);
if db_message.message.status.solved:
# yay, message is solved
self.db.set_solution(count=count, solution=True);
else:
# not solved, but all steps done
self.db.set_solution(count=count, solution=False);
# try again
self._do_again(db_message.message)
def _hash_solution(self, s:RiddleSolution) -> int:
return hash((s.solution, s.explanation, tuple((d.file_plain, d.type) for d in s.used_data)))
def _add_solutions(self, riddle_id:str, solution:List[RiddleSolution], trial:int) -> bool:
# do not do anything, if all solutions available
if len(solution) >= len(self.AGENTS_SOLVE):
return True
# try again
self._do_again(db_message.message)
contained = set(self._hash_solution(s) for s in solution)
# search db for solutions from before
for row in self.db.iterate(
id=riddle_id,
limit=min(self.db.len(id=riddle_id), 250)
):
# make sure to only use solutions from same "trial"
if row.message.status.trial == trial:
for s in row.message.solution:
h = self._hash_solution(s)
if h not in contained:
# add the 'new' solution
solution.append(s)
contained.add(h)
# all solutions found ?
if len(solution) >= len(self.AGENTS_SOLVE):
break
return len(solution) >= len(self.AGENTS_SOLVE)
def _hash_data(self, d:RiddleData) -> int:
return hash((d.file_plain, d.type, d.prompt))
def _add_extractions(self, riddle_id:str, data:List[RiddleData]) -> bool:
# get all the data items without extraction
empty_data = {}
for i, d in enumerate(data):
if d.file_extracted is None:
empty_data[self._hash_data(d)] = i
# do not do anything if fully extracted
if len(empty_data) == 0:
return True
# search db for extractions already available
for row in self.db.iterate(
@ -178,13 +225,13 @@ class MessageProcessor():
# increment trial
message.status.trial += 1
# append current solution als old one
if not message.solution is None:
message.riddle.solutions_before.append(
# append current solution(s) als old one(s)
if len(message.solution) > 0:
message.riddle.solutions_before.extend(
message.solution
)
# reset current solution
message.solution = None
message.solution = []
# add the riddle as new to management
self._send_message(self.MANAGEMENT_URL, message)