vanna 0.7.1__py3-none-any.whl → 0.7.3__py3-none-any.whl

vanna/azuresearch/azuresearch_vector.py

@@ -186,10 +186,10 @@ class AzureAISearch_VectorStore(VannaBase):
             result = df["document"].tolist()
         return result
 
-    def get_similar_question_sql(self, text: str) -> List[str]:
+    def get_similar_question_sql(self, question: str) -> List[str]:
         result = []
         # Vectorize the text
-        vector_query = VectorizedQuery(vector=self.generate_embedding(text), 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,

vanna/faiss/__init__.py

@@ -0,0 +1 @@
+from .faiss import FAISS

vanna/faiss/faiss.py

@@ -0,0 +1,176 @@
+import os 
+import json
+import uuid
+from typing import List, Dict, Any
+
+import faiss
+import numpy as np
+import pandas as pd
+
+from ..base import VannaBase
+from ..exceptions import DependencyError
+
+class FAISS(VannaBase):
+    def __init__(self, config=None):
+        if config is None:
+            config = {}
+        
+        VannaBase.__init__(self, config=config)
+        
+        try:
+            import faiss
+        except ImportError:
+            raise DependencyError(
+                "FAISS is not installed. Please install it with 'pip install faiss-cpu' or 'pip install faiss-gpu'"
+            )
+
+        try:
+            from sentence_transformers import SentenceTransformer
+        except ImportError:
+            raise DependencyError(
+                "SentenceTransformer is not installed. Please install it with 'pip install sentence-transformers'."
+            )
+        
+        self.path = config.get("path", ".")
+        self.embedding_dim = config.get('embedding_dim', 384)
+        self.n_results_sql = config.get('n_results_sql', config.get("n_results", 10))
+        self.n_results_ddl = config.get('n_results_ddl', config.get("n_results", 10))
+        self.n_results_documentation = config.get('n_results_documentation', config.get("n_results", 10))
+        self.curr_client = config.get("client", "persistent")
+
+        if self.curr_client == 'persistent':
+            self.sql_index = self._load_or_create_index('sql_index.faiss')
+            self.ddl_index = self._load_or_create_index('ddl_index.faiss')
+            self.doc_index = self._load_or_create_index('doc_index.faiss')
+        elif self.curr_client == 'in-memory':
+            self.sql_index = faiss.IndexFlatL2(self.embedding_dim)
+            self.ddl_index = faiss.IndexFlatL2(self.embedding_dim)
+            self.doc_index = faiss.IndexFlatL2(self.embedding_dim)
+        elif isinstance(self.curr_client, list) and len(self.curr_client) == 3 and all(isinstance(idx, faiss.Index) for idx in self.curr_client):
+            self.sql_index = self.curr_client[0]
+            self.ddl_index = self.curr_client[1]
+            self.doc_index = self.curr_client[2]
+        else:
+            raise ValueError(f"Unsupported storage type was set in config: {self.curr_client}")
+
+        self.sql_metadata: List[Dict[str, Any]] = self._load_or_create_metadata('sql_metadata.json')
+        self.ddl_metadata: List[Dict[str, str]] = self._load_or_create_metadata('ddl_metadata.json')
+        self.doc_metadata: List[Dict[str, str]] = self._load_or_create_metadata('doc_metadata.json')
+
+        model_name = config.get('embedding_model', 'all-MiniLM-L6-v2')
+        self.embedding_model = SentenceTransformer(model_name)
+
+    def _load_or_create_index(self, filename):
+        filepath = os.path.join(self.path, filename)
+        if os.path.exists(filepath):
+            return faiss.read_index(filepath)
+        return faiss.IndexFlatL2(self.embedding_dim)
+
+    def _load_or_create_metadata(self, filename):
+        filepath = os.path.join(self.path, filename)
+        if os.path.exists(filepath):
+            with open(filepath, 'r') as f:
+                return json.load(f)
+        return []
+
+    def _save_index(self, index, filename):
+        if self.curr_client == 'persistent':
+            filepath = os.path.join(self.path, filename)
+            faiss.write_index(index, filepath)
+
+    def _save_metadata(self, metadata, filename):
+        if self.curr_client == 'persistent':
+            filepath = os.path.join(self.path, filename)
+            with open(filepath, 'w') as f:
+                json.dump(metadata, f)
+
+    def generate_embedding(self, data: str, **kwargs) -> List[float]:
+        embedding = self.embedding_model.encode(data)
+        assert embedding.shape[0] == self.embedding_dim, \
+            f"Embedding dimension mismatch: expected {self.embedding_dim}, got {embedding.shape[0]}"
+        return embedding.tolist()
+
+    def _add_to_index(self, index, metadata_list, text, extra_metadata=None) -> str:
+        embedding = self.generate_embedding(text)
+        index.add(np.array([embedding], dtype=np.float32))
+        entry_id = str(uuid.uuid4())
+        metadata_list.append({"id": entry_id, **(extra_metadata or {})})
+        return entry_id
+    
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        entry_id = self._add_to_index(self.sql_index, self.sql_metadata, question + " " + sql, {"question": question, "sql": sql})
+        self._save_index(self.sql_index, 'sql_index.faiss')
+        self._save_metadata(self.sql_metadata, 'sql_metadata.json')
+        return entry_id
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        entry_id = self._add_to_index(self.ddl_index, self.ddl_metadata, ddl, {"ddl": ddl})
+        self._save_index(self.ddl_index, 'ddl_index.faiss')
+        self._save_metadata(self.ddl_metadata, 'ddl_metadata.json')
+        return entry_id
+
+    def add_documentation(self, documentation: str, **kwargs) -> str:
+        entry_id = self._add_to_index(self.doc_index, self.doc_metadata, documentation, {"documentation": documentation})
+        self._save_index(self.doc_index, 'doc_index.faiss')
+        self._save_metadata(self.doc_metadata, 'doc_metadata.json')
+        return entry_id
+
+    def _get_similar(self, index, metadata_list, text, n_results) -> list:
+        embedding = self.generate_embedding(text)
+        D, I = index.search(np.array([embedding], dtype=np.float32), k=n_results)
+        return [] if len(I[0]) == 0 or I[0][0] == -1 else [metadata_list[i] for i in I[0]]
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> list:
+        return self._get_similar(self.sql_index, self.sql_metadata, question, self.n_results_sql)
+    
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        return [metadata["ddl"] for metadata in self._get_similar(self.ddl_index, self.ddl_metadata, question, self.n_results_ddl)]
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        return [metadata["documentation"] for metadata in self._get_similar(self.doc_index, self.doc_metadata, question, self.n_results_documentation)]
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        sql_data = pd.DataFrame(self.sql_metadata)
+        sql_data['training_data_type'] = 'sql'
+
+        ddl_data = pd.DataFrame(self.ddl_metadata)
+        ddl_data['training_data_type'] = 'ddl'
+
+        doc_data = pd.DataFrame(self.doc_metadata)
+        doc_data['training_data_type'] = 'documentation'
+
+        return pd.concat([sql_data, ddl_data, doc_data], ignore_index=True)
+
+    def remove_training_data(self, id: str, **kwargs) -> bool:
+        for metadata_list, index, index_name in [
+            (self.sql_metadata, self.sql_index, 'sql_index.faiss'),
+            (self.ddl_metadata, self.ddl_index, 'ddl_index.faiss'),
+            (self.doc_metadata, self.doc_index, 'doc_index.faiss')
+        ]:
+            for i, item in enumerate(metadata_list):
+                if item['id'] == id:
+                    del metadata_list[i]
+                    new_index = faiss.IndexFlatL2(self.embedding_dim)
+                    embeddings = [self.generate_embedding(json.dumps(m)) for m in metadata_list]
+                    if embeddings:
+                        new_index.add(np.array(embeddings, dtype=np.float32))
+                    setattr(self, index_name.split('.')[0], new_index)
+                    
+                    if self.curr_client == 'persistent':
+                        self._save_index(new_index, index_name)
+                        self._save_metadata(metadata_list, f"{index_name.split('.')[0]}_metadata.json")
+                    
+                    return True
+        return False
+
+    def remove_collection(self, collection_name: str) -> bool:
+        if collection_name in ["sql", "ddl", "documentation"]:
+            setattr(self, f"{collection_name}_index", faiss.IndexFlatL2(self.embedding_dim))
+            setattr(self, f"{collection_name}_metadata", [])
+            
+            if self.curr_client == 'persistent':
+                self._save_index(getattr(self, f"{collection_name}_index"), f"{collection_name}_index.faiss")
+                self._save_metadata([], f"{collection_name}_metadata.json")
+            
+            return True
+        return False

vanna/qdrant/qdrant.py

@@ -3,7 +3,6 @@ from typing import List, Tuple
 
 import pandas as pd
 from qdrant_client import QdrantClient, grpc, models
-from qdrant_client.http.models.models import UpdateStatus
 
 from ..base import VannaBase
 from ..utils import deterministic_uuid
@@ -234,32 +233,32 @@ class Qdrant_VectorStore(VannaBase):
         return len(self.generate_embedding("ABCDEF"))
 
     def get_similar_question_sql(self, question: str, **kwargs) -> list:
-        results = self._client.search(
+        results = self._client.query_points(
             self.sql_collection_name,
-            query_vector=self.generate_embedding(question),
+            query=self.generate_embedding(question),
             limit=self.n_results,
             with_payload=True,
-        )
+        ).points
 
         return [dict(result.payload) for result in results]
 
     def get_related_ddl(self, question: str, **kwargs) -> list:
-        results = self._client.search(
+        results = self._client.query_points(
             self.ddl_collection_name,
-            query_vector=self.generate_embedding(question),
+            query=self.generate_embedding(question),
             limit=self.n_results,
             with_payload=True,
-        )
+        ).points
 
         return [result.payload["ddl"] for result in results]
 
     def get_related_documentation(self, question: str, **kwargs) -> list:
-        results = self._client.search(
+        results = self._client.query_points(
             self.documentation_collection_name,
-            query_vector=self.generate_embedding(question),
+            query=self.generate_embedding(question),
             limit=self.n_results,
             with_payload=True,
-        )
+        ).points
 
         return [result.payload["documentation"] for result in results]
 

{vanna-0.7.1.dist-info → vanna-0.7.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.7.1
+Version: 0.7.3
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9

{vanna-0.7.1.dist-info → vanna-0.7.3.dist-info}/RECORD

@@ -9,7 +9,7 @@ vanna/advanced/__init__.py,sha256=oDj9g1JbrbCfp4WWdlr_bhgdMqNleyHgr6VXX6DcEbo,65
 vanna/anthropic/__init__.py,sha256=85s_2mAyyPxc0T_0JEvYeAkEKWJwkwqoyUwSC5dw9Gk,43
 vanna/anthropic/anthropic_chat.py,sha256=7X3x8SYwDY28aGyBnt0YNRMG8YY1p_t-foMfKGj8_Oo,2627
 vanna/azuresearch/__init__.py,sha256=tZfvsrCJESiL3EnxA4PrOc5NoO8MXEzCfHX_hnj8n-c,58
-vanna/azuresearch/azuresearch_vector.py,sha256=fc7w_bE6IgOsWb1vkA8t8k-AtuzLyH6zlW4ej4d2lE8,9558
+vanna/azuresearch/azuresearch_vector.py,sha256=_-t53PUnJM914GYbTYlyee06ocfu7l2NkZerBQtlJcs,9566
 vanna/base/__init__.py,sha256=Sl-HM1RRYzAZoSqmL1CZQmF3ZF-byYTCFQP3JZ2A5MU,28
 vanna/base/base.py,sha256=j5xQmK-MeFKAuPjgYLSl1ThCHZieG-ab-RFFSkDlbiw,73679
 vanna/bedrock/__init__.py,sha256=hRT2bgJbHEqViLdL-t9hfjSfFdIOkPU2ADBt-B1En-8,46
@@ -17,6 +17,8 @@ vanna/bedrock/bedrock_converse.py,sha256=Nx5kYm-diAfYmsWAnTP5xnv7V84Og69-AP9b3se
 vanna/chromadb/__init__.py,sha256=-iL0nW_g4uM8nWKMuWnNePfN4nb9uk8P3WzGvezOqRg,50
 vanna/chromadb/chromadb_vector.py,sha256=eKyPck99Y6Jt-BNWojvxLG-zvAERzLSm-3zY-bKXvaA,8792
 vanna/exceptions/__init__.py,sha256=dJ65xxxZh1lqBeg6nz6Tq_r34jLVmjvBvPO9Q6hFaQ8,685
+vanna/faiss/__init__.py,sha256=MXuojmLPt4kUtkES9XKWJcCDHVa4L5a6YF5gebhmKLw,24
+vanna/faiss/faiss.py,sha256=HLUO5PQdnJio9OXJiJcgmRuxVWXvg_XRBnnohS21Z0w,8304
 vanna/flask/__init__.py,sha256=jcdaau1tQ142nL1ZsDklk0ilMkEyRxgQZdmsl1IN4LQ,43866
 vanna/flask/assets.py,sha256=af-vact_5HSftltugBpPxzLkAI14Z0lVWcObyVe6eKE,453462
 vanna/flask/auth.py,sha256=UpKxh7W5cd43W0LGch0VqhncKwB78L6dtOQkl1JY5T0,1246
@@ -45,7 +47,7 @@ vanna/opensearch/opensearch_vector.py,sha256=VhIcrSyNzWR9ZrqrJnyGFOyuQZs3swfbhr8
 vanna/pinecone/__init__.py,sha256=eO5l8aX8vKL6aIUMgAXGPt1jdqKxB_Hic6cmoVAUrD0,90
 vanna/pinecone/pinecone_vector.py,sha256=mpq1lzo3KRj2QfJEw8pwFclFQK1Oi_Nx-lDkx9Gp0mw,11448
 vanna/qdrant/__init__.py,sha256=PX_OsDOiPMvwCJ2iGER1drSdQ9AyM8iN5PEBhRb6qqY,73
-vanna/qdrant/qdrant.py,sha256=qkTWhGrVSAngJZkrcRQ8YFVHcI9j_ZoOGbF6ZVUUdsU,12567
+vanna/qdrant/qdrant.py,sha256=Acl_jN-ZrtoQav_G3FuKypXiuYSo_hlP5lyOOwTxCWM,12527
 vanna/qianfan/Qianfan_Chat.py,sha256=Z-s9MwH22T4KMR8AViAjms6qoj67pHeQkMsbK-aXf1M,5273
 vanna/qianfan/Qianfan_embeddings.py,sha256=TYynAJXlyuZfmoj49h8nU6bXu_GjlXREp3tgfQUca04,954
 vanna/qianfan/__init__.py,sha256=QpR43BjZQZcrcDRkyYcYiS-kyqtYmu23AHDzK0Wy1D0,90
@@ -59,6 +61,6 @@ vanna/vllm/__init__.py,sha256=aNlUkF9tbURdeXAJ8ytuaaF1gYwcG3ny1MfNl_cwQYg,23
 vanna/vllm/vllm.py,sha256=oM_aA-1Chyl7T_Qc_yRKlL6oSX1etsijY9zQdjeMGMQ,2827
 vanna/weaviate/__init__.py,sha256=HL6PAl7ePBAkeG8uln-BmM7IUtWohyTPvDfcPzSGSCg,46
 vanna/weaviate/weaviate_vector.py,sha256=tUJIZjEy2mda8CB6C8zeN2SKkEO-UJdLsIqy69skuF0,7584
-vanna-0.7.1.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-vanna-0.7.1.dist-info/METADATA,sha256=3u_bTszt64VtEBVQbfmUCBvUV7OZk2df07eAhN8mZa4,12407
-vanna-0.7.1.dist-info/RECORD,,
+vanna-0.7.3.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+vanna-0.7.3.dist-info/METADATA,sha256=BOfBtwy1ENcdHApatLWXjqvKj8Zl3bti1hlueVoplR8,12407
+vanna-0.7.3.dist-info/RECORD,,