Updated query store

Author: BenConstable9

deploy_ai_search/text_2_sql_query_cache.py

@@ -53,7 +53,7 @@ def get_index_fields(self) -> list[SearchableField]:
                 vector_search_profile_name=self.vector_search_profile_name,
             ),
             SearchableField(
-                name="Query", type=SearchFieldDataType.String, filterable=True
+                name="SQLQuery", type=SearchFieldDataType.String, filterable=True
             ),
             ComplexField(
                 name="Schemas",
@@ -69,13 +69,13 @@ def get_index_fields(self) -> list[SearchableField]:
                         collection=True,
                         fields=[
                             SearchableField(
-                                name="Name", type=SearchFieldDataType.String
+                                name="ColumnName", type=SearchFieldDataType.String
                             ),
                             SearchableField(
-                                name="Definition", type=SearchFieldDataType.String
+                                name="ColumnDefinition", type=SearchFieldDataType.String
                             ),
                             SearchableField(
-                                name="Type", type=SearchFieldDataType.String
+                                name="DataType", type=SearchFieldDataType.String
                             ),
                             SearchableField(
                                 name="AllowedValues",

deploy_ai_search/text_2_sql_schema_store.py

@@ -88,7 +88,7 @@ def get_index_fields(self) -> list[SearchableField]:
                 name="Columns",
                 collection=True,
                 fields=[
-                    SearchableField(name="ColumnName", type=SearchFieldDataType.String),
+                    SearchableField(name="Name", type=SearchFieldDataType.String),
                     SearchableField(name="Definition", type=SearchFieldDataType.String),
                     SearchableField(name="DataType", type=SearchFieldDataType.String),
                     SearchableField(
@@ -125,7 +125,7 @@ def get_index_fields(self) -> list[SearchableField]:
                 # This is needed to enable semantic searching against the column names as complex field types are not used.
             ),
             ComplexField(
-                name="ImmediateRelationships",
+                name="ImmediateEntityRelationships",
                 collection=True,
                 fields=[
                     SearchableField(
@@ -150,6 +150,11 @@ def get_index_fields(self) -> list[SearchableField]:
                     ),
                 ],
             ),
+            SimpleField(
+                name="CompleteEntityRelationships",
+                type=SearchFieldDataType.String,
+                collection=True,
+            ),
             SimpleField(
                 name="DateLastModified",
                 type=SearchFieldDataType.DateTimeOffset,

text_2_sql/data_dictionary/data_dictionary_creator.py

@@ -22,7 +22,7 @@ class ColumnItem(BaseModel):
     """A class to represent a column item."""
 
     name: str = Field(..., alias="Name")
-    type: str = Field(..., alias="Type")
+    data_type: str = Field(..., alias="DataType")
     definition: Optional[str] = Field(..., alias="Definition")
     distinct_values: Optional[list[any]] = Field(
         None, alias="DistinctValues", exclude=True
@@ -37,15 +37,17 @@ def from_sql_row(cls, row, columns):
         """A method to create a ColumnItem from a SQL row."""
         result = dict(zip(columns, row))
         return cls(
-            name=result["Name"], type=result["Type"], definition=result["Definition"]
+            name=result["Name"],
+            type=result["DataType"],
+            definition=result["Definition"],
         )
 
 
 class EntityItem(BaseModel):
     """A class to represent an entity item."""
 
     entity: str = Field(..., alias="Entity")
-    description: Optional[str] = Field(..., alias="Description")
+    definition: Optional[str] = Field(..., alias="Definition")
     name: str = Field(..., alias="Name", exclude=True)
     entity_schema: str = Field(..., alias="Schema", exclude=True)
     entity_name: Optional[str] = Field(default=None, alias="EntityName")
@@ -67,7 +69,7 @@ def from_sql_row(cls, row, columns):
             entity=entity,
             name=result["Entity"],
             entity_schema=result["EntitySchema"],
-            description=result["Description"],
+            definition=result["Definition"],
         )
 
 
@@ -102,20 +104,20 @@ def __init__(
     def extract_table_entities_sql_query(self) -> str:
         """An abstract property to extract table entities from a database.
 
-        Must return 3 columns: Entity, EntitySchema, Description."""
+        Must return 3 columns: Entity, EntitySchema, Definition."""
 
     @property
     @abstractmethod
     def extract_view_entities_sql_query(self) -> str:
         """An abstract property to extract view entities from a database.
 
-        Must return 3 columns: Entity, EntitySchema, Description."""
+        Must return 3 columns: Entity, EntitySchema, Definition."""
 
     @abstractmethod
     def extract_columns_sql_query(self, entity: EntityItem) -> str:
         """An abstract method to extract column information from a database.
 
-        Must return 3 columns: Name, Type, Definition."""
+        Must return 3 columns: Name, DataType, Definition."""
 
     def extract_distinct_values_sql_query(
         self, entity: EntityItem, column: ColumnItem

text_2_sql/data_dictionary/requirements.txt

@@ -3,3 +3,4 @@ azure-identity
 python-dotenv
 pydantic
 openai
+networkx

text_2_sql/data_dictionary/sql_sever_data_dictionary_creator.py

@@ -52,7 +52,7 @@ def extract_view_entities_sql_query(self) -> str:
         return """SELECT
     v.TABLE_NAME AS Entity,
     v.TABLE_SCHEMA AS EntitySchema,
-    CAST(ep.value AS NVARCHAR(500)) AS Description
+    CAST(ep.value AS NVARCHAR(500)) AS Definition
 FROM
     INFORMATION_SCHEMA.VIEWS v
 LEFT JOIN
@@ -66,7 +66,7 @@ def extract_columns_sql_query(self, entity: EntityItem) -> str:
         """A property to extract column information from a SQL Server database."""
         return f"""SELECT
     c.COLUMN_NAME AS Name,
-    c.DATA_TYPE AS Type,
+    c.DATA_TYPE AS DataType,
     CAST(ep.value AS NVARCHAR(500)) AS Definition
 FROM
     INFORMATION_SCHEMA.COLUMNS c