vanna 0.7.9__py3-none-any.whl → 2.0.0rc1__py3-none-any.whl

vanna/{azuresearch → legacy/azuresearch}/azuresearch_vector.py

@@ -7,16 +7,16 @@ from azure.core.credentials import AzureKeyCredential
 from azure.search.documents import SearchClient
 from azure.search.documents.indexes import SearchIndexClient
 from azure.search.documents.indexes.models import (
-  ExhaustiveKnnAlgorithmConfiguration,
-  ExhaustiveKnnParameters,
-  SearchableField,
-  SearchField,
-  SearchFieldDataType,
-  SearchIndex,
-  VectorSearch,
-  VectorSearchAlgorithmKind,
-  VectorSearchAlgorithmMetric,
-  VectorSearchProfile,
+    ExhaustiveKnnAlgorithmConfiguration,
+    ExhaustiveKnnParameters,
+    SearchableField,
+    SearchField,
+    SearchFieldDataType,
+    SearchIndex,
+    VectorSearch,
+    VectorSearchAlgorithmKind,
+    VectorSearchAlgorithmMetric,
+    VectorSearchProfile,
 )
 from azure.search.documents.models import VectorFilterMode, VectorizedQuery
 from fastembed import TextEmbedding
@@ -44,6 +44,7 @@ class AzureAISearch_VectorStore(VannaBase):
     Raises:
         ValueError: If config is None, or if 'azure_search_api_key' is not provided in the config.
     """
+
     def __init__(self, config=None):
         VannaBase.__init__(self, config=config)
 
@@ -54,7 +55,9 @@ class AzureAISearch_VectorStore(VannaBase):
                 "config is required, pass an API key, 'azure_search_api_key', in the config."
             )
 
-        azure_search_endpoint = config.get("azure_search_endpoint", "https://azcognetive.search.windows.net")
+        azure_search_endpoint = config.get(
+            "azure_search_endpoint", "https://azcognetive.search.windows.net"
+        )
         azure_search_api_key = config.get("azure_search_api_key")
 
         self.dimensions = config.get("dimensions", 384)
@@ -64,22 +67,24 @@ class AzureAISearch_VectorStore(VannaBase):
 
         self.n_results_ddl = config.get("n_results_ddl", config.get("n_results", 10))
         self.n_results_sql = config.get("n_results_sql", config.get("n_results", 10))
-        self.n_results_documentation = config.get("n_results_documentation", config.get("n_results", 10))
+        self.n_results_documentation = config.get(
+            "n_results_documentation", config.get("n_results", 10)
+        )
 
         if not azure_search_api_key:
             raise ValueError(
                 "'azure_search_api_key' is required in config to use AzureAISearch_VectorStore"
-        )
+            )
 
         self.index_client = SearchIndexClient(
             endpoint=azure_search_endpoint,
-            credential=AzureKeyCredential(azure_search_api_key)
+            credential=AzureKeyCredential(azure_search_api_key),
         )
 
         self.search_client = SearchClient(
             endpoint=azure_search_endpoint,
             index_name=self.index_name,
-            credential=AzureKeyCredential(azure_search_api_key)
+            credential=AzureKeyCredential(azure_search_api_key),
         )
 
         if self.index_name not in self._get_indexes():
@@ -87,10 +92,28 @@ class AzureAISearch_VectorStore(VannaBase):
 
     def _create_index(self) -> bool:
         fields = [
-            SearchableField(name="id", type=SearchFieldDataType.String, key=True, filterable=True),
-            SearchableField(name="document", type=SearchFieldDataType.String, searchable=True, filterable=True),
-            SearchField(name="type", type=SearchFieldDataType.String, filterable=True, searchable=True),
-            SearchField(name="document_vector", type=SearchFieldDataType.Collection(SearchFieldDataType.Single), searchable=True, vector_search_dimensions=self.dimensions, vector_search_profile_name="ExhaustiveKnnProfile"),
+            SearchableField(
+                name="id", type=SearchFieldDataType.String, key=True, filterable=True
+            ),
+            SearchableField(
+                name="document",
+                type=SearchFieldDataType.String,
+                searchable=True,
+                filterable=True,
+            ),
+            SearchField(
+                name="type",
+                type=SearchFieldDataType.String,
+                filterable=True,
+                searchable=True,
+            ),
+            SearchField(
+                name="document_vector",
+                type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
+                searchable=True,
+                vector_search_dimensions=self.dimensions,
+                vector_search_profile_name="ExhaustiveKnnProfile",
+            ),
         ]
 
         vector_search = VectorSearch(
@@ -100,7 +123,7 @@ class AzureAISearch_VectorStore(VannaBase):
                     kind=VectorSearchAlgorithmKind.EXHAUSTIVE_KNN,
                     parameters=ExhaustiveKnnParameters(
                         metric=VectorSearchAlgorithmMetric.COSINE
-                    )
+                    ),
                 )
             ],
             profiles=[
@@ -108,12 +131,14 @@ class AzureAISearch_VectorStore(VannaBase):
                     name="ExhaustiveKnnProfile",
                     algorithm_configuration_name="ExhaustiveKnn",
                 )
-            ]
+            ],
         )
 
-        index = SearchIndex(name=self.index_name, fields=fields, vector_search=vector_search)
+        index = SearchIndex(
+            name=self.index_name, fields=fields, vector_search=vector_search
+        )
         result = self.index_client.create_or_update_index(index)
-        print(f'{result.name} created')
+        print(f"{result.name} created")
 
     def _get_indexes(self) -> list:
         return [index for index in self.index_client.list_index_names()]
@@ -124,7 +149,7 @@ class AzureAISearch_VectorStore(VannaBase):
             "id": id,
             "document": ddl,
             "type": "ddl",
-            "document_vector": self.generate_embedding(ddl)
+            "document_vector": self.generate_embedding(ddl),
         }
         self.search_client.upload_documents(documents=[document])
         return id
@@ -135,32 +160,36 @@ class AzureAISearch_VectorStore(VannaBase):
             "id": id,
             "document": doc,
             "type": "doc",
-            "document_vector": self.generate_embedding(doc)
+            "document_vector": self.generate_embedding(doc),
         }
         self.search_client.upload_documents(documents=[document])
         return id
 
     def add_question_sql(self, question: str, sql: str) -> str:
-        question_sql_json = json.dumps({"question": question, "sql": sql}, ensure_ascii=False)
+        question_sql_json = json.dumps(
+            {"question": question, "sql": sql}, ensure_ascii=False
+        )
         id = deterministic_uuid(question_sql_json) + "-sql"
         document = {
             "id": id,
             "document": question_sql_json,
             "type": "sql",
-            "document_vector": self.generate_embedding(question_sql_json)
+            "document_vector": self.generate_embedding(question_sql_json),
         }
         self.search_client.upload_documents(documents=[document])
         return id
 
     def get_related_ddl(self, text: str) -> List[str]:
         result = []
-        vector_query = VectorizedQuery(vector=self.generate_embedding(text), fields="document_vector")
+        vector_query = VectorizedQuery(
+            vector=self.generate_embedding(text), fields="document_vector"
+        )
         df = pd.DataFrame(
             self.search_client.search(
                 top=self.n_results_ddl,
                 vector_queries=[vector_query],
                 select=["id", "document", "type"],
-                filter=f"type eq 'ddl'"
+                filter=f"type eq 'ddl'",
             )
         )
 
@@ -170,7 +199,9 @@ class AzureAISearch_VectorStore(VannaBase):
 
     def get_related_documentation(self, text: str) -> List[str]:
         result = []
-        vector_query = VectorizedQuery(vector=self.generate_embedding(text), fields="document_vector")
+        vector_query = VectorizedQuery(
+            vector=self.generate_embedding(text), fields="document_vector"
+        )
 
         df = pd.DataFrame(
             self.search_client.search(
@@ -178,7 +209,7 @@ class AzureAISearch_VectorStore(VannaBase):
                 vector_queries=[vector_query],
                 select=["id", "document", "type"],
                 filter=f"type eq 'doc'",
-                vector_filter_mode=VectorFilterMode.PRE_FILTER
+                vector_filter_mode=VectorFilterMode.PRE_FILTER,
             )
         )
 
@@ -189,42 +220,49 @@ class AzureAISearch_VectorStore(VannaBase):
     def get_similar_question_sql(self, question: str) -> List[str]:
         result = []
         # Vectorize the text
-        vector_query = VectorizedQuery(vector=self.generate_embedding(question), fields="document_vector")
+        vector_query = VectorizedQuery(
+            vector=self.generate_embedding(question), fields="document_vector"
+        )
         df = pd.DataFrame(
             self.search_client.search(
                 top=self.n_results_sql,
                 vector_queries=[vector_query],
                 select=["id", "document", "type"],
-                filter=f"type eq 'sql'"
+                filter=f"type eq 'sql'",
             )
         )
 
-        if len(df): # Check if there is similar query and the result is not empty
+        if len(df):  # Check if there is similar query and the result is not empty
             result = [ast.literal_eval(element) for element in df["document"].tolist()]
 
         return result
 
     def get_training_data(self) -> List[str]:
-
         search = self.search_client.search(
             search_text="*",
-            select=['id', 'document', 'type'],
-            filter=f"(type eq 'sql') or (type eq 'ddl') or (type eq 'doc')"
+            select=["id", "document", "type"],
+            filter=f"(type eq 'sql') or (type eq 'ddl') or (type eq 'doc')",
         ).by_page()
 
         df = pd.DataFrame([item for page in search for item in page])
 
         if len(df):
-            df.loc[df["type"] == "sql", "question"] = df.loc[df["type"] == "sql"]["document"].apply(lambda x: json.loads(x)["question"])
-            df.loc[df["type"] == "sql", "content"]  = df.loc[df["type"] == "sql"]["document"].apply(lambda x: json.loads(x)["sql"])
-            df.loc[df["type"] != "sql", "content"]  = df.loc[df["type"] != "sql"]["document"]
+            df.loc[df["type"] == "sql", "question"] = df.loc[df["type"] == "sql"][
+                "document"
+            ].apply(lambda x: json.loads(x)["question"])
+            df.loc[df["type"] == "sql", "content"] = df.loc[df["type"] == "sql"][
+                "document"
+            ].apply(lambda x: json.loads(x)["sql"])
+            df.loc[df["type"] != "sql", "content"] = df.loc[df["type"] != "sql"][
+                "document"
+            ]
 
             return df[["id", "question", "content", "type"]]
 
         return pd.DataFrame()
 
     def remove_training_data(self, id: str) -> bool:
-        result = self.search_client.delete_documents(documents=[{'id':id}])
+        result = self.search_client.delete_documents(documents=[{"id": id}])
         return result[0].succeeded
 
     def remove_index(self):