Improve building of autogen

Author: BenConstable9

text_2_sql/autogen/agentic_text_2_sql.py

@@ -2,7 +2,7 @@
 # Licensed under the MIT License.
 from autogen_agentchat.task import TextMentionTermination, MaxMessageTermination
 from autogen_agentchat.teams import SelectorGroupChat
-from utils.models import MINI_MODEL
+from utils.llm_model_creator import LLMModelCreator
 from utils.llm_agent_creator import LLMAgentCreator
 import logging
 from agents.custom_agents.sql_query_cache_agent import SqlQueryCacheAgent
@@ -86,13 +86,17 @@ def selector(messages):
             and messages[-1].content is not None
         ):
             cache_result = json.loads(messages[-1].content)
-            if cache_result.get("cached_questions_and_schemas") is not None:
+            if cache_result.get(
+                "cached_questions_and_schemas"
+            ) is not None and cache_result.get("contains_pre_run_results"):
                 decision = "sql_query_correction_agent"
+            if (
+                cache_result.get("cached_questions_and_schemas") is not None
+                and cache_result.get("contains_pre_run_results") is False
+            ):
+                decision = "sql_query_generation_agent"
             else:
-                decision = "sql_schema_selection_agent"
-
-        elif messages[-1].source == "sql_query_cache_agent":
-            decision = "question_decomposition_agent"
+                decision = "question_decomposition_agent"
 
         elif messages[-1].source == "question_decomposition_agent":
             decomposition_result = json.loads(messages[-1].content)
@@ -129,7 +133,7 @@ def agentic_flow(self):
         agentic_flow = SelectorGroupChat(
             self.agents,
             allow_repeated_speaker=False,
-            model_client=MINI_MODEL,
+            model_client=LLMModelCreator.get_model("4o-mini"),
             termination_condition=self.termination_condition,
             selector_func=AgenticText2Sql.selector,
         )

text_2_sql/autogen/utils/llm_agent_creator.py

@@ -4,10 +4,9 @@
 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 GPT_4O_MINI_MODEL, GPT_4O_MODEL
+from utils.llm_model_creator import LLMModelCreator
 from jinja2 import Template
 from datetime import datetime
-from autogen_ext.models import AzureOpenAIChatCompletionClient
 
 
 class LLMAgentCreator:
@@ -27,23 +26,6 @@ def load_agent_file(cls, name: str) -> dict:
 
         return file
 
-    @classmethod
-    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: str) -> FunctionTool:
         """Retrieves the tool based on the tool name.
@@ -120,7 +102,7 @@ def create(cls, name: str, **kwargs) -> AssistantAgent:
         agent = AssistantAgent(
             name=name,
             tools=tools,
-            model_client=cls.get_model(agent_file["model"]),
+            model_client=LLMModelCreator.get_model(agent_file["model"]),
             description=cls.get_property_and_render_parameters(
                 agent_file, "description", kwargs
             ),

text_2_sql/autogen/utils/llm_model_creator.py

@@ -0,0 +1,86 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from autogen_ext.models import AzureOpenAIChatCompletionClient
+from environment import IdentityType, get_identity_type
+
+from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+import os
+import dotenv
+
+dotenv.load_dotenv()
+
+
+class LLMModelCreator:
+    @classmethod
+    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 cls.gpt_4o_mini_model()
+        elif model_name == "4o":
+            return cls.gpt_4o_model()
+        else:
+            raise ValueError(f"Model {model_name} not found")
+
+    @classmethod
+    def get_authentication_properties(cls) -> dict:
+        if get_identity_type() == IdentityType.SYSTEM_ASSIGNED:
+            # Create the token provider
+            api_key = None
+            token_provider = get_bearer_token_provider(
+                DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
+            )
+        elif get_identity_type() == IdentityType.USER_ASSIGNED:
+            # Create the token provider
+            api_key = None
+            token_provider = get_bearer_token_provider(
+                DefaultAzureCredential(
+                    managed_identity_client_id=os.environ["ClientId"]
+                ),
+                "https://cognitiveservices.azure.com/.default",
+            )
+        else:
+            token_provider = None
+            api_key = os.environ["OpenAI__ApiKey"]
+
+        return token_provider, api_key
+
+    @classmethod
+    def gpt_4o_mini_model(cls) -> AzureOpenAIChatCompletionClient:
+        token_provider, api_key = cls.get_authentication_properties()
+        return AzureOpenAIChatCompletionClient(
+            azure_deployment=os.environ["OpenAI__MiniCompletionDeployment"],
+            model=os.environ["OpenAI__MiniCompletionDeployment"],
+            api_version="2024-08-01-preview",
+            azure_endpoint=os.environ["OpenAI__Endpoint"],
+            azure_ad_token_provider=token_provider,
+            api_key=api_key,
+            model_capabilities={
+                "vision": False,
+                "function_calling": True,
+                "json_output": True,
+            },
+        )
+
+    @classmethod
+    def gpt_4o_model(cls) -> AzureOpenAIChatCompletionClient:
+        token_provider, api_key = cls.get_authentication_properties()
+        return AzureOpenAIChatCompletionClient(
+            azure_deployment=os.environ["OpenAI__CompletionDeployment"],
+            model=os.environ["OpenAI__CompletionDeployment"],
+            api_version="2024-08-01-preview",
+            azure_endpoint=os.environ["OpenAI__Endpoint"],
+            azure_ad_token_provider=token_provider,
+            api_key=api_key,
+            model_capabilities={
+                "vision": False,
+                "function_calling": True,
+                "json_output": True,
+            },
+        )