Work work on agentic

Author: BenConstable9

text_2_sql/autogen/agents.py

@@ -0,0 +1,18 @@
+from autogen_agentchat.agents import ToolUseAssistantAgent
+from utils.models import MINI_MODEL
+
+ANSWER_AGENT = ToolUseAssistantAgent(
+    name="Answer_Revision_Agent",
+    registered_tools=[],
+    model_client=MINI_MODEL,
+    description="An agent that takes the user's question, the outputs from the SQL queries to provide an answer to the user's question.",
+    system_message="You are a helpful AI assistant. Take the user's question and the outputs from the SQL queries to provide an answer to the user's question.",
+)
+
+QUERY_DECOMPOSITION_AGENT = ToolUseAssistantAgent(
+    name="Query_Decomposition_Agent",
+    registered_tools=[],
+    model_client=MINI_MODEL,
+    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.",
+    system_message="You are a helpful AI assistant. 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.",
+)

text_2_sql/autogen/requirements.txt

@@ -1,7 +1,9 @@
+autogen-core
 autogen-agentchat
-autogen-ext[openai]
+autogen-ext[openai,azure]
 aioodbc
 azure-search
 azure-search-documents==11.6.0b5
 azure-identity
 python-dotenv
+openai

text_2_sql/autogen/sql_tools.py

@@ -0,0 +1,55 @@
+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
+
+SQL_QUERY_EXECUTION_TOOL = FunctionTool(
+    query_execution,
+    description="Runs an SQL query against the SQL Database to extract information",
+)
+
+SQL_GET_ENTITY_SCHEMAS_TOOL = 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.",
+)
+
+SQL_QUERY_CACHE_TOOLS = FunctionTool(
+    fetch_queries_from_cache,
+    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",
+    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.",
+)
+
+SQL_SCHEMA_EXTRACTION_AGENT = ToolUseAssistantAgent(
+    name="SQL_Schema_Extraction_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.",
+)
+
+SQL_QUERY_CORRECTION_AGENT = ToolUseAssistantAgent(
+    name="SQL_Query_Correction_Agent",
+    registered_tools=[SQL_QUERY_EXECUTION_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="",
+)
+
+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="",
+)

text_2_sql/autogen/sql_tools_and_agents.py

@@ -0,0 +1,22 @@
+from autogen_ext.models import AzureOpenAIChatCompletionClient
+from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+import os
+
+# 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",
+    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.
+    model_capabilities={
+        "vision": False,
+        "function_calling": True,
+        "json_output": True,
+    },
+)

text_2_sql/autogen/sql_utils.py text_2_sql/autogen/utils/__init__.pytext_2_sql/autogen/sql_utils.py renamed to text_2_sql/autogen/utils/__init__.py

@@ -0,0 +1,133 @@
+import logging
+import os
+import aioodbc
+from typing import Annotated
+from utils.ai_search_utils import run_ai_search_query
+import json
+import asyncio
+
+USE_QUERY_CACHE = os.environ.get("Text2Sql__UseQueryCache", "False").lower() == "true"
+
+PRE_RUN_QUERY_CACHE = (
+    os.environ.get("Text2Sql__PreRunQueryCache", "False").lower() == "true"
+)
+
+
+async def get_entity_schemas(
+    text: Annotated[
+        str,
+        "The text to run a semantic search against. Relevant entities will be returned.",
+    ],
+) -> str:
+    """Gets the schema of a view or table in the SQL Database by selecting the most relevant entity based on the search term. Several entities may be returned.
+
+    Args:
+    ----
+        text (str): The text to run the search against.
+
+    Returns:
+        str: The schema of the views or tables in JSON format.
+    """
+
+    schemas = await run_ai_search_query(
+        text,
+        ["DescriptionEmbedding"],
+        ["Entity", "EntityName", "Description", "Columns"],
+        os.environ["AIService__AzureSearchOptions__Text2Sql__Index"],
+        os.environ["AIService__AzureSearchOptions__Text2Sql__SemanticConfig"],
+        top=3,
+    )
+
+    for schema in schemas:
+        entity = schema["Entity"]
+        database = os.environ["Text2Sql__DatabaseName"]
+        schema["SelectFromEntity"] = f"{database}.{entity}"
+
+    return json.dumps(schemas, default=str)
+
+
+async def query_execution(sql_query: str) -> list[dict]:
+    """Run the SQL query against the database.
+
+    Args:
+    ----
+        sql_query (str): The SQL query to run against the database.
+
+    Returns:
+    -------
+        list[dict]: The results of the SQL query.
+    """
+    connection_string = os.environ["Text2Sql__DatabaseConnectionString"]
+    async with await aioodbc.connect(dsn=connection_string) as sql_db_client:
+        async with sql_db_client.cursor() as cursor:
+            await cursor.execute(sql_query)
+
+            columns = [column[0] for column in cursor.description]
+
+            rows = await cursor.fetchall()
+            results = [dict(zip(columns, returned_row)) for returned_row in rows]
+
+    logging.debug("Results: %s", results)
+    return results
+
+
+async def fetch_queries_from_cache(question: str) -> str:
+    """Fetch the queries from the cache based on the question.
+
+    Args:
+    ----
+        question (str): The question to use to fetch the queries.
+
+    Returns:
+    -------
+        str: The formatted string of the queries fetched from the cache. This is injected into the prompt.
+    """
+    cached_schemas = await run_ai_search_query(
+        question,
+        ["QuestionEmbedding"],
+        ["Question", "SqlQueryDecomposition", "Schemas"],
+        os.environ["AIService__AzureSearchOptions__Text2SqlQueryCache__Index"],
+        os.environ["AIService__AzureSearchOptions__Text2SqlQueryCache__SemanticConfig"],
+        top=1,
+        include_scores=True,
+        minimum_score=1.5,
+    )
+
+    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}"
+
+    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")
+
+            sql_queries = cached_schemas[0]["SqlQueryDecomposition"]
+            query_result_store = {}
+
+            query_tasks = []
+
+            for sql_query in sql_queries:
+                logging.info("SQL Query: %s", sql_query)
+
+                # Run the SQL query
+                query_tasks.append(query_execution(sql_query["SqlQuery"]))
+
+            sql_results = await asyncio.gather(*query_tasks)
+
+            for sql_query, sql_result in zip(sql_queries, sql_results):
+                query_result_store[sql_query["SqlQuery"]] = {
+                    "result": sql_result,
+                    "schemas": sql_queries["schemas"],
+                }
+
+            return query_result_store
+
+    return {"cached_questions": cached_schemas}