Move to agent creator

BenConstable9 · BenConstable9 · commit 92b236828604 · 2024-11-20T11:49:40.000Z
diff --git a/text_2_sql/autogen/agentic_text_2_sql.ipynb b/text_2_sql/autogen/agentic_text_2_sql.ipynb
@@ -2,7 +2,7 @@
   "cells": [
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 8,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -13,7 +13,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 9,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -22,18 +22,95 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 10,
       "metadata": {},
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "True"
+            ]
+          },
+          "execution_count": 10,
+          "metadata": {},
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "dotenv.load_dotenv()"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 11,
       "metadata": {},
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "INFO:autogen_core:Publishing message of type GroupChatStart to all subscribers: {'message': TextMessage(source='user', models_usage=None, content='What are the total number of sales within 2008?')}\n",
+            "INFO:autogen_core:Calling message handler for SQL_Query_Generation_Agent with message type GroupChatStart published by Unknown\n",
+            "INFO:autogen_core:Calling message handler for SQL_Schema_Selection_Agent with message type GroupChatStart published by Unknown\n",
+            "INFO:autogen_core:Calling message handler for SQL_Query_Correction_Agent with message type GroupChatStart published by Unknown\n",
+            "INFO:autogen_core:Calling message handler for SQL_Query_Cache_Agent with message type GroupChatStart published by Unknown\n",
+            "INFO:autogen_core:Calling message handler for Answer_Agent with message type GroupChatStart published by Unknown\n",
+            "INFO:autogen_core:Calling message handler for Question_Decomposition_Agent with message type GroupChatStart published by Unknown\n",
+            "INFO:autogen_core:Calling message handler for SQL_Query_Generation_Agent with message type GroupChatStart published by Unknown\n",
+            "INFO:autogen_core:Calling message handler for SQL_Schema_Selection_Agent with message type GroupChatStart published by Unknown\n",
+            "INFO:autogen_core:Calling message handler for SQL_Query_Correction_Agent with message type GroupChatStart published by Unknown\n",
+            "INFO:autogen_core:Calling message handler for SQL_Query_Cache_Agent with message type GroupChatStart published by Unknown\n",
+            "INFO:autogen_core:Calling message handler for Answer_Agent with message type GroupChatStart published by Unknown\n",
+            "INFO:autogen_core:Calling message handler for Question_Decomposition_Agent with message type GroupChatStart published by Unknown\n",
+            "INFO:autogen_core:Calling message handler for group_chat_manager with message type GroupChatStart published by Unknown\n",
+            "INFO:autogen_core:Publishing message of type GroupChatStart to all subscribers: {'message': TextMessage(source='user', models_usage=None, content='What are the total number of sales within 2008?')}\n",
+            "INFO:autogen_core:Calling message handler for collect_output_messages with message type GroupChatStart published by group_chat_manager/3b4d5ab6-7cfa-4179-8079-28bf0d1e0db4\n",
+            "INFO:autogen_agentchat.events:source='user' models_usage=None content='What are the total number of sales within 2008?'\n",
+            "INFO:httpx:HTTP Request: POST https://aoai-text2sql-adi.openai.azure.com/openai/deployments/gpt-4o-mini/chat/completions?api-version=2024-08-01-preview \"HTTP/1.1 200 OK\"\n",
+            "INFO:autogen_core.events:{\"prompt_tokens\": 1801, \"completion_tokens\": 6, \"type\": \"LLMCall\"}\n",
+            "INFO:autogen_core:Publishing message of type GroupChatRequestPublish to all subscribers: {}\n",
+            "INFO:autogen_core:Calling message handler for SQL_Query_Cache_Agent with message type GroupChatRequestPublish published by group_chat_manager/3b4d5ab6-7cfa-4179-8079-28bf0d1e0db4\n",
+            "INFO:httpx:HTTP Request: POST https://aoai-text2sql-adi.openai.azure.com/openai/deployments/gpt-4o-mini/chat/completions?api-version=2024-08-01-preview \"HTTP/1.1 200 OK\"\n",
+            "INFO:autogen_core.events:{\"prompt_tokens\": 1582, \"completion_tokens\": 28, \"type\": \"LLMCall\"}\n",
+            "INFO:autogen_core:Publishing message of type GroupChatMessage to all subscribers: {'message': ToolCallMessage(source='SQL_Query_Cache_Agent', models_usage=RequestUsage(prompt_tokens=1582, completion_tokens=28), content=[FunctionCall(id='call_CSvANg2842Hg57erUpUGssoe', arguments='{\"question\":\"What is the total number of sales orders in 2008?\"}', name='fetch_queries_from_cache')])}\n",
+            "INFO:autogen_core:Calling message handler for collect_output_messages with message type GroupChatMessage published by SQL_Query_Cache_Agent/3b4d5ab6-7cfa-4179-8079-28bf0d1e0db4\n",
+            "INFO:autogen_agentchat.events:source='SQL_Query_Cache_Agent' models_usage=RequestUsage(prompt_tokens=1582, completion_tokens=28) content=[FunctionCall(id='call_CSvANg2842Hg57erUpUGssoe', arguments='{\"question\":\"What is the total number of sales orders in 2008?\"}', name='fetch_queries_from_cache')]\n",
+            "INFO:httpx:HTTP Request: POST https://aoai-text2sql-adi.openai.azure.com/openai/deployments/text-embedding-ada-002/embeddings?api-version=2023-03-15-preview \"HTTP/1.1 200 OK\"\n",
+            "INFO:azure.core.pipeline.policies.http_logging_policy:Request URL: 'https://aisearch-text2sql-adi.search.windows.net/indexes('text-2-sql-query-cache-index')/docs/search.post.search?api-version=REDACTED'\n",
+            "Request method: 'POST'\n",
+            "Request headers:\n",
+            "    'Content-Type': 'application/json'\n",
+            "    'Content-Length': '34704'\n",
+            "    'api-key': 'REDACTED'\n",
+            "    'Accept': 'application/json;odata.metadata=none'\n",
+            "    'x-ms-client-request-id': '6a968312-a735-11ef-a762-0242ac110002'\n",
+            "    'User-Agent': 'azsdk-python-search-documents/11.6.0b5 Python/3.12.6 (Linux-5.15.153.1-microsoft-standard-WSL2-x86_64-with-glibc2.36)'\n",
+            "A body is sent with the request\n",
+            "INFO:azure.core.pipeline.policies.http_logging_policy:Response status: 200\n",
+            "Response headers:\n",
+            "    'Transfer-Encoding': 'chunked'\n",
+            "    'Content-Type': 'application/json; odata.metadata=none; odata.streaming=true; charset=utf-8'\n",
+            "    'Content-Encoding': 'REDACTED'\n",
+            "    'Vary': 'REDACTED'\n",
+            "    'Server': 'Microsoft-IIS/10.0'\n",
+            "    'Strict-Transport-Security': 'REDACTED'\n",
+            "    'Preference-Applied': 'REDACTED'\n",
+            "    'OData-Version': 'REDACTED'\n",
+            "    'request-id': '6a968312-a735-11ef-a762-0242ac110002'\n",
+            "    'elapsed-time': 'REDACTED'\n",
+            "    'Strict-Transport-Security': 'REDACTED'\n",
+            "    'Date': 'Wed, 20 Nov 2024 11:48:53 GMT'\n",
+            "INFO:root:Results: []\n",
+            "INFO:autogen_core:Publishing message of type GroupChatMessage to all subscribers: {'message': ToolCallResultMessage(source='SQL_Query_Cache_Agent', models_usage=None, content=[FunctionExecutionResult(content=\"{'cached_questions': None}\", call_id='call_CSvANg2842Hg57erUpUGssoe')])}\n",
+            "INFO:autogen_core:Calling message handler for collect_output_messages with message type GroupChatMessage published by SQL_Query_Cache_Agent/3b4d5ab6-7cfa-4179-8079-28bf0d1e0db4\n",
+            "INFO:autogen_agentchat.events:source='SQL_Query_Cache_Agent' models_usage=None content=[FunctionExecutionResult(content=\"{'cached_questions': None}\", call_id='call_CSvANg2842Hg57erUpUGssoe')]\n",
+            "INFO:httpx:HTTP Request: POST https://aoai-text2sql-adi.openai.azure.com/openai/deployments/gpt-4o-mini/chat/completions?api-version=2024-08-01-preview \"HTTP/1.1 500 model_error\"\n",
+            "INFO:openai._base_client:Retrying request to /chat/completions in 0.951392 seconds\n",
+            "INFO:httpx:HTTP Request: POST https://aoai-text2sql-adi.openai.azure.com/openai/deployments/gpt-4o-mini/chat/completions?api-version=2024-08-01-preview \"HTTP/1.1 500 model_error\"\n",
+            "INFO:openai._base_client:Retrying request to /chat/completions in 1.995870 seconds\n"
+          ]
+        }
+      ],
       "source": [
         "result = await SELECTOR.run(task=\"What are the total number of sales within 2008?\")"
       ]
