From 3fa3751a18705601dfed62a5b505fc32f6300773 Mon Sep 17 00:00:00 2001 From: Antonino Lorenzo <94693967+antoninoLorenzo@users.noreply.github.com> Date: Wed, 22 Jan 2025 19:53:10 +0100 Subject: [PATCH] refactoring: updated usage of llm and memory module. --- src/agent/agent.py | 10 ++-- .../architectures/default/architecture.py | 48 ++++++++++--------- src/routers/sessions.py | 2 +- 3 files changed, 32 insertions(+), 28 deletions(-) diff --git a/src/agent/agent.py b/src/agent/agent.py index 7b893ac..dd657e8 100644 --- a/src/agent/agent.py +++ b/src/agent/agent.py @@ -91,21 +91,21 @@ def new_session(self, sid: int): def get_session(self, sid: int): """Open existing conversation""" - return self.agent.memory.get_session(sid) + return self.agent.memory.get_conversation(sid) def get_sessions(self): """Returns list of Session objects""" - return self.agent.memory.get_sessions() + return self.agent.memory.get_conversations() def save_session(self, sid: int): """Saves the specified session to JSON""" - self.agent.memory.save_session(sid) + self.agent.memory.save_conversation(sid) def delete_session(self, sid: int): """Deletes the specified session""" - self.agent.memory.delete_session(sid) + self.agent.memory.delete_conversation(sid) def rename_session(self, sid: int, session_name: str): """Rename the specified session""" - self.agent.memory.rename_session(sid, session_name) + self.agent.memory.rename_conversation(sid, session_name) diff --git a/src/agent/architectures/default/architecture.py b/src/agent/architectures/default/architecture.py index ebf9572..5fdbd34 100644 --- a/src/agent/architectures/default/architecture.py +++ b/src/agent/architectures/default/architecture.py @@ -6,6 +6,7 @@ from tool_parse import ToolRegistry from src.agent import AgentArchitecture +from src.core import Conversation from src.core.llm import LLM from src.core.memory import Message, Role from src.utils import get_logger, LOGS_PATH @@ -124,7 +125,7 @@ def query( :returns: Generator with response text in chunks.""" # TODO: yield (chunk, context_length) # create a new conversation if not exists - if not self.memory.get_session(session_id): + if not self.memory.get_conversation(session_id): self.new_session(session_id) # route query @@ -143,7 +144,7 @@ def query( tool_call_str: str | None = None for tool_call_execution in self.__tool_call( user_input, - self.memory.get_session(session_id).message_dict, + self.memory.get_conversation(session_id), ): tool_call_state = tool_call_execution['state'] if tool_call_state == 'error': @@ -163,9 +164,9 @@ def query( assistant_index = 1 # Replace system prompt with the one built for specific assistant type - history = self.memory.get_session(session_id) + history = self.memory.get_conversation(session_id) history.messages[0] = Message(role=Role.SYS, content=prompt) - history.add_message( + history.add( Message( role=Role.USER, content=user_input_with_tool_call @@ -174,7 +175,7 @@ def query( # note: history.message_dict doesn't care about context length response = '' - for chunk, ctx_length in self.llm.query(history.message_dict): + for chunk, ctx_length in self.llm.query(history): if ctx_length: self.token_logger.info( f'Session: {session_id}; Tokens: {ctx_length}' @@ -194,9 +195,9 @@ def query( # remove tool call result from user input and add response to history history.messages[-1].content = user_input - history.add_message( + history.add( Message( - Role.ASSISTANT, + role=Role.ASSISTANT, content=response ) ) @@ -221,10 +222,13 @@ def __get_assistant_index( :param user_input: The user's input query. :return: An index to choose the proper prompt. """ - route_messages = [ - {'role': 'system', 'content': self.__prompts['router']}, - {'role': 'user', 'content': user_input} - ] + route_messages = Conversation( + name='get_assistant_index', + messages=[ + {'role': 'system', 'content': self.__prompts['router']}, + {'role': 'user', 'content': user_input} + ] + ) assistant_index_buffer = '' for chunk, _ in self.llm.query(route_messages): if not chunk: @@ -241,26 +245,26 @@ def __get_assistant_index( def __tool_call( self, user_input: str, - message_history: list + conversation: Conversation ): """Query a LLM for a tool call and executes it. :param user_input: The user's input query. - :param message_history: The conversation history. + :param conversation: The conversation history. :returns: Result of the tool execution.""" # replace system prompt and generate tool call - message_history[0] = { - 'role': 'system', - 'content': self.__prompts['tool'] - } - message_history.append({ - 'role': 'user', - 'content': user_input - }) + conversation.messages[0] = Message( + role='system', + content=self.__prompts['tool'] + ) + conversation.add(Message( + role='user', + content=user_input + )) tool_call_response = '' - for chunk, _ in self.llm.query(message_history): + for chunk, _ in self.llm.query(conversation): tool_call_response += chunk # extract tool call and run it diff --git a/src/routers/sessions.py b/src/routers/sessions.py index d821808..310e047 100644 --- a/src/routers/sessions.py +++ b/src/routers/sessions.py @@ -54,7 +54,7 @@ async def get_session(sid: int, agent: Agent = Depends(get_agent)): return { 'sid': sid, 'name': session.name, - 'messages': session.message_dict + 'messages': session.model_dump() }