Updated agent

Author: BenConstable9

text_2_sql/autogen/agentic_text_2_sql.ipynb

@@ -8,7 +8,8 @@
       "source": [
         "import dotenv\n",
         "import logging\n",
-        "from agentic_text_2_sql import text_2_sql_generator, text_2_sql_cache_updater"
+        "from autogen_agentchat.task import Console\n",
+        "from agentic_text_2_sql import text_2_sql_generator"
       ]
     },
     {
@@ -35,7 +36,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "result = await text_2_sql_generator.run(task=\"What are the total number of sales within 2008?\")"
+        "result = text_2_sql_generator.run_stream(task=\"What are the total number of sales within 2008?\")"
       ]
     },
     {
@@ -44,7 +45,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "print(result.messages[-1].content)"
+        "await Console(result)"
       ]
     },
     {

text_2_sql/autogen/agentic_text_2_sql.py

@@ -2,8 +2,9 @@
 from autogen_agentchat.teams import SelectorGroupChat
 from utils.models import MINI_MODEL
 from utils.llm_agent_creator import LLMAgentCreator
-from autogen_core.components.models import FunctionExecutionResult
 import logging
+from custom_agents.sql_query_cache_agent import SqlQueryCacheAgent
+import json
 
 SQL_QUERY_GENERATION_AGENT = LLMAgentCreator.create(
     "sql_query_generation_agent",
@@ -16,39 +17,47 @@
     target_engine="Microsoft SQL Server",
     engine_specific_rules="Use TOP X to limit the number of rows returned instead of LIMIT X. NEVER USE LIMIT X as it produces a syntax error.",
 )
-SQL_QUERY_CACHE_AGENT = LLMAgentCreator.create("sql_query_cache_agent")
+SQL_QUERY_CACHE_AGENT = SqlQueryCacheAgent()
 ANSWER_AGENT = LLMAgentCreator.create("answer_agent")
 QUESTION_DECOMPOSITION_AGENT = LLMAgentCreator.create("question_decomposition_agent")
 
 
 def text_2_sql_generator_selector_func(messages):
     logging.info("Messages: %s", messages)
+    decision = None  # Initialize decision variable
+
     if len(messages) == 1:
-        return "sql_query_cache_agent"
+        decision = "sql_query_cache_agent"
 
     elif (
         messages[-1].source == "sql_query_cache_agent"
-        and isinstance(messages[-1].content, FunctionExecutionResult)
-        and messages[-1].content.content is not None
+        and messages[-1].content is not None
     ):
-        return "sql_query_correction_agent"
+        cache_result = json.loads(messages[-1].content)
+        if cache_result.get("cached_questions_and_schemas") is not None:
+            decision = "sql_query_correction_agent"
+        else:
+            decision = "sql_schema_selection_agent"
 
     elif messages[-1].source == "question_decomposition_agent":
-        return "sql_schema_selection_agent"
+        decision = "sql_schema_selection_agent"
 
     elif messages[-1].source == "sql_schema_selection_agent":
-        return "sql_query_generation_agent"
+        decision = "sql_query_generation_agent"
 
     elif (
         messages[-1].source == "sql_query_correction_agent"
         and messages[-1].content == "VALIDATED"
     ):
-        return "answer_agent"
+        decision = "answer_agent"
 
     elif messages[-1].source == "sql_query_correction_agent":
-        return "sql_query_correction_agent"
+        decision = "sql_query_correction_agent"
+
+    # Log the decision
+    logging.info("Decision: %s", decision)
 
-    return None
+    return decision
 
 
 termination = TextMentionTermination("TERMINATE") | MaxMessageTermination(10)
@@ -61,6 +70,7 @@ def text_2_sql_generator_selector_func(messages):
         ANSWER_AGENT,
         QUESTION_DECOMPOSITION_AGENT,
     ],
+    allow_repeated_speaker=False,
     model_client=MINI_MODEL,
     termination_condition=termination,
     selector_func=text_2_sql_generator_selector_func,

text_2_sql/autogen/custom_agents/__init__.py

