vanna 0.3.4__py3-none-any.whl → 0.4.1__py3-none-any.whl

vanna/flask/auth.py

@@ -0,0 +1,55 @@
+from abc import ABC, abstractmethod
+
+import flask
+
+
+class AuthInterface(ABC):
+    @abstractmethod
+    def get_user(self, flask_request) -> any:
+        pass
+
+    @abstractmethod
+    def is_logged_in(self, user: any) -> bool:
+        pass
+
+    @abstractmethod
+    def override_config_for_user(self, user: any, config: dict) -> dict:
+        pass
+
+    @abstractmethod
+    def login_form(self) -> str:
+        pass
+
+    @abstractmethod
+    def login_handler(self, flask_request) -> str:
+        pass
+
+    @abstractmethod
+    def callback_handler(self, flask_request) -> str:
+        pass
+
+    @abstractmethod
+    def logout_handler(self, flask_request) -> str:
+        pass
+
+class NoAuth(AuthInterface):
+    def get_user(self, flask_request) -> any:
+        return {}
+
+    def is_logged_in(self, user: any) -> bool:
+        return True
+
+    def override_config_for_user(self, user: any, config: dict) -> dict:
+        return config
+
+    def login_form(self) -> str:
+        return ''
+
+    def login_handler(self, flask_request) -> str:
+        return 'No login required'
+
+    def callback_handler(self, flask_request) -> str:
+        return 'No login required'
+
+    def logout_handler(self, flask_request) -> str:
+        return 'No login required'

vanna/google/__init__.py

@@ -0,0 +1 @@
+from .gemini_chat import GoogleGeminiChat

vanna/google/gemini_chat.py

@@ -0,0 +1,52 @@
+import os
+from ..base import VannaBase
+
+
+class GoogleGeminiChat(VannaBase):
+    def __init__(self, config=None):
+        VannaBase.__init__(self, config=config)
+
+        # default temperature - can be overrided using config
+        self.temperature = 0.7      
+
+        if "temperature" in config:
+            self.temperature = config["temperature"]
+
+        if "model_name" in config:
+            model_name = config["model_name"]
+        else:
+            model_name = "gemini-1.0-pro"
+
+        self.google_api_key = None
+
+        if "api_key" in config or os.getenv("GOOGLE_API_KEY"):
+            """
+            If Google api_key is provided through config
+            or set as an environment variable, assign it.
+            """
+            import google.generativeai as genai
+
+            genai.configure(api_key=config["api_key"])
+            self.chat_model = genai.GenerativeModel(model_name)
+        else:
+            # Authenticate using VertexAI
+            from vertexai.preview.generative_models import GenerativeModel
+            self.chat_model = GenerativeModel("gemini-pro")
+
+    def system_message(self, message: str) -> any:
+        return message
+
+    def user_message(self, message: str) -> any:
+        return message
+
+    def assistant_message(self, message: str) -> any:
+        return message
+
+    def submit_prompt(self, prompt, **kwargs) -> str:
+        response = self.chat_model.generate_content(
+            prompt,
+            generation_config={
+                "temperature": self.temperature,
+            },
+        )
+        return response.text

vanna/qdrant/__init__.py

@@ -0,0 +1,3 @@
+from .qdrant import Qdrant_VectorStore
+
+__all__ = ["Qdrant_VectorStore"]

vanna/qdrant/qdrant.py

@@ -0,0 +1,324 @@
+from functools import cached_property
+from typing import List, Tuple
+
+import pandas as pd
+from qdrant_client import QdrantClient, grpc, models
+
+from ..base import VannaBase
+from ..utils import deterministic_uuid
+
+DOCUMENTATION_COLLECTION_NAME = "documentation"
+DDL_COLLECTION_NAME = "ddl"
+SQL_COLLECTION_NAME = "sql"
+SCROLL_SIZE = 1000
+
+ID_SUFFIXES = {
+    DDL_COLLECTION_NAME: "ddl",
+    DOCUMENTATION_COLLECTION_NAME: "doc",
+    SQL_COLLECTION_NAME: "sql",
+}
+
+
+class Qdrant_VectorStore(VannaBase):
+    """Vectorstore implementation using Qdrant - https://qdrant.tech/"""
+
+    def __init__(
+        self,
+        config={},
+    ):
+        """
+        Vectorstore implementation using Qdrant - https://qdrant.tech/
+
+        Args:
+            - config (dict, optional): Dictionary of `Qdrant_VectorStore config` options. Defaults to `{}`.
+                - client: A `qdrant_client.QdrantClient` instance. Overrides other config options.
+                - location: If `":memory:"` - use in-memory Qdrant instance. If `str` - use it as a `url` parameter.
+                - url: Either host or str of "Optional[scheme], host, Optional[port], Optional[prefix]". Eg. `"http://localhost:6333"`.
+                - prefer_grpc: If `true` - use gPRC interface whenever possible in custom methods.
+                - https: If `true` - use HTTPS(SSL) protocol. Default: `None`
+                - api_key: API key for authentication in Qdrant Cloud. Default: `None`
+                - timeout: Timeout for REST and gRPC API requests. Defaults to 5 seconds for REST and unlimited for gRPC.
+                - path: Persistence path for QdrantLocal. Default: `None`.
+                - prefix: Prefix to the REST URL paths. Example: `service/v1` will result in `http://localhost:6333/service/v1/{qdrant-endpoint}`.
+                - n_results: Number of results to return from similarity search. Defaults to 10.
+                - fastembed_model: [Model](https://qdrant.github.io/fastembed/examples/Supported_Models/#supported-text-embedding-models) to use for `fastembed.TextEmbedding`.
+                  Defaults to `"BAAI/bge-small-en-v1.5"`.
+                - collection_params: Additional parameters to pass to `qdrant_client.QdrantClient#create_collection()` method.
+                - distance_metric: Distance metric to use when creating collections. Defaults to `qdrant_client.models.Distance.COSINE`.
+
+        Raises:
+            TypeError: If config["client"] is not a `qdrant_client.QdrantClient` instance
+        """
+        VannaBase.__init__(self, config=config)
+        client = config.get("client")
+
+        if client is None:
+            self._client = QdrantClient(
+                location=config.get("location", None),
+                url=config.get("url", None),
+                prefer_grpc=config.get("prefer_grpc", False),
+                https=config.get("https", None),
+                api_key=config.get("api_key", None),
+                timeout=config.get("timeout", None),
+                path=config.get("path", None),
+                prefix=config.get("prefix", None),
+            )
+        elif not isinstance(client, QdrantClient):
+            raise TypeError(
+                f"Unsupported client of type {client.__class__} was set in config"
+            )
+
+        else:
+            self._client = client
+
+        self.n_results = config.get("n_results", 10)
+        self.fastembed_model = config.get("fastembed_model", "BAAI/bge-small-en-v1.5")
+        self.collection_params = config.get("collection_params", {})
+        self.distance_metric = config.get("distance_metric", models.Distance.COSINE)
+
+        self._setup_collections()
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        question_answer = format("Question: {0}\n\nSQL: {1}", question, sql)
+        id = deterministic_uuid(question_answer)
+
+        self._client.upsert(
+            SQL_COLLECTION_NAME,
+            points=[
+                models.PointStruct(
+                    id=id,
+                    vector=self.generate_embedding(question_answer),
+                    payload={
+                        "question": question,
+                        "sql": sql,
+                    },
+                )
+            ],
+        )
+
+        return self._format_point_id(id, SQL_COLLECTION_NAME)
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        id = deterministic_uuid(ddl)
+        self._client.upsert(
+            DDL_COLLECTION_NAME,
+            points=[
+                models.PointStruct(
+                    id=id,
+                    vector=self.generate_embedding(ddl),
+                    payload={
+                        "ddl": ddl,
+                    },
+                )
+            ],
+        )
+        return self._format_point_id(id, DDL_COLLECTION_NAME)
+
+    def add_documentation(self, documentation: str, **kwargs) -> str:
+        id = deterministic_uuid(documentation)
+
+        self._client.upsert(
+            DOCUMENTATION_COLLECTION_NAME,
+            points=[
+                models.PointStruct(
+                    id=id,
+                    vector=self.generate_embedding(documentation),
+                    payload={
+                        "documentation": documentation,
+                    },
+                )
+            ],
+        )
+
+        return self._format_point_id(id, DOCUMENTATION_COLLECTION_NAME)
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        df = pd.DataFrame()
+
+        if sql_data := self._get_all_points(SQL_COLLECTION_NAME):
+            question_list = [data.payload["question"] for data in sql_data]
+            sql_list = [data.payload["sql"] for data in sql_data]
+            id_list = [
+                self._format_point_id(data.id, SQL_COLLECTION_NAME) for data in sql_data
+            ]
+
+            df_sql = pd.DataFrame(
+                {
+                    "id": id_list,
+                    "question": question_list,
+                    "content": sql_list,
+                }
+            )
+
+            df_sql["training_data_type"] = "sql"
+
+            df = pd.concat([df, df_sql])
+
+        if ddl_data := self._get_all_points(DDL_COLLECTION_NAME):
+            ddl_list = [data.payload["ddl"] for data in ddl_data]
+            id_list = [
+                self._format_point_id(data.id, DDL_COLLECTION_NAME) for data in sql_data
+            ]
+
+            df_ddl = pd.DataFrame(
+                {
+                    "id": id_list,
+                    "question": [None for _ in ddl_list],
+                    "content": ddl_list,
+                }
+            )
+
+            df_ddl["training_data_type"] = "ddl"
+
+            df = pd.concat([df, df_ddl])
+
+        doc_data = self.documentation_collection.get()
+
+        if doc_data := self._get_all_points(DOCUMENTATION_COLLECTION_NAME):
+            document_list = [data.payload["documentation"] for data in doc_data]
+            id_list = [
+                self._format_point_id(data.id, DOCUMENTATION_COLLECTION_NAME)
+                for data in doc_data
+            ]
+
+            df_doc = pd.DataFrame(
+                {
+                    "id": id_list,
+                    "question": [None for _ in document_list],
+                    "content": document_list,
+                }
+            )
+
+            df_doc["training_data_type"] = "documentation"
+
+            df = pd.concat([df, df_doc])
+
+        return df
+
+    def remove_training_data(self, id: str, **kwargs) -> bool:
+        try:
+            id, collection_name = self._parse_point_id(id)
+            self._client.delete(collection_name, points_selector=[id])
+        except ValueError:
+            return False
+
+    def remove_collection(self, collection_name: str) -> bool:
+        """
+        This function can reset the collection to empty state.
+
+        Args:
+            collection_name (str): sql or ddl or documentation
+
+        Returns:
+            bool: True if collection is deleted, False otherwise
+        """
+        if collection_name in ID_SUFFIXES.keys():
+            self._client.delete_collection(collection_name)
+            self._setup_collections()
+            return True
+        else:
+            return False
+
+    @cached_property
+    def embeddings_dimension(self):
+        return len(self.generate_embedding("ABCDEF"))
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> list:
+        results = self._client.search(
+            SQL_COLLECTION_NAME,
+            query_vector=self.generate_embedding(question),
+            limit=self.n_results,
+            with_payload=True,
+        )
+
+        return [dict(result.payload) for result in results]
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        results = self._client.search(
+            DDL_COLLECTION_NAME,
+            query_vector=self.generate_embedding(question),
+            limit=self.n_results,
+            with_payload=True,
+        )
+
+        return [result.payload["ddl"] for result in results]
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        results = self._client.search(
+            DOCUMENTATION_COLLECTION_NAME,
+            query_vector=self.generate_embedding(question),
+            limit=self.n_results,
+            with_payload=True,
+        )
+
+        return [result.payload["documentation"] for result in results]
+
+    def generate_embedding(self, data: str, **kwargs) -> List[float]:
+        embedding_model = self._client._get_or_init_model(
+            model_name=self.fastembed_model
+        )
+        embedding = next(embedding_model.embed(data))
+
+        return embedding.tolist()
+
+    def _get_all_points(self, collection_name: str):
+        results: List[models.Record] = []
+        next_offset = None
+        stop_scrolling = False
+        while not stop_scrolling:
+            records, next_offset = self._client.scroll(
+                collection_name,
+                limit=SCROLL_SIZE,
+                offset=next_offset,
+                with_payload=True,
+                with_vectors=False,
+            )
+            stop_scrolling = next_offset is None or (
+                isinstance(next_offset, grpc.PointId)
+                and next_offset.num == 0
+                and next_offset.uuid == ""
+            )
+
+            results.extend(records)
+
+        return results
+
+    def _setup_collections(self):
+        if not self._client.collection_exists(SQL_COLLECTION_NAME):
+            self._client.create_collection(
+                collection_name=SQL_COLLECTION_NAME,
+                vectors_config=models.VectorParams(
+                    size=self.embeddings_dimension,
+                    distance=self.distance_metric,
+                ),
+                **self.collection_params,
+            )
+
+        if not self._client.collection_exists(DDL_COLLECTION_NAME):
+            self._client.create_collection(
+                collection_name=DDL_COLLECTION_NAME,
+                vectors_config=models.VectorParams(
+                    size=self.embeddings_dimension,
+                    distance=self.distance_metric,
+                ),
+                **self.collection_params,
+            )
+        if not self._client.collection_exists(DOCUMENTATION_COLLECTION_NAME):
+            self._client.create_collection(
+                collection_name=DOCUMENTATION_COLLECTION_NAME,
+                vectors_config=models.VectorParams(
+                    size=self.embeddings_dimension,
+                    distance=self.distance_metric,
+                ),
+                **self.collection_params,
+            )
+
+    def _format_point_id(self, id: str, collection_name: str) -> str:
+        return "{0}-{1}".format(id, ID_SUFFIXES[collection_name])
+
+    def _parse_point_id(self, id: str) -> Tuple[str, str]:
+        id, suffix = id.rsplit("-", 1)
+        for collection_name, suffix in ID_SUFFIXES.items():
+            if type == suffix:
+                return id, collection_name
+        raise ValueError(f"Invalid id {id}")

