added state store

Author: Kristian Nylund

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -12,6 +12,7 @@
 from autogen_text_2_sql.custom_agents.parallel_query_solving_agent import (
     ParallelQuerySolvingAgent,
 )
+from state_store import StateStore
 from autogen_agentchat.messages import TextMessage
 import json
 import os
@@ -31,9 +32,13 @@
 
 
 class AutoGenText2Sql:
-    def __init__(self, **kwargs):
+    def __init__(self, state_store : StateStore, **kwargs):
         self.target_engine = os.environ["Text2Sql__DatabaseEngine"].upper()
 
+        if not state_store:
+            raise ValueError("State store must be provided")
+        self.state_store = state_store
+
         if "use_case" not in kwargs:
             logging.warning(
                 "No use case provided. It is advised to provide a use case to help the LLM reason."
@@ -250,43 +255,31 @@ def extract_answer_payload(self, messages: list) -> AnswerWithSourcesPayload:
 
     async def process_user_message(
         self,
+        thread_id: str,
         message_payload: UserMessagePayload,
-        chat_history: list[InteractionPayload] = None,
     ) -> AsyncGenerator[InteractionPayload, None]:
         """Process the complete message through the unified system.
 
         Args:
         ----
+            thread_id (str): The ID of the thread the message belongs to.
             task (str): The user message to process.
-            chat_history (list[str], optional): The chat history. Defaults to None. The last message is the most recent message.
             injected_parameters (dict, optional): Parameters to pass to agents. Defaults to None.
 
         Returns:
         -------
             dict: The response from the system.
         """
         logging.info("Processing message: %s", message_payload.body.user_message)
-        logging.info("Chat history: %s", chat_history)
 
         agent_input = {
             "message": message_payload.body.user_message,
             "injected_parameters": message_payload.body.injected_parameters,
         }
 
-        latest_state = None
-        if chat_history is not None:
-            # Update input
-            for chat in reversed(chat_history):
-                if chat.root.payload_type in [
-                    PayloadType.ANSWER_WITH_SOURCES,
-                    PayloadType.DISAMBIGUATION_REQUESTS,
-                ]:
-                    latest_state = chat.body.assistant_state
-                    break
-
-        # TODO: Trim the chat history to the last message from the user
-        if latest_state is not None:
-            await self.agentic_flow.load_state(latest_state)
+        state = self.state_store.get_state(thread_id)
+        if state is not None:
+            await self.agentic_flow.load_state(state)
 
         async for message in self.agentic_flow.run_stream(task=json.dumps(agent_input)):
             logging.debug("Message: %s", message)
@@ -340,7 +333,7 @@ async def process_user_message(
         ):
             # Get the state
             assistant_state = await self.agentic_flow.save_state()
-            payload.body.assistant_state = assistant_state
+            self.state_store.save_state(thread_id, assistant_state)
 
             logging.debug("Final Payload: %s", payload)
 

text_2_sql/autogen/src/autogen_text_2_sql/state_store.py

@@ -0,0 +1,22 @@
+from abc import ABC, abstractmethod
+
+class StateStore(ABC):
+    @abstractmethod
+    def get_state(self, thread_id):
+        pass
+
+    @abstractmethod
+    def save_state(self, thread_id, state):
+        pass
+
+
+class InMemoryStateStore(StateStore):
+    def __init__(self):
+        # Replace with a caching library or something to have some sort of expiry for entries so this doesn't grow forever
+        self.cache = {}
+
+    def get_state(self, thread_id):
+        return self.cache.get(thread_id)
+
+    def save_state(self, thread_id, state):
+        self.cache[thread_id] = state

text_2_sql/autogen/src/autogen_text_2_sql/test.py

@@ -0,0 +1,6 @@
+from state_store import InMemoryStateStore
+
+x=InMemoryStateStore()
+print(x.get_state("1"))
+x.save_state("1", {'x':2})
+print(x.get_state("1"))

text_2_sql/text_2_sql_core/src/text_2_sql_core/payloads/interaction_payloads.py

@@ -57,7 +57,6 @@ class DismabiguationRequest(InteractionPayloadBase):
         decomposed_user_messages: list[list[str]] = Field(
             default_factory=list, alias="decomposedUserMessages"
         )
-        assistant_state: dict | None = Field(default=None, alias="assistantState")
 
     payload_type: Literal[PayloadType.DISAMBIGUATION_REQUESTS] = Field(
         PayloadType.DISAMBIGUATION_REQUESTS, alias="payloadType"
@@ -86,7 +85,6 @@ class Source(InteractionPayloadBase):
             default_factory=list, alias="decomposedUserMessages"
         )
         sources: list[Source] = Field(default_factory=list)
-        assistant_state: dict | None = Field(default=None, alias="assistantState")
 
     payload_type: Literal[PayloadType.ANSWER_WITH_SOURCES] = Field(
         PayloadType.ANSWER_WITH_SOURCES, alias="payloadType"