Merge branch 'feature/structured-outputs' into fix/model-flow-errors

Author: BenConstable9

text_2_sql/.env.example

@@ -5,6 +5,7 @@ Text2Sql__DatabaseEngine=<DatabaseEngine> # TSQL or Postgres or Snowflake or Dat
 Text2Sql__UseQueryCache=<Determines if the Query Cache will be used to speed up query generation. Defaults to True.> # True or False
 Text2Sql__PreRunQueryCache=<Determines if the results from the Query Cache will be pre-run to speed up answer generation. Defaults to True.> # True or False
 Text2Sql__UseColumnValueStore=<Determines if the Column Value Store will be used for schema selection Defaults to True.> # True or False
+Text2Sql__GenerateFollowUpQuestions=<Determines if follow up questions will be generated. Defaults to True.> # True or False
 
 # Open AI Connection Details
 OpenAI__CompletionDeployment=<openAICompletionDeploymentId. Used for data dictionary creator>

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

@@ -7,6 +7,10 @@
 from autogen_text_2_sql.creators.llm_model_creator import LLMModelCreator
 from jinja2 import Template
 import logging
+from text_2_sql_core.structured_outputs import (
+    AnswerAgentWithFollowUpQuestionsAgentOutput,
+    UserMessageRewriteAgentOutput,
+)
 
 
 class LLMAgentCreator:
@@ -106,10 +110,20 @@ def create(cls, name: str, **kwargs) -> AssistantAgent:
             for tool in agent_file["tools"]:
                 tools.append(cls.get_tool(sql_helper, tool))
 
+        structured_output = None
+        if agent_file.get("structured_output", False):
+            # Import the structured output agent
+            if name == "answer_agent_with_follow_up_questions":
+                structured_output = AnswerAgentWithFollowUpQuestionsAgentOutput
+            elif name == "user_message_rewrite_agent":
+                structured_output = UserMessageRewriteAgentOutput
+
         agent = AssistantAgent(
             name=name,
             tools=tools,
-            model_client=LLMModelCreator.get_model(agent_file["model"]),
+            model_client=LLMModelCreator.get_model(
+                agent_file["model"], structured_output=structured_output
+            ),
             description=cls.get_property_and_render_parameters(
                 agent_file, "description", kwargs
             ),

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

@@ -12,7 +12,9 @@
 
 class LLMModelCreator:
     @classmethod
-    def get_model(cls, model_name: str) -> AzureOpenAIChatCompletionClient:
+    def get_model(
+        cls, model_name: str, structured_output=None
+    ) -> AzureOpenAIChatCompletionClient:
         """Retrieves the model based on the model name.
 
         Args:
@@ -22,9 +24,9 @@ def get_model(cls, model_name: str) -> AzureOpenAIChatCompletionClient:
         Returns:
             AzureOpenAIChatCompletionClient: The model client."""
         if model_name == "4o-mini":
-            return cls.gpt_4o_mini_model()
+            return cls.gpt_4o_mini_model(structured_output=structured_output)
         elif model_name == "4o":
-            return cls.gpt_4o_model()
+            return cls.gpt_4o_model(structured_output=structured_output)
         else:
             raise ValueError(f"Model {model_name} not found")
 
@@ -46,7 +48,9 @@ def get_authentication_properties(cls) -> dict:
         return token_provider, api_key
 
     @classmethod
-    def gpt_4o_mini_model(cls) -> AzureOpenAIChatCompletionClient:
+    def gpt_4o_mini_model(
+        cls, structured_output=None
+    ) -> AzureOpenAIChatCompletionClient:
         token_provider, api_key = cls.get_authentication_properties()
         return AzureOpenAIChatCompletionClient(
             azure_deployment=os.environ["OpenAI__MiniCompletionDeployment"],
@@ -61,10 +65,11 @@ def gpt_4o_mini_model(cls) -> AzureOpenAIChatCompletionClient:
                 "json_output": True,
             },
             temperature=0,
+            response_format=structured_output,
         )
 
     @classmethod
-    def gpt_4o_model(cls) -> AzureOpenAIChatCompletionClient:
+    def gpt_4o_model(cls, structured_output=None) -> AzureOpenAIChatCompletionClient:
         token_provider, api_key = cls.get_authentication_properties()
         return AzureOpenAIChatCompletionClient(
             azure_deployment=os.environ["OpenAI__CompletionDeployment"],
@@ -79,4 +84,5 @@ def gpt_4o_model(cls) -> AzureOpenAIChatCompletionClient:
                 "json_output": True,
             },
             temperature=0,
