Limit number of results

Author: BenConstable9

text_2_sql/ai_search.py

@@ -17,6 +17,7 @@ async def run_ai_search_query(
     retrieval_fields: list[str],
     index_name: str,
     semantic_config: str,
+    top=5,
 ):
     """Run the AI search query."""
     identity_type = get_identity_type()
@@ -56,8 +57,7 @@ async def run_ai_search_query(
         credential=credential,
     ) as search_client:
         results = await search_client.search(
-            top=5,
-            query_type="semantic",
+            top=top,
             semantic_configuration_name=semantic_config,
             search_text=query,
             select=",".join(retrieval_fields),

text_2_sql/plugins/vector_based_sql_plugin/vector_based_sql_plugin.py

@@ -101,6 +101,7 @@ async def get_entity_schemas(
             ["Entity", "EntityName", "Description", "Columns"],
             os.environ["AIService__AzureSearchOptions__Text2Sql__Index"],
             os.environ["AIService__AzureSearchOptions__Text2Sql__SemanticConfig"],
+            top=3,
         )
 
         for schema in schemas:

text_2_sql/rag_with_vector_based_text_2_sql_query_cache.ipynb

@@ -252,6 +252,7 @@
         "        [\"Question\", \"Query\", \"Schemas\"],\n",
         "        os.environ[\"AIService__AzureSearchOptions__Text2SqlQueryCache__Index\"],\n",
         "        os.environ[\"AIService__AzureSearchOptions__Text2SqlQueryCache__SemanticConfig\"],\n",
+        "        top=1,\n",
         "    )\n",
         "\n",
         "    # Create important information prompt that contains the SQL database information.\n",