Add some more agents

BenConstable9 · BenConstable9 · commit 372e9951c7e0 · 2024-11-20T09:21:17.000Z
diff --git a/text_2_sql/autogen/agentic_text_2_sql.ipynb b/text_2_sql/autogen/agentic_text_2_sql.ipynb
@@ -0,0 +1,70 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import dotenv\n",
+        "import logging\n",
+        "from selector import SELECTOR"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "logging.basicConfig(level=logging.INFO)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "dotenv.load_dotenv()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "await SELECTOR.run(task=\"What are the top sales this month?\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": []
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
+      "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython3",
+      "version": "3.12.6"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 2
+}
diff --git a/text_2_sql/autogen/agents.py b/text_2_sql/autogen/agents.py
diff --git a/text_2_sql/autogen/prompts/answer_agent.yaml b/text_2_sql/autogen/prompts/answer_agent.yaml
@@ -0,0 +1,6 @@
+description:
+  An agent that takes the final results from the SQL query and writes the answer to the user's question
+system_message:
+  Write a data-driven answer that directly addresses the user's question. Use the results from the SQL query to provide the answer. Do not make up or guess the answer.
+
+  End your answer with 'TERMINATE'
diff --git a/text_2_sql/autogen/prompts/question_decomposition_agent.yaml b/text_2_sql/autogen/prompts/question_decomposition_agent.yaml
@@ -0,0 +1,8 @@
+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.
+
+  Only use this agent once per user question and after the 'Query Cache Agent' if the results are none.
+system_message:
+  You are a helpful AI Assistant that specialises in decomposing complex user questions into smaller parts that can be used in SQL queries.
+
+  Break down the user's question into smaller parts that can be used in SQL queries.
diff --git a/text_2_sql/autogen/prompts/sql_query_cache_agent.yaml b/text_2_sql/autogen/prompts/sql_query_cache_agent.yaml
@@ -0,0 +1,8 @@
+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.
+
+  Only use this agent once per user question and at the beginning of the conversation.
+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.
diff --git a/text_2_sql/autogen/prompts/sql_query_correction_agent.yaml b/text_2_sql/autogen/prompts/sql_query_correction_agent.yaml
@@ -0,0 +1,8 @@
+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:
+  You are a helpful AI Assistant that specialises in correcting invalid SQL queries or queries that do not return the expected results.
+
+  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'
diff --git a/text_2_sql/autogen/prompts/sql_query_generation_agent.yaml b/text_2_sql/autogen/prompts/sql_query_generation_agent.yaml
@@ -0,0 +1,18 @@
+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:
+  "You are a helpful AI Assistant that specialises in writing and executing SQL Queries to answer a given user's question.
+
+  Only use schema / column information provided when constructing a SQL query. Do not use any other entities and columns in your SQL query, other than those defined above.
+  Do not makeup or guess column names.
+
+  The target database engine is SQL Server, SQL queries must be able compatible to run on Target Engine.
+  You must only provide SELECT SQL queries.
+  For a given entity, use the 'SelectFromEntity' property returned in the schema in the SELECT FROM part of the SQL query. If the property is {{'SelectFromEntity': 'test_schema.test_table'}}, the select statement will be formulated from 'SELECT <VALUES> FROM test_schema.test_table WHERE <CONDITION>.
+
+  If you don't know how the value is formatted in a column, run a query against the column to get the unique values that might match your query or use the corresponding lookup values. Use a 'like' operator to match the values, rather than a direct match unless you are sure of the value.
+  Some columns in the schema may have the properties 'AllowedValues' or 'SampleValues'. Use these values to determine the possible values that can be used in the SQL query.
+
+  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."
diff --git a/text_2_sql/autogen/prompts/sql_schema_selection_agent.yaml b/text_2_sql/autogen/prompts/sql_schema_selection_agent.yaml
@@ -0,0 +1,6 @@
+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.
+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.
diff --git a/text_2_sql/autogen/qna_agents.py b/text_2_sql/autogen/qna_agents.py
@@ -0,0 +1,19 @@
+from autogen_agentchat.agents import ToolUseAssistantAgent
+from utils.models import MINI_MODEL
+from utils.prompts import load_system_message, load_description
+
+ANSWER_AGENT = ToolUseAssistantAgent(
+    name="Answer_Agent",
+    registered_tools=[],
+    model_client=MINI_MODEL,
+    description=load_description("answer_agent"),
+    system_message=load_system_message("answer_agent"),
+)
+
+QUESTION_DECOMPOSITION_AGENT = ToolUseAssistantAgent(
+    name="Question_Decomposition_Agent",
+    registered_tools=[],
+    model_client=MINI_MODEL,
+    description=load_description("question_decomposition_agent"),
+    system_message=load_system_message("question_decomposition_agent"),
+)
diff --git a/text_2_sql/autogen/requirements.txt b/text_2_sql/autogen/requirements.txt
@@ -1,6 +1,6 @@
-autogen-core
-autogen-agentchat
-autogen-ext[openai,azure]
+autogen-core==0.4.0.dev6
+autogen-agentchat==0.4.0.dev6
+autogen-ext[openai,azure]==0.4.0.dev6
 aioodbc
 azure-search
 azure-search-documents==11.6.0b5