diff --git a/text_2_sql/autogen/agents/answer_agent.yaml b/text_2_sql/autogen/agents/answer_agent.yaml
@@ -1,3 +1,5 @@
+model:
+  gpt-4o-mini
 description:
   An agent that takes the final results from the SQL query and writes the answer to the user's question
 system_message:
diff --git a/text_2_sql/autogen/agents/question_decomposition_agent.yaml b/text_2_sql/autogen/agents/question_decomposition_agent.yaml
@@ -1,3 +1,5 @@
+model:
+  gpt-4o-mini
 description:
   An agent that will decompose the user's question into smaller parts to be used in the SQL queries. Use this agent when the user's question is too complex to be answered in one SQL query. Only use if the user's question is too complex to be answered in one SQL query.
 
diff --git a/text_2_sql/autogen/agents/sql_query_cache_agent.yaml b/text_2_sql/autogen/agents/sql_query_cache_agent.yaml
@@ -1,3 +1,5 @@
+model:
+  gpt-4o-mini
 description:
   An agent that can take a user's question and use a pre-built SQL query cache to quickly fetch results for the user. Use this agent before all other agents to quickly fetch results for the user. If the results are not as expected, use the SQL Query Correction Agent to correct the SQL query if the schemas are similar. If no results are found, use the SQL Schema Selection Agent to select the correct schema.
 
