Reduce temperature

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -1,6 +1,10 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-from autogen_agentchat.conditions import TextMentionTermination, MaxMessageTermination
+from autogen_agentchat.conditions import (
+    TextMentionTermination,
+    MaxMessageTermination,
+    SourceMatchTermination,
+)
 from autogen_agentchat.teams import SelectorGroupChat
 from autogen_text_2_sql.creators.llm_model_creator import LLMModelCreator
 from autogen_text_2_sql.creators.llm_agent_creator import LLMAgentCreator
@@ -89,7 +93,11 @@ def agents(self):
     @property
     def termination_condition(self):
         """Define the termination condition for the chat."""
-        termination = TextMentionTermination("TERMINATE") | MaxMessageTermination(20)
+        termination = (
+            TextMentionTermination("TERMINATE")
+            | MaxMessageTermination(20)
+            | SourceMatchTermination(["answer_agent"])
+        )
         return termination
 
     @staticmethod

text_2_sql/autogen/src/autogen_text_2_sql/creators/llm_model_creator.py

@@ -1,8 +1,9 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from autogen_ext.models import AzureOpenAIChatCompletionClient
-from text_2_sql_core.connectors.factory import ConnectorFactory
+from text_2_sql_core.utils.environment import IdentityType, get_identity_type
 
+from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 import os
 import dotenv
 
@@ -27,12 +28,32 @@ def get_model(cls, model_name: str) -> AzureOpenAIChatCompletionClient:
         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,
-        ) = ConnectorFactory.get_open_ai_connector().get_authentication_properties()
+        token_provider, api_key = cls.get_authentication_properties()
         return AzureOpenAIChatCompletionClient(
             azure_deployment=os.environ["OpenAI__MiniCompletionDeployment"],
             model=os.environ["OpenAI__MiniCompletionDeployment"],
@@ -45,14 +66,12 @@ def gpt_4o_mini_model(cls) -> AzureOpenAIChatCompletionClient:
                 "function_calling": True,
                 "json_output": True,
             },
+            temperature=0,
         )
 
     @classmethod
     def gpt_4o_model(cls) -> AzureOpenAIChatCompletionClient:
-        (
-            token_provider,
-            api_key,
-        ) = ConnectorFactory.get_open_ai_connector().get_authentication_properties()
+        token_provider, api_key = cls.get_authentication_properties()
         return AzureOpenAIChatCompletionClient(
             azure_deployment=os.environ["OpenAI__CompletionDeployment"],
             model=os.environ["OpenAI__CompletionDeployment"],
@@ -65,4 +84,5 @@ def gpt_4o_model(cls) -> AzureOpenAIChatCompletionClient:
                 "function_calling": True,
                 "json_output": True,
             },
+            temperature=0,
         )

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/answer_agent.yaml

@@ -16,5 +16,4 @@ system_message:
   }
 
   Title is the entity name of the schema, chunk is the result of the SQL query and reference is the SQL query used to generate the answer.
-
-  End your answer with 'TERMINATE'"
+  "