diff --git a/text_2_sql/autogen/selector.py b/text_2_sql/autogen/selector.py
@@ -0,0 +1,25 @@
+from autogen_agentchat.task import TextMentionTermination
+from autogen_agentchat.teams import SelectorGroupChat
+from utils.models import MINI_MODEL
+
+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
+
+termination = TextMentionTermination("TERMINATE")
+SELECTOR = SelectorGroupChat(
+    [
+        SQL_QUERY_GENERATION_AGENT,
+        SQL_SCHEMA_SELECTION_AGENT,
+        SQL_QUERY_CORRECTION_AGENT,
+        SQL_QUERY_CACHE_AGENT,
+        ANSWER_AGENT,
+        QUESTION_DECOMPOSITION_AGENT,
+    ],
+    model_client=MINI_MODEL,
+    termination_condition=termination,
+)
diff --git a/text_2_sql/autogen/sql_tools_and_agents.py b/text_2_sql/autogen/sql_tools_and_agents.py
@@ -6,6 +6,7 @@
     fetch_queries_from_cache,
 )
 from utils.models import MINI_MODEL
+from utils.prompts import load_system_message, load_description
 
 SQL_QUERY_EXECUTION_TOOL = FunctionTool(
     query_execution,
@@ -22,34 +23,34 @@
     description="Fetch the pre-assembled queries, and potential results from the cache based on the user's question.",
 )
 