+            response_format=structured_output,
         )

text_2_sql/text_2_sql_core/src/text_2_sql_core/payloads/interaction_payloads.py

@@ -81,6 +81,10 @@ class Source(InteractionPayloadBase):
         answer: str
         steps: list[list[str]] = Field(default_factory=list, alias="Steps")
         sources: list[Source] = Field(default_factory=list)
+        follow_up_questions: list[str] | None = Field(
+            default=None, alias="followUpQuestions"
+        )
+        assistant_state: dict | None = Field(default=None, alias="assistantState")
 
     payload_type: Literal[PayloadType.ANSWER_WITH_SOURCES] = Field(
         PayloadType.ANSWER_WITH_SOURCES, alias="payloadType"

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/answer_agent_with_follow_up_questions.yaml

@@ -0,0 +1,34 @@
+model: "4o-mini"
+description: "An agent that generates a response to a user's question."
+system_message: |
+  <role_and_objective>
+    You are Senior Data Analystm, specializing in providing data driven answers to a user's question. Use the general business use case of '{{ use_case }}' to aid understanding of the user's question. You should provide a clear and concise response based on the information obtained from the SQL queries and their results. Adopt a data-driven approach to generate the response.
+  </role_and_objective>
+
+  <system_information>
+    You are part of an overall system that provides Text2SQL and subsequent data analysis functionality only. You will be passed a result from multiple SQL queries, you must formulate a response to the user's question using this information.
+    You can assume that the SQL queries are correct and that the results are accurate.
+    You and the wider system can only generate SQL queries and process the results of these queries. You cannot access any external resources.
+    The main ability of the system is to perform natural language understanding and generate SQL queries from the user's question. These queries are then automatically run against the database and the results are passed to you.
+  </system_information>
+
+  <instructions>
+
+    Use the information obtained to generate a response to the user's question. The question has been broken down into a series of SQL queries and you need to generate a response based on the results of these queries.
+
+    Do not use any external resources to generate the response. The response should be based solely on the information provided in the SQL queries and their results.
+
+    You have no access to the internet or any other external resources. You can only use the information provided in the SQL queries and their results, to generate the response.
+
+    You can use Markdown and Markdown tables to format the response. You MUST use the information obtained from the SQL queries to generate the response.
+
+    If the user is asking about your capabilities, use the <system_information> to explain what you do.
+
+    Make sure your response directly addresses every part of the user's question.
+
+    Finally, generate 3 data driven follow-up questions based on the information obtained from the SQL queries and their results. Think carefully about what questions may arise from the data and how they can be used to further analyze the data.
+
+  </instructions>
+
+  <output_
+structured_output: true

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/sql_schema_selection_agent.yaml

@@ -96,3 +96,4 @@ system_message: |
   <key_relationships>
     {{ relationship_paths }}
   </key_relationships>
+structured_output: true

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/user_message_rewrite_agent.yaml

@@ -166,3 +166,4 @@ system_message: |
       }
       ```
   </examples>
+structured_output: true

text_2_sql/text_2_sql_core/src/text_2_sql_core/structured_outputs/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from text_2_sql_core.structured_outputs.sql_schema_selection_agent import (
+    SQLSchemaSelectionAgentOutput,
+)
+from text_2_sql_core.structured_outputs.user_message_rewrite_agent import (
+    UserMessageRewriteAgentOutput,
+)
+from text_2_sql_core.structured_outputs.answer_agent_with_follow_up_questions import (
+    AnswerAgentWithFollowUpQuestionsAgentOutput,
+)
+
+__all__ = [
+    "AnswerAgentWithFollowUpQuestionsAgentOutput",
+    "SQLSchemaSelectionAgentOutput",
+    "UserMessageRewriteAgentOutput",
+]

text_2_sql/text_2_sql_core/src/text_2_sql_core/structured_outputs/answer_agent_with_follow_up_questions.py

@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from pydantic import BaseModel
+
+
+class AnswerAgentWithFollowUpQuestionsAgentOutput(BaseModel):
+    answer: str
+    follow_up_questions: list[str]

text_2_sql/text_2_sql_core/src/text_2_sql_core/structured_outputs/user_message_rewrite_agent.py

@@ -0,0 +1,10 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from pydantic import BaseModel
+
+
+class UserMessageRewriteAgentOutput(BaseModel):
+    decomposed_user_messages: list[list[str]]
+    combination_logic: str
+    query_type: str
+    all_non_database_query: bool