Update the prompts

Author: BenConstable9

text_2_sql/autogen/agentic_text_2_sql.ipynb

@@ -35,9 +35,25 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "await SELECTOR.run(task=\"What are the top sales this month?\")"
+        "result = await SELECTOR.run(task=\"What are the total number of sales within 2008?\")"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "print(result.messages[-1].content)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": []
+    },
     {
       "cell_type": "code",
       "execution_count": null,

text_2_sql/autogen/prompts/sql_query_generation_agent.yaml

@@ -3,6 +3,8 @@ description:
 system_message:
   "You are a helpful AI Assistant that specialises in writing and executing SQL Queries to answer a given user's question.
 
+  If you need more information from the user to generate the SQL query, ask the user for the information you need with a question and end your answer with 'TERMINATE'.
+
   Only use schema / column information provided when constructing a SQL query. Do not use any other entities and columns in your SQL query, other than those defined above.
   Do not makeup or guess column names.
 

text_2_sql/autogen/prompts/sql_schema_selection_agent.yaml

@@ -1,5 +1,7 @@
 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.
+
+  Call this in parallel if needed multiple times. Limit the use of this agent where possible.
 system_message:
   You are a helpful AI Assistant that specialises in selecting relevant SQL schemas to answer a given user's question.
 

text_2_sql/autogen/utils/sql_utils.py

@@ -18,6 +18,10 @@ async def get_entity_schemas(
         str,
         "The text to run a semantic search against. Relevant entities will be returned.",
     ],
+    excluded_entities: Annotated[
+        list[str],
+        "The entities to exclude from the search results. Pass the entity property of entities (e.g. 'SalesLT.Address') you already have the schemas for to avoid getting repeated entities.",
+    ] = [],
 ) -> 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.
 
@@ -50,6 +54,13 @@ async def get_entity_schemas(
         database = os.environ["Text2Sql__DatabaseName"]
         schema["SelectFromEntity"] = f"{database}.{entity}"
 
+        filtered_schemas = []
+        for excluded_entity in excluded_entities:
+            if excluded_entity.lower() == entity.lower():
+                logging.info("Excluded entity: %s", excluded_entity)
+            else:
+                filtered_schemas.append(schema)
+
     return json.dumps(schemas, default=str)
 
 
@@ -92,7 +103,7 @@ async def fetch_queries_from_cache(question: str) -> str:
     cached_schemas = await run_ai_search_query(
         question,
         ["QuestionEmbedding"],
-        ["Question", "SqlQueryDecomposition", "Schemas"],
+        ["Question", "SqlQueryDecomposition"],
         os.environ["AIService__AzureSearchOptions__Text2SqlQueryCache__Index"],
         os.environ["AIService__AzureSearchOptions__Text2SqlQueryCache__SemanticConfig"],
         top=1,
@@ -101,7 +112,7 @@ async def fetch_queries_from_cache(question: str) -> str:
     )
 
     if len(cached_schemas) == 0:
-        return None
+        return {"cached_questions": None}
 
     logging.info("Cached schemas: %s", cached_schemas)
     if PRE_RUN_QUERY_CACHE and len(cached_schemas) > 0: