vanna 0.6.0__py3-none-any.whl → 0.6.1__py3-none-any.whl

vanna/base/base.py

@@ -992,6 +992,7 @@ class VannaBase(ABC):
         def run_sql_mysql(sql: str) -> Union[pd.DataFrame, None]:
             if conn:
                 try:
+                    conn.ping(reconnect=True)
                     cs = conn.cursor()
                     cs.execute(sql)
                     results = cs.fetchall()

vanna/milvus/__init__.py

@@ -0,0 +1 @@
+from .milvus_vector import Milvus_VectorStore

vanna/milvus/milvus_vector.py

@@ -0,0 +1,305 @@
+import uuid
+from typing import List
+
+import pandas as pd
+from pymilvus import DataType, MilvusClient, model
+
+from ..base import VannaBase
+
+# Setting the URI as a local file, e.g.`./milvus.db`,
+# is the most convenient method, as it automatically utilizes Milvus Lite
+# to store all data in this file.
+#
+# If you have large scale of data such as more than a million docs, we
+# recommend setting up a more performant Milvus server on docker or kubernetes.
+# When using this setup, please use the server URI,
+# e.g.`http://localhost:19530`, as your URI.
+
+DEFAULT_MILVUS_URI = "./milvus.db"
+# DEFAULT_MILVUS_URI = "http://localhost:19530"
+
+MAX_LIMIT_SIZE = 10_000
+
+
+class Milvus_VectorStore(VannaBase):
+    """
+    Vectorstore implementation using Milvus - https://milvus.io/docs/quickstart.md
+
+    Args:
+        - config (dict, optional): Dictionary of `Milvus_VectorStore config` options. Defaults to `None`.
+            - milvus_client: A `pymilvus.MilvusClient` instance.
+            - embedding_function:
+                A `milvus_model.base.BaseEmbeddingFunction` instance. Defaults to `DefaultEmbeddingFunction()`.
+                For more models, please refer to:
+                https://milvus.io/docs/embeddings.md
+    """
+    def __init__(self, config=None):
+        VannaBase.__init__(self, config=config)
+
+        if "milvus_client" in config:
+            self.milvus_client = config["milvus_client"]
+        else:
+            self.milvus_client = MilvusClient(uri=DEFAULT_MILVUS_URI)
+
+        if "embedding_function" in config:
+            self.embedding_function = config.get("embedding_function")
+        else:
+            self.embedding_function = model.DefaultEmbeddingFunction()
+        self._embedding_dim = self.embedding_function.encode_documents(["foo"])[0].shape[0]
+        self._create_collections()
+        self.n_results = config.get("n_results", 10)
+
+    def _create_collections(self):
+        self._create_sql_collection("vannasql")
+        self._create_ddl_collection("vannaddl")
+        self._create_doc_collection("vannadoc")
+
+
+    def generate_embedding(self, data: str, **kwargs) -> List[float]:
+        return self.embedding_function.encode_documents(data).tolist()
+
+
+    def _create_sql_collection(self, name: str):
+        if not self.milvus_client.has_collection(collection_name=name):
+            vannasql_schema = MilvusClient.create_schema(
+                auto_id=False,
+                enable_dynamic_field=False,
+            )
+            vannasql_schema.add_field(field_name="id", datatype=DataType.VARCHAR, max_length=65535, is_primary=True)
+            vannasql_schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=65535)
+            vannasql_schema.add_field(field_name="sql", datatype=DataType.VARCHAR, max_length=65535)
+            vannasql_schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=self._embedding_dim)
+
+            vannasql_index_params = self.milvus_client.prepare_index_params()
+            vannasql_index_params.add_index(
+                field_name="vector",
+                index_name="vector",
+                index_type="AUTOINDEX",
+                metric_type="L2",
+            )
+            self.milvus_client.create_collection(
+                collection_name=name,
+                schema=vannasql_schema,
+                index_params=vannasql_index_params,
+                consistency_level="Strong"
+            )
+
+    def _create_ddl_collection(self, name: str):
+        if not self.milvus_client.has_collection(collection_name=name):
+            vannaddl_schema = MilvusClient.create_schema(
+                auto_id=False,
+                enable_dynamic_field=False,
+            )
+            vannaddl_schema.add_field(field_name="id", datatype=DataType.VARCHAR, max_length=65535, is_primary=True)
+            vannaddl_schema.add_field(field_name="ddl", datatype=DataType.VARCHAR, max_length=65535)
+            vannaddl_schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=self._embedding_dim)
+
+            vannaddl_index_params = self.milvus_client.prepare_index_params()
+            vannaddl_index_params.add_index(
+                field_name="vector",
+                index_name="vector",
+                index_type="AUTOINDEX",
+                metric_type="L2",
+            )
+            self.milvus_client.create_collection(
+                collection_name=name,
+                schema=vannaddl_schema,
+                index_params=vannaddl_index_params,
+                consistency_level="Strong"
+            )
+
+    def _create_doc_collection(self, name: str):
+        if not self.milvus_client.has_collection(collection_name=name):
+            vannadoc_schema = MilvusClient.create_schema(
+                auto_id=False,
+                enable_dynamic_field=False,
+            )
+            vannadoc_schema.add_field(field_name="id", datatype=DataType.VARCHAR, max_length=65535, is_primary=True)
+            vannadoc_schema.add_field(field_name="doc", datatype=DataType.VARCHAR, max_length=65535)
+            vannadoc_schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=self._embedding_dim)
+
+            vannadoc_index_params = self.milvus_client.prepare_index_params()
+            vannadoc_index_params.add_index(
+                field_name="vector",
+                index_name="vector",
+                index_type="AUTOINDEX",
+                metric_type="L2",
+            )
+            self.milvus_client.create_collection(
+                collection_name=name,
+                schema=vannadoc_schema,
+                index_params=vannadoc_index_params,
+                consistency_level="Strong"
+            )
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        if len(question) == 0 or len(sql) == 0:
+            raise Exception("pair of question and sql can not be null")
+        _id = str(uuid.uuid4()) + "-sql"
+        embedding = self.embedding_function.encode_documents([question])[0]
+        self.milvus_client.insert(
+            collection_name="vannasql",
+            data={
+                "id": _id,
+                "text": question,
+                "sql": sql,
+                "vector": embedding
+            }
+        )
+        return _id
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        if len(ddl) == 0:
+            raise Exception("ddl can not be null")
+        _id = str(uuid.uuid4()) + "-ddl"
+        embedding = self.embedding_function.encode_documents([ddl])[0]
+        self.milvus_client.insert(
+            collection_name="vannaddl",
+            data={
+                "id": _id,
+                "ddl": ddl,
+                "vector": embedding
+            }
+        )
+        return _id
+
+    def add_documentation(self, documentation: str, **kwargs) -> str:
+        if len(documentation) == 0:
+            raise Exception("documentation can not be null")
+        _id = str(uuid.uuid4()) + "-doc"
+        embedding = self.embedding_function.encode_documents([documentation])[0]
+        self.milvus_client.insert(
+            collection_name="vannadoc",
+            data={
+                "id": _id,
+                "doc": documentation,
+                "vector": embedding
+            }
+        )
+        return _id
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        sql_data = self.milvus_client.query(
+            collection_name="vannasql",
+            output_fields=["*"],
+            limit=MAX_LIMIT_SIZE,
+        )
+        df = pd.DataFrame()
+        df_sql = pd.DataFrame(
+            {
+                "id": [doc["id"] for doc in sql_data],
+                "question": [doc["text"] for doc in sql_data],
+                "content": [doc["sql"] for doc in sql_data],
+            }
+        )
+        df = pd.concat([df, df_sql])
+
+        ddl_data = self.milvus_client.query(
+            collection_name="vannaddl",
+            output_fields=["*"],
+            limit=MAX_LIMIT_SIZE,
+        )
+
+        df_ddl = pd.DataFrame(
+            {
+                "id": [doc["id"] for doc in ddl_data],
+                "question": [None for doc in ddl_data],
+                "content": [doc["ddl"] for doc in ddl_data],
+            }
+        )
+        df = pd.concat([df, df_ddl])
+
+        doc_data = self.milvus_client.query(
+            collection_name="vannadoc",
+            output_fields=["*"],
+            limit=MAX_LIMIT_SIZE,
+        )
+
+        df_doc = pd.DataFrame(
+            {
+                "id": [doc["id"] for doc in doc_data],
+                "question": [None for doc in doc_data],
+                "content": [doc["doc"] for doc in doc_data],
+            }
+        )
+        df = pd.concat([df, df_doc])
+        return df
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> list:
+        search_params = {
+            "metric_type": "L2",
+            "params": {"nprobe": 128},
+        }
+        embeddings = self.embedding_function.encode_queries([question])
+        res = self.milvus_client.search(
+            collection_name="vannasql",
+            anns_field="vector",
+            data=embeddings,
+            limit=self.n_results,
+            output_fields=["text", "sql"],
+            search_params=search_params
+        )
+        res = res[0]
+
+        list_sql = []
+        for doc in res:
+            dict = {}
+            dict["question"] = doc["entity"]["text"]
+            dict["sql"] = doc["entity"]["sql"]
+            list_sql.append(dict)
+        return list_sql
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        search_params = {
+            "metric_type": "L2",
+            "params": {"nprobe": 128},
+        }
+        embeddings = self.embedding_function.encode_queries([question])
+        res = self.milvus_client.search(
+            collection_name="vannaddl",
+            anns_field="vector",
+            data=embeddings,
+            limit=self.n_results,
+            output_fields=["ddl"],
+            search_params=search_params
+        )
+        res = res[0]
+
+        list_ddl = []
+        for doc in res:
+            list_ddl.append(doc["entity"]["ddl"])
+        return list_ddl
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        search_params = {
+            "metric_type": "L2",
+            "params": {"nprobe": 128},
+        }
+        embeddings = self.embedding_function.encode_queries([question])
+        res = self.milvus_client.search(
+            collection_name="vannadoc",
+            anns_field="vector",
+            data=embeddings,
+            limit=self.n_results,
+            output_fields=["doc"],
+            search_params=search_params
+        )
+        res = res[0]
+
+        list_doc = []
+        for doc in res:
+            list_doc.append(doc["entity"]["doc"])
+        return list_doc
+
+    def remove_training_data(self, id: str, **kwargs) -> bool:
+        if id.endswith("-sql"):
+            self.milvus_client.delete(collection_name="vannasql", ids=[id])
+            return True
+        elif id.endswith("-ddl"):
+            self.milvus_client.delete(collection_name="vannaddl", ids=[id])
+            return True
+        elif id.endswith("-doc"):
+            self.milvus_client.delete(collection_name="vannadoc", ids=[id])
+            return True
+        else:
+            return False

vanna/qdrant/qdrant.py

@@ -39,16 +39,6 @@ class Qdrant_VectorStore(VannaBase):
         TypeError: If config["client"] is not a `qdrant_client.QdrantClient` instance
     """
 
-    documentation_collection_name = "documentation"
-    ddl_collection_name = "ddl"
-    sql_collection_name = "sql"
-
-    id_suffixes = {
-        ddl_collection_name: "ddl",
-        documentation_collection_name: "doc",
-        sql_collection_name: "sql",
-    }
-
     def __init__(
         self,
         config={},
@@ -80,15 +70,21 @@ class Qdrant_VectorStore(VannaBase):
         self.collection_params = config.get("collection_params", {})
         self.distance_metric = config.get("distance_metric", models.Distance.COSINE)
         self.documentation_collection_name = config.get(
-            "documentation_collection_name", self.documentation_collection_name
+            "documentation_collection_name", "documentation"
         )
         self.ddl_collection_name = config.get(
-            "ddl_collection_name", self.ddl_collection_name
+            "ddl_collection_name", "ddl"
         )
         self.sql_collection_name = config.get(
-            "sql_collection_name", self.sql_collection_name
+            "sql_collection_name", "sql"
         )
 
+        self.id_suffixes = {
+            self.ddl_collection_name: "ddl",
+            self.documentation_collection_name: "doc",
+            self.sql_collection_name: "sql",
+        }
+
         self._setup_collections()
 
     def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
@@ -212,7 +208,7 @@ class Qdrant_VectorStore(VannaBase):
         try:
             id, collection_name = self._parse_point_id(id)
             res = self._client.delete(collection_name, points_selector=[id])
-            res == UpdateStatus.COMPLETED
+            return True
         except ValueError:
             return False
 

{vanna-0.6.0.dist-info → vanna-0.6.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.6.0
+Version: 0.6.1
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -39,6 +39,7 @@ Requires-Dist: opensearch-py ; extra == "all"
 Requires-Dist: opensearch-dsl ; extra == "all"
 Requires-Dist: transformers ; extra == "all"
 Requires-Dist: pinecone-client ; extra == "all"
+Requires-Dist: pymilvus[model] ; extra == "all"
 Requires-Dist: anthropic ; extra == "anthropic"
 Requires-Dist: google-cloud-bigquery ; extra == "bigquery"
 Requires-Dist: chromadb ; extra == "chromadb"
@@ -49,6 +50,7 @@ Requires-Dist: google-generativeai ; extra == "google"
 Requires-Dist: google-cloud-aiplatform ; extra == "google"
 Requires-Dist: transformers ; extra == "hf"
 Requires-Dist: marqo ; extra == "marqo"
+Requires-Dist: pymilvus[model] ; extra == "milvus"
 Requires-Dist: mistralai ; extra == "mistralai"
 Requires-Dist: PyMySQL ; extra == "mysql"
 Requires-Dist: ollama ; extra == "ollama"
@@ -78,6 +80,7 @@ Provides-Extra: gemini
 Provides-Extra: google
 Provides-Extra: hf
 Provides-Extra: marqo
+Provides-Extra: milvus
 Provides-Extra: mistralai
 Provides-Extra: mysql
 Provides-Extra: ollama

{vanna-0.6.0.dist-info → vanna-0.6.1.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=Wk0o-NMW1uvR2fhSWxrR_2FqNh-dLprNG4uuVqpqAkY,2615
 vanna/base/__init__.py,sha256=Sl-HM1RRYzAZoSqmL1CZQmF3ZF-byYTCFQP3JZ2A5MU,28
-vanna/base/base.py,sha256=XiBJU4LOKoFnY1vnkr_jz3soAdVdoL0D1WIZdzZCVyU,70942
+vanna/base/base.py,sha256=l1H0TKsK9DN3n5XgDkUckdLois4dTCAUwrVsRa_6SlQ,70988
 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
@@ -22,6 +22,8 @@ vanna/hf/__init__.py,sha256=vD0bIhfLkA1UsvVSF4MAz3Da8aQunkQo3wlDztmMuj0,19
 vanna/hf/hf.py,sha256=v1v6sZnbj5xcrjgmvLP_ytS9NM7E5d0GyMfXXtr6BMU,2703
 vanna/marqo/__init__.py,sha256=GaAWtJ0B-H5rTY607iLCCrLD7T0zMYM5qWIomEB9gLk,37
 vanna/marqo/marqo.py,sha256=W7WTtzWp4RJjZVy6OaXHqncUBIPdI4Q7qH7BRCxZ1_A,5242
+vanna/milvus/__init__.py,sha256=VBasJG2eTKbJI6CEand7kPLNBrqYrn0QCAhSYVz814s,46
+vanna/milvus/milvus_vector.py,sha256=Mq0eaSh0UcTYhgh8mTm0fvS6rbfL6tQONVnDZGemWoM,11268
 vanna/mistral/__init__.py,sha256=70rTY-69Z2ehkkMj84dNMCukPo6AWdflBGvIB_pztS0,29
 vanna/mistral/mistral.py,sha256=DAEqAT9SzC91rfMM_S3SuzBZ34MrKHw9qAj6EP2MGVk,1508
 vanna/mock/__init__.py,sha256=nYR2WfcV5NdwpK3V64QGOWHBGc3ESN9uV68JLS76aRw,97
@@ -38,12 +40,12 @@ 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=RyICUvOO_jt8u9MB4oIYhqv3BicZ0d9pQkSFwkIfUjg,12719
+vanna/qdrant/qdrant.py,sha256=qkTWhGrVSAngJZkrcRQ8YFVHcI9j_ZoOGbF6ZVUUdsU,12567
 vanna/types/__init__.py,sha256=Qhn_YscKtJh7mFPCyCDLa2K8a4ORLMGVnPpTbv9uB2U,4957
 vanna/vannadb/__init__.py,sha256=C6UkYocmO6dmzfPKZaWojN0mI5YlZZ9VIbdcquBE58A,48
 vanna/vannadb/vannadb_vector.py,sha256=N8poMYvAojoaOF5gI4STD5pZWK9lBKPvyIjbh9dPBa0,14189
 vanna/vllm/__init__.py,sha256=aNlUkF9tbURdeXAJ8ytuaaF1gYwcG3ny1MfNl_cwQYg,23
 vanna/vllm/vllm.py,sha256=oM_aA-1Chyl7T_Qc_yRKlL6oSX1etsijY9zQdjeMGMQ,2827
-vanna-0.6.0.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-vanna-0.6.0.dist-info/METADATA,sha256=Nu63PMm3HT3HAkpxKU4fte9CKpIDgnL90Ww2AljWSOs,11506
-vanna-0.6.0.dist-info/RECORD,,
+vanna-0.6.1.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+vanna-0.6.1.dist-info/METADATA,sha256=BApaZrir1-x-Y99Ufh4eoVZ9_poaGuG2uh-YVQfBP3Y,11628
+vanna-0.6.1.dist-info/RECORD,,