{vanna-0.3.4.dist-info → vanna-0.4.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.3.4
+Version: 0.4.1
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -28,17 +28,24 @@ Requires-Dist: chromadb ; extra == "all"
 Requires-Dist: anthropic ; extra == "all"
 Requires-Dist: zhipuai ; extra == "all"
 Requires-Dist: marqo ; extra == "all"
+Requires-Dist: google-generativeai ; extra == "all"
+Requires-Dist: google-cloud-aiplatform ; extra == "all"
+Requires-Dist: qdrant-client ; extra == "all"
+Requires-Dist: fastembed ; extra == "all"
 Requires-Dist: anthropic ; extra == "anthropic"
 Requires-Dist: google-cloud-bigquery ; extra == "bigquery"
 Requires-Dist: chromadb ; extra == "chromadb"
 Requires-Dist: duckdb ; extra == "duckdb"
 Requires-Dist: google-generativeai ; extra == "gemini"
+Requires-Dist: google-generativeai ; extra == "google"
+Requires-Dist: google-cloud-aiplatform ; extra == "google"
 Requires-Dist: marqo ; extra == "marqo"
 Requires-Dist: mistralai ; extra == "mistralai"
 Requires-Dist: PyMySQL ; extra == "mysql"
 Requires-Dist: openai ; extra == "openai"
 Requires-Dist: psycopg2-binary ; extra == "postgres"
 Requires-Dist: db-dtypes ; extra == "postgres"
+Requires-Dist: qdrant-client ; extra == "qdrant"
 Requires-Dist: snowflake-connector-python ; extra == "snowflake"
 Requires-Dist: tox ; extra == "test"
 Requires-Dist: zhipuai ; extra == "zhipuai"
@@ -50,11 +57,13 @@ Provides-Extra: bigquery
 Provides-Extra: chromadb
 Provides-Extra: duckdb
 Provides-Extra: gemini
+Provides-Extra: google
 Provides-Extra: marqo
 Provides-Extra: mistralai
 Provides-Extra: mysql
 Provides-Extra: openai
 Provides-Extra: postgres
+Provides-Extra: qdrant
 Provides-Extra: snowflake
 Provides-Extra: test
 Provides-Extra: zhipuai

{vanna-0.3.4.dist-info → vanna-0.4.1.dist-info}/RECORD

@@ -8,12 +8,15 @@ vanna/ZhipuAI/__init__.py,sha256=NlsijtcZp5Tj9jkOe9fNcOQND_QsGgu7otODsCLBPr0,116
 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=89XPWy97YVx6090mNmu1zvn4k8X1pusCuAIypHHexNc,58100
+vanna/base/base.py,sha256=_2vANGAcUe6IrsEhZyFnE6FdO8NMLylfyGlRI4XujWE,58143
 vanna/chromadb/__init__.py,sha256=-iL0nW_g4uM8nWKMuWnNePfN4nb9uk8P3WzGvezOqRg,50
-vanna/chromadb/chromadb_vector.py,sha256=1n4U4XpXThCFqyJf0zAYVA7mQu9rUkjOFtYn9e04JAo,8461
+vanna/chromadb/chromadb_vector.py,sha256=eKyPck99Y6Jt-BNWojvxLG-zvAERzLSm-3zY-bKXvaA,8792
 vanna/exceptions/__init__.py,sha256=N76unE7sjbGGBz6LmCrPQAugFWr9cUFv8ErJxBrCTts,717
-vanna/flask/__init__.py,sha256=tpwpA8596Uyn60FAy7I5oJ81c7kgCB2JG9X044P0_SA,21211
-vanna/flask/assets.py,sha256=pOOtPV8aWtFsTuxJneFHcfrXhXh6cOSvS-Y8JO2HYrY,180924
+vanna/flask/__init__.py,sha256=5Du2oK5s-VSLicRPvxMAL1-Gh_jdX847FjJOf5AVapo,23721
+vanna/flask/assets.py,sha256=ZESgn0-XrJl4_YV69Lu7Dr7-Y9Eql7xqb7PsoMzrofw,183889
+vanna/flask/auth.py,sha256=UpKxh7W5cd43W0LGch0VqhncKwB78L6dtOQkl1JY5T0,1246
+vanna/google/__init__.py,sha256=M-dCxCZcKL4bTQyMLj6r6VRs65YNX9Tl2aoPCuqGm-8,41
+vanna/google/gemini_chat.py,sha256=ps3A-afFbCo3HeFTLL_nMoQO1PsGvRUUPRUppbMcDew,1584
 vanna/marqo/__init__.py,sha256=GaAWtJ0B-H5rTY607iLCCrLD7T0zMYM5qWIomEB9gLk,37
 vanna/marqo/marqo.py,sha256=W7WTtzWp4RJjZVy6OaXHqncUBIPdI4Q7qH7BRCxZ1_A,5242
 vanna/mistral/__init__.py,sha256=70rTY-69Z2ehkkMj84dNMCukPo6AWdflBGvIB_pztS0,29
@@ -23,9 +26,11 @@ vanna/ollama/ollama.py,sha256=jfW9VQHAcmzDeo4jF3HJjOMYwAWmptknKqEJaQ0MTno,2418
 vanna/openai/__init__.py,sha256=tGkeQ7wTIPsando7QhoSHehtoQVdYLwFbKNlSmCmNeQ,86
 vanna/openai/openai_chat.py,sha256=lm-hUsQxu6Q1t06A2csC037zI4VkMk0wFbQ-_Lj74Wg,4764
 vanna/openai/openai_embeddings.py,sha256=g4pNh9LVcYP9wOoO8ecaccDFWmCUYMInebfHucAa2Gc,1260
+vanna/qdrant/__init__.py,sha256=PX_OsDOiPMvwCJ2iGER1drSdQ9AyM8iN5PEBhRb6qqY,73
+vanna/qdrant/qdrant.py,sha256=XlesB0UniR0LmiXvQ8Ct7o19EdCZPuvUgvyzYSlp94c,11940
 vanna/types/__init__.py,sha256=Qhn_YscKtJh7mFPCyCDLa2K8a4ORLMGVnPpTbv9uB2U,4957
 vanna/vannadb/__init__.py,sha256=C6UkYocmO6dmzfPKZaWojN0mI5YlZZ9VIbdcquBE58A,48
 vanna/vannadb/vannadb_vector.py,sha256=9YwTO3Lh5owWQE7KPMBqLp2EkiGV0RC1sEYhslzJzgI,6168
-vanna-0.3.4.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-vanna-0.3.4.dist-info/METADATA,sha256=FEg4vs5ZiSAvd5YkF5oEfFqod9n3UoNfi51Q_2WKotA,10107
-vanna-0.3.4.dist-info/RECORD,,
+vanna-0.4.1.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+vanna-0.4.1.dist-info/METADATA,sha256=SwU235jbylf7gJ4VwDXtz7bLhUF5uTV0IdPCnK1HH68,10512
+vanna-0.4.1.dist-info/RECORD,,