@@ -6,3 +8,5 @@ system_message:
   You are a helpful AI Assistant that specialises in using a pre-built SQL query cache to quickly fetch results for the user.
 
   Use the tool and access to the pre-built query cache to ask the pass the user's question to the cache and fetch the results.
+tools:
+  - sql_query_cache_tool
diff --git a/text_2_sql/autogen/agents/sql_query_correction_agent.yaml b/text_2_sql/autogen/agents/sql_query_correction_agent.yaml
@@ -1,3 +1,5 @@
+model:
+  gpt-4o-mini
 description:
   An agent that will look at the SQL query, SQL query results and correct any mistakes in the SQL query to ensure the correct results are returned. Use this agent AFTER the SQL query has been executed and the results are not as expected.
 system_message:
@@ -6,3 +8,6 @@ system_message:
   Review the SQL query provided and correct any errors or issues that you find. Ensure that the corrected query returns the expected results.
 
   If you are consistently unable to correct the SQL query and cannot use the schemas to answer the question. Say 'I am unable to correct the SQL query. Please ask another question.' and then end your answer with 'TERMINATE'
+tools:
+  - sql_get_entity_schemas_tool
+  - sql_query_execution_tool
diff --git a/text_2_sql/autogen/agents/sql_query_generation_agent.yaml b/text_2_sql/autogen/agents/sql_query_generation_agent.yaml
@@ -1,3 +1,5 @@
+model:
+  gpt-4o-mini
 description:
   An agent that can generate SQL queries once given the schema and the user's question. It will run the SQL query to fetch the results. Use this agent after the SQL Schema Selection Agent has selected the correct schema.
 system_message:
@@ -18,3 +20,5 @@ system_message:
   The complete entity relationship graph shows you all the entities and their relationships. You can use this information to get a better understanding of the schema and the relationships between the entities and request more schema information if needed.
 
   Always run any SQL query you generate to return the results."
+tools:
+  - sql_query_execution_tool
diff --git a/text_2_sql/autogen/agents/sql_schema_selection_agent.yaml b/text_2_sql/autogen/agents/sql_schema_selection_agent.yaml
@@ -1,3 +1,5 @@
+model:
+  gpt-4o-mini
 description:
   An agent that can take a user's question and extract the schema of a view or table in the SQL Database by selecting the most relevant entity based on the search term.
 
