Update agent logic

Author: BenConstable9

…t_2_sql/autogen/agentic_text_2_sql.ipynb … 5 - Agentic Vector Based Text2SQL.ipynbtext_2_sql/autogen/agentic_text_2_sql.ipynb renamed to text_2_sql/autogen/Iteration 5 - Agentic Vector Based Text2SQL.ipynb text_2_sql/autogen/agentic_text_2_sql.ipynb renamed to text_2_sql/autogen/Iteration 5 - Agentic Vector Based Text2SQL.ipynb

@@ -1,5 +1,5 @@
 model:
-  gpt-4o-mini
+  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:

text_2_sql/autogen/agents/llm_agents/answer_agent.yaml

@@ -1,5 +1,5 @@
 model:
-  gpt-4o-mini
+  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."
 system_message:

text_2_sql/autogen/agents/llm_agents/question_decomposition_agent.yaml

@@ -1,5 +1,5 @@
 model:
-  gpt-4o-mini
+  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:

text_2_sql/autogen/agents/llm_agents/sql_query_correction_agent.yaml

@@ -1,5 +1,5 @@
 model:
-  gpt-4o-mini
+  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:

text_2_sql/autogen/agents/llm_agents/sql_query_generation_agent.yaml

@@ -1,5 +1,5 @@
 model:
-  gpt-4o-mini
+  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.
 

text_2_sql/autogen/agents/llm_agents/sql_schema_selection_agent.yaml

@@ -4,28 +4,56 @@
 from autogen_core.components.tools import FunctionTool
 from autogen_agentchat.agents import AssistantAgent
 from utils.sql import query_execution, get_entity_schemas, query_validation
-from utils.models import MINI_MODEL
+from utils.models import GPT_4O_MINI_MODEL, GPT_4O_MODEL
 from jinja2 import Template
 from datetime import datetime
+from autogen_ext.models import AzureOpenAIChatCompletionClient
 
 
 class LLMAgentCreator:
     @classmethod
-    def load_agent_file(cls, name):
+    def load_agent_file(cls, name: str) -> dict:
+        """Loads the agent file based on the agent name.
+
+        Args:
+        ----
+            name (str): The name of the agent to load.
+
+        Returns:
+        -------
+            dict: The agent file."""
         with open(f"./agents/llm_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
+    def get_model(cls, model_name: str) -> AzureOpenAIChatCompletionClient:
+        """Retrieves the model based on the model name.
+
+        Args:
+        ----
+            model_name (str): The name of the model to retrieve.
+
+        Returns:
+            AzureOpenAIChatCompletionClient: The model client."""
+        if model_name == "4o-mini":
+            return GPT_4O_MINI_MODEL
+        elif model_name == "4o":
+            return GPT_4O_MODEL
         else:
             raise ValueError(f"Model {model_name} not found")
 
     @classmethod
-    def get_tool(cls, tool_name):
+    def get_tool(cls, tool_name: str) -> FunctionTool:
+        """Retrieves the tool based on the tool name.
+
+        Args:
+        ----
+            tool_name (str): The name of the tool to retrieve.
+
+        Returns:
+            FunctionTool: The tool."""
         if tool_name == "sql_query_execution_tool":
             return FunctionTool(
                 query_execution,
@@ -50,15 +78,38 @@ def get_tool(cls, tool_name):
             raise ValueError(f"Tool {tool_name} not found")
 
     @classmethod
-    def get_property_and_render_parameters(cls, agent_file, property, parameters):
+    def get_property_and_render_parameters(
+        cls, agent_file: dict, property: str, parameters: dict
+    ) -> str:
+        """Gets the property from the agent file and renders the parameters.
+
+        Args:
+        ----
+            agent_file (dict): The agent file.
+            property (str): The property to retrieve.
+            parameters (dict): The parameters to render.
+
+        Returns:
+        -------
+            str: The rendered property."""
         unrendered_parameters = agent_file[property]
 
         rendered_template = Template(unrendered_parameters).render(parameters)
 
         return rendered_template
 
     @classmethod
-    def create(cls, name: str, **kwargs):
+    def create(cls, name: str, **kwargs) -> AssistantAgent:
+        """Creates an assistant agent based on the agent name.
+
+        Args:
+        ----
+            name (str): The name of the agent to create.
+            **kwargs: The parameters to render.
+
+        Returns:
+        -------
+            AssistantAgent: The assistant agent."""
         agent_file = cls.load_agent_file(name)
 
         tools = []

text_2_sql/autogen/utils/llm_agent_creator.py

@@ -13,7 +13,7 @@
 #     DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
 # )
 
-MINI_MODEL = AzureOpenAIChatCompletionClient(
+GPT_4O_MINI_MODEL = AzureOpenAIChatCompletionClient(
     azure_deployment=os.environ["OpenAI__MiniCompletionDeployment"],
     model=os.environ["OpenAI__MiniCompletionDeployment"],
     api_version="2024-08-01-preview",
@@ -27,3 +27,18 @@
         "json_output": True,
     },
 )
+
+GPT_4O_MODEL = AzureOpenAIChatCompletionClient(
+    azure_deployment=os.environ["OpenAI__CompletionDeployment"],
+    model=os.environ["OpenAI__CompletionDeployment"],
+    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=os.environ["OpenAI__ApiKey"],  # For key-based authentication.
+    model_capabilities={
+        "vision": False,
+        "function_calling": True,
+        "json_output": True,
+    },
+)

text_2_sql/autogen/utils/models.py

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