-SQL_QUERY_AGENT = ToolUseAssistantAgent(
-    name="SQL_Query_Agent",
+SQL_QUERY_GENERATION_AGENT = ToolUseAssistantAgent(
+    name="SQL_Query_Generation_Agent",
     registered_tools=[SQL_QUERY_EXECUTION_TOOL],
     model_client=MINI_MODEL,
-    description="An agent that can take a user's question and run an SQL query against the SQL Database to extract information",
-    system_message="You are a helpful AI assistant. Solve tasks using your tools. Specifically, you can take into consideration the user's request and run an SQL query against the SQL Database to extract information.",
+    description=load_description("sql_query_generation_agent"),
+    system_message=load_system_message("sql_query_generation_agent"),
 )
 
-SQL_SCHEMA_EXTRACTION_AGENT = ToolUseAssistantAgent(
-    name="SQL_Schema_Extraction_Agent",
+SQL_SCHEMA_SELECTION_AGENT = ToolUseAssistantAgent(
+    name="SQL_Schema_Selection_Agent",
     registered_tools=[SQL_GET_ENTITY_SCHEMAS_TOOL],
     model_client=MINI_MODEL,
-    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",
-    system_message="You are a helpful AI assistant. Solve tasks using your tools. Specifically, you can take into consideration the user's request and extract the schema of a view or table in the SQL Database by selecting the most relevant entity based on the search term.",
+    description=load_description("sql_schema_selection_agent"),
+    system_message=load_system_message("sql_schema_selection_agent"),
 )
 
 SQL_QUERY_CORRECTION_AGENT = ToolUseAssistantAgent(
     name="SQL_Query_Correction_Agent",
-    registered_tools=[SQL_QUERY_EXECUTION_TOOL],
+    registered_tools=[SQL_QUERY_EXECUTION_TOOL, SQL_GET_ENTITY_SCHEMAS_TOOL],
     model_client=MINI_MODEL,
-    description="An agent that will look at the SQL query, SQL query results and correct any mistakes in the SQL query",
-    system_message="",
+    description=load_description("sql_query_correction_agent"),
+    system_message=load_system_message("sql_query_correction_agent"),
 )
 
 SQL_QUERY_CACHE_AGENT = ToolUseAssistantAgent(
     name="SQL_Query_Cache_Agent",
     registered_tools=[SQL_QUERY_CACHE_TOOLS],
     model_client=MINI_MODEL,
-    description="An agent that will fetch the queries from the cache based on the user's question.",
-    system_message="",
+    description=load_description("sql_query_cache_agent"),
+    system_message=load_system_message("sql_query_cache_agent"),
 )
diff --git a/text_2_sql/autogen/utils/models.py b/text_2_sql/autogen/utils/models.py
@@ -1,19 +1,23 @@
 from autogen_ext.models import AzureOpenAIChatCompletionClient
-from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+
+# from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 import os
+import dotenv
 
-# Create the token provider
-token_provider = get_bearer_token_provider(
-    DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
-)
+dotenv.load_dotenv()
+
+# # Create the token provider
+# token_provider = get_bearer_token_provider(
+#     DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
+# )
 
 MINI_MODEL = AzureOpenAIChatCompletionClient(
-    model="{your-azure-deployment}",
-    api_version="2024-06-01",
+    model=os.environ["OpenAI__MiniCompletionDeployment"],
+    api_version="2024-08-01-preview",
     azure_endpoint=os.environ["OpenAI__Endpoint"],
-    # Optional if you choose key-based authentication.
-    azure_ad_token_provider=token_provider,
-    # api_key="sk-...", # For key-based authentication.
+    # # Optional if you choose key-based authentication.
+    # azure_ad_token_provider=token_provider,
+    api_key=os.environ["OpenAI__ApiKey"],  # For key-based authentication.
     model_capabilities={
         "vision": False,
         "function_calling": True,
diff --git a/text_2_sql/autogen/utils/prompts.py b/text_2_sql/autogen/utils/prompts.py
@@ -0,0 +1,15 @@
+import yaml
+
+
+def load_system_message(name):
+    with open(f"prompts/{name.lower()}.yaml", "r") as file:
+        file = yaml.safe_load(file)
+
+    return file["system_message"]
+
+
+def load_description(name):
+    with open(f"prompts/{name.lower()}.yaml", "r") as file:
+        file = yaml.safe_load(file)
+
+    return file["description"]
diff --git a/text_2_sql/autogen/utils/sql_utils.py b/text_2_sql/autogen/utils/sql_utils.py
@@ -31,8 +31,15 @@ async def get_entity_schemas(
 
     schemas = await run_ai_search_query(
         text,
-        ["DescriptionEmbedding"],
-        ["Entity", "EntityName", "Description", "Columns"],
+        ["DefinitionEmbedding"],
+        [
+            "Entity",
+            "EntityName",
+            "Definition",
+            "Columns",
+            "EntityRelationships",
+            "CompleteEntityRelationshipsGraph",
+        ],
         os.environ["AIService__AzureSearchOptions__Text2Sql__Index"],
         os.environ["AIService__AzureSearchOptions__Text2Sql__SemanticConfig"],
         top=3,
@@ -95,16 +102,9 @@ async def fetch_queries_from_cache(question: str) -> str:
 
     if len(cached_schemas) == 0:
         return None
-    else:
-        database = os.environ["Text2Sql__DatabaseName"]
-        for entry in cached_schemas:
-            for schema in entry["Schemas"]:
-                entity = schema["Entity"]
-                schema["SelectFromEntity"] = f"{database}.{entity}"
 
+    logging.info("Cached schemas: %s", cached_schemas)
     if PRE_RUN_QUERY_CACHE and len(cached_schemas) > 0:
-        logging.info("Cached schemas: %s", cached_schemas)
-
         # check the score
         if cached_schemas[0]["@search.reranker_score"] > 2.75:
             logging.info("Score is greater than 3")
@@ -125,7 +125,7 @@ async def fetch_queries_from_cache(question: str) -> str:
             for sql_query, sql_result in zip(sql_queries, sql_results):
                 query_result_store[sql_query["SqlQuery"]] = {
                     "result": sql_result,
-                    "schemas": sql_queries["schemas"],
+                    "schemas": sql_query["Schemas"],
                 }
 
             return query_result_store
diff --git a/text_2_sql/semantic_kernel/plugins/vector_based_sql_plugin/vector_based_sql_plugin.py b/text_2_sql/semantic_kernel/plugins/vector_based_sql_plugin/vector_based_sql_plugin.py