@@ -6,3 +8,5 @@ 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.
+tools:
+  - sql_get_entity_schemas_tool
diff --git a/text_2_sql/autogen/qna_agents.py b/text_2_sql/autogen/qna_agents.py
diff --git a/text_2_sql/autogen/selector.py b/text_2_sql/autogen/selector.py
@@ -1,14 +1,14 @@
 from autogen_agentchat.task import TextMentionTermination
 from autogen_agentchat.teams import SelectorGroupChat
 from utils.models import MINI_MODEL
+from utils.agent_creator import AgentCreator
 
-from sql_tools_and_agents import (
-    SQL_QUERY_GENERATION_AGENT,
-    SQL_SCHEMA_SELECTION_AGENT,
-    SQL_QUERY_CORRECTION_AGENT,
-    SQL_QUERY_CACHE_AGENT,
-)
-from qna_agents import ANSWER_AGENT, QUESTION_DECOMPOSITION_AGENT
+SQL_QUERY_GENERATION_AGENT = AgentCreator.create("sql_query_generation_agent")
+SQL_SCHEMA_SELECTION_AGENT = AgentCreator.create("sql_schema_selection_agent")
+SQL_QUERY_CORRECTION_AGENT = AgentCreator.create("sql_query_correction_agent")
+SQL_QUERY_CACHE_AGENT = AgentCreator.create("sql_query_cache_agent")
+ANSWER_AGENT = AgentCreator.create("answer_agent")
+QUESTION_DECOMPOSITION_AGENT = AgentCreator.create("question_decomposition_agent")
 
 termination = TextMentionTermination("TERMINATE")
 SELECTOR = SelectorGroupChat(
diff --git a/text_2_sql/autogen/sql_tools_and_agents.py b/text_2_sql/autogen/sql_tools_and_agents.py
diff --git a/text_2_sql/autogen/utils/agent_creator.py b/text_2_sql/autogen/utils/agent_creator.py
@@ -0,0 +1,63 @@
+import yaml
+from autogen_core.components.tools import FunctionTool
+from autogen_agentchat.agents import ToolUseAssistantAgent
+from utils.sql_utils import (
+    query_execution,
+    get_entity_schemas,
+    fetch_queries_from_cache,
+)
+from utils.models import MINI_MODEL
+
+
+class AgentCreator:
+    def load_agent_file(cls, name):
+        with open(f"agents/{name.lower()}.yaml", "r") as file:
+            file = yaml.safe_load(file)
+
+        return file
+
+    @classmethod
+    def get_model(cls, model_name):
+        if model_name == "gpt-4o-mini":
+            return MINI_MODEL
+        else:
+            raise ValueError(f"Model {model_name} not found")
+
+    @classmethod
+    def get_tool(cls, tool_name):
+        if tool_name == "query_execution":
+            return FunctionTool(
+                query_execution,
+                description="Runs an SQL query against the SQL Database to extract information",
+            )
+        elif tool_name == "get_entity_schemas":
+            return FunctionTool(
+                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 == "fetch_queries_from_cache":
+            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")
+
+    @classmethod
+    def create(cls, name: str):
+        agent_file = cls.load_agent_file(name)
+
+        tools = []
+        if "tools" in agent_file and len(agent_file["tools"]):
+            for tool in agent_file["tools"]:
+                tools.append(cls.get_tool(tool))
+
+        agent = ToolUseAssistantAgent(
+            name=agent_file["name"],
+            registered_tools=[],
+            model_client=cls.get_model(agent_file["model"]),
+            description=agent_file["description"],
+            system_message=agent_file["system_message"],
+        )
+
+        return agent
diff --git a/text_2_sql/autogen/utils/prompts.py b/text_2_sql/autogen/utils/prompts.py

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,5 @@`
	`1`	`+model:`
	`2`	`+ gpt-4o-mini`
`1`	`3`	`description:`
`2`	`4`	`An agent that takes the final results from the SQL query and writes the answer to the user's question`
`3`	`5`	`system_message:`