@@ -0,0 +1,51 @@
+from typing import AsyncGenerator, List, Sequence
+
+from autogen_agentchat.agents import BaseChatAgent
+from autogen_agentchat.base import Response
+from autogen_agentchat.messages import AgentMessage, ChatMessage, TextMessage
+from autogen_core.base import CancellationToken
+from utils.sql_utils import fetch_queries_from_cache
+import json
+import logging
+
+
+class SqlQueryCacheAgent(BaseChatAgent):
+    def __init__(self):
+        super().__init__(
+            "sql_query_cache_agent",
+            "An agent that fetches the queries from the cache based on the user question.",
+        )
+
+    @property
+    def produced_message_types(self) -> List[type[ChatMessage]]:
+        return [TextMessage]
+
+    async def on_messages(
+        self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
+    ) -> Response:
+        # Calls the on_messages_stream.
+        response: Response | None = None
+        async for message in self.on_messages_stream(messages, cancellation_token):
+            if isinstance(message, Response):
+                response = message
+        assert response is not None
+        return response
+
+    async def on_messages_stream(
+        self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
+    ) -> AsyncGenerator[AgentMessage | Response, None]:
+        user_question = messages[0].content
+
+        # Fetch the queries from the cache based on the user question.
+        logging.info("Fetching queries from cache based on the user question...")
+
+        cached_queries = await fetch_queries_from_cache(user_question)
+
+        yield Response(
+            chat_message=TextMessage(
+                content=json.dumps(cached_queries), source=self.name
+            )
+        )
+
+    async def on_reset(self, cancellation_token: CancellationToken) -> None:
+        pass

text_2_sql/autogen/custom_agents/sql_query_cache_agent.py

@@ -7,6 +7,10 @@ description:
 system_message:
   "You are a helpful AI Assistant that specialises in selecting relevant SQL schemas to answer a given user's question.
 
-  Use the tools available to you to select the correct schemas that will help. Extract key terms from the user's question and use them to search for the correct schema."
+  Use the tools available to you to select the correct schemas that will help. Extract key terms from the user's question and use them to search for the correct schema.
+
+  Limit the number of calls to the 'sql_get_entity_schemas_tool' tool to avoid unnecessary calls.
+
+  If you are unsure about the schema, you can ask the user for more information or ask for clarification."
 tools:
   - sql_get_entity_schemas_tool

text_2_sql/autogen/llm_agents/sql_query_cache_agent.yaml

@@ -1,10 +1,9 @@
 import yaml
 from autogen_core.components.tools import FunctionTool
-from autogen_agentchat.agents import ToolUseAssistantAgent
+from autogen_agentchat.agents import AssistantAgent
 from utils.sql_utils import (
     query_execution,
     get_entity_schemas,
-    fetch_queries_from_cache,
 )
 from utils.models import MINI_MODEL
 from jinja2 import Template
@@ -37,11 +36,6 @@ def get_tool(cls, tool_name):
                 get_entity_schemas,
                 description="Gets the schema of a view or table in the SQL Database by selecting the most relevant entity based on the search term. Extract key terms from the user question and use these as the search term. Several entities may be returned. Only use when the provided schemas in the system prompt are not sufficient to answer the question.",
             )
-        elif tool_name == "sql_query_cache_tool":
-            return FunctionTool(
-                fetch_queries_from_cache,
-                description="Fetch the pre-assembled queries, and potential results from the cache based on the user's question.",
-            )
         else:
             raise ValueError(f"Tool {tool_name} not found")
 
@@ -62,9 +56,9 @@ def create(cls, name: str, **kwargs):
             for tool in agent_file["tools"]:
                 tools.append(cls.get_tool(tool))
 
-        agent = ToolUseAssistantAgent(
+        agent = AssistantAgent(
             name=name,
-            registered_tools=tools,
+            tools=tools,
             model_client=cls.get_model(agent_file["model"]),
             description=cls.get_property_and_render_parameters(
                 agent_file, "description", kwargs

text_2_sql/autogen/llm_agents/sql_schema_selection_agent.yaml

@@ -112,7 +112,7 @@ async def fetch_queries_from_cache(question: str) -> str:
     )
 
     if len(cached_schemas) == 0:
-        return None
+        return {"cached_questions_and_schemas": None}
 
     logging.info("Cached schemas: %s", cached_schemas)
     if PRE_RUN_QUERY_CACHE and len(cached_schemas) > 0:
@@ -139,6 +139,6 @@ async def fetch_queries_from_cache(question: str) -> str:
                     "schemas": sql_query["Schemas"],
                 }
 
-            return query_result_store
+            return {"cached_questions_and_schemas": query_result_store}
 
-    return {"cached_questions": cached_schemas}
+    return {"cached_questions_and_schemas": cached_schemas}