feat: pinecone vectorstore

.gitignore

@@ -17,3 +17,4 @@ dist
 htmlcov
 chroma.sqlite3
 *.bin
+.coverage.*

pyproject.toml

@@ -33,7 +33,7 @@ bigquery = ["google-cloud-bigquery"]
 snowflake = ["snowflake-connector-python"]
 duckdb = ["duckdb"]
 google = ["google-generativeai", "google-cloud-aiplatform"]
-all = ["psycopg2-binary", "db-dtypes", "PyMySQL", "google-cloud-bigquery", "snowflake-connector-python", "duckdb", "openai", "mistralai", "chromadb", "anthropic", "zhipuai", "marqo", "google-generativeai", "google-cloud-aiplatform", "qdrant-client", "fastembed", "ollama", "httpx", "opensearch-py", "opensearch-dsl", "transformers"]
+all = ["psycopg2-binary", "db-dtypes", "PyMySQL", "google-cloud-bigquery", "snowflake-connector-python", "duckdb", "openai", "mistralai", "chromadb", "anthropic", "zhipuai", "marqo", "google-generativeai", "google-cloud-aiplatform", "qdrant-client", "fastembed", "ollama", "httpx", "opensearch-py", "opensearch-dsl", "transformers", "pinecone-client"]
 test = ["tox"]
 chromadb = ["chromadb"]
 openai = ["openai"]
@@ -45,5 +45,6 @@ zhipuai = ["zhipuai"]
 ollama = ["ollama", "httpx"]
 qdrant = ["qdrant-client", "fastembed"]
 vllm = ["vllm"]
+pinecone = ["pinecone-client", "fastembed"]
 opensearch = ["opensearch-py", "opensearch-dsl"]
 hf = ["transformers"]

src/vanna/mock/__init__.py

@@ -0,0 +1,3 @@
+from .embedding import MockEmbedding
+from .llm import MockLLM
+from .vectordb import MockVectorDB

src/vanna/mock/embedding.py

@@ -0,0 +1,11 @@
+from typing import List
+
+from ..base import VannaBase
+
+
+class MockEmbedding(VannaBase):
+    def __init__(self, config=None):
+        pass
+
+    def generate_embedding(self, data: str, **kwargs) -> List[float]:
+        return [1.0, 2.0, 3.0, 4.0, 5.0]

src/vanna/mock/llm.py

@@ -0,0 +1,19 @@
+
+from ..base import VannaBase
+
+
+class MockLLM(VannaBase):
+    def __init__(self, config=None):
+        pass
+
+    def system_message(self, message: str) -> any:
+        return {"role": "system", "content": message}
+
+    def user_message(self, message: str) -> any:
+        return {"role": "user", "content": message}
+
+    def assistant_message(self, message: str) -> any:
+        return {"role": "assistant", "content": message}
+
+    def submit_prompt(self, prompt, **kwargs) -> str:
+        return "Mock LLM response"

src/vanna/mock/vectordb.py

@@ -0,0 +1,55 @@
+import pandas as pd
+
+from ..base import VannaBase
+
+
+class MockVectorDB(VannaBase):
+    def __init__(self, config=None):
+        pass
+
+    def _get_id(self, value: str, **kwargs) -> str:
+        # Hash the value and return the ID
+        return str(hash(value))
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        return self._get_id(ddl)
+
+    def add_documentation(self, doc: str, **kwargs) -> str:
+        return self._get_id(doc)
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        return self._get_id(question)
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        return []
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        return []
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> list:
+        return []
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        return pd.DataFrame({'id': {0: '19546-ddl',
+  1: '91597-sql',
+  2: '133976-sql',
+  3: '59851-doc',
+  4: '73046-sql'},
+ 'training_data_type': {0: 'ddl',
+  1: 'sql',
+  2: 'sql',
+  3: 'documentation',
+  4: 'sql'},
+ 'question': {0: None,
+  1: 'What are the top selling genres?',
+  2: 'What are the low 7 artists by sales?',
+  3: None,
+  4: 'What is the total sales for each customer?'},
+ 'content': {0: 'CREATE TABLE [Invoice]\n(\n    [InvoiceId] INTEGER  NOT NULL,\n    [CustomerId] INTEGER  NOT NULL,\n    [InvoiceDate] DATETIME  NOT NULL,\n    [BillingAddress] NVARCHAR(70),\n    [BillingCity] NVARCHAR(40),\n    [BillingState] NVARCHAR(40),\n    [BillingCountry] NVARCHAR(40),\n    [BillingPostalCode] NVARCHAR(10),\n    [Total] NUMERIC(10,2)  NOT NULL,\n    CONSTRAINT [PK_Invoice] PRIMARY KEY  ([InvoiceId]),\n    FOREIGN KEY ([CustomerId]) REFERENCES [Customer] ([CustomerId]) \n\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n)',
+  1: 'SELECT g.Name AS Genre, SUM(il.Quantity) AS TotalSales\nFROM Genre g\nJOIN Track t ON g.GenreId = t.GenreId\nJOIN InvoiceLine il ON t.TrackId = il.TrackId\nGROUP BY g.GenreId, g.Name\nORDER BY TotalSales DESC;',
+  2: 'SELECT a.ArtistId, a.Name, SUM(il.Quantity) AS TotalSales\nFROM Artist a\nINNER JOIN Album al ON a.ArtistId = al.ArtistId\nINNER JOIN Track t ON al.AlbumId = t.AlbumId\nINNER JOIN InvoiceLine il ON t.TrackId = il.TrackId\nGROUP BY a.ArtistId, a.Name\nORDER BY TotalSales ASC\nLIMIT 7;',
+  3: 'This is a SQLite database. For dates rememeber to use SQLite syntax.',
+  4: 'SELECT c.CustomerId, c.FirstName, c.LastName, SUM(i.Total) AS TotalSales\nFROM Customer c\nJOIN Invoice i ON c.CustomerId = i.CustomerId\nGROUP BY c.CustomerId, c.FirstName, c.LastName;'}})
+
+    def remove_training_data(id: str, **kwargs) -> bool:
+        return True

src/vanna/pinecone/__init__.py

@@ -0,0 +1,3 @@
+from .pinecone_vector import PineconeDB_VectorStore
+
+__all__ = ["PineconeDB_VectorStore"]

src/vanna/pinecone/pinecone_vector.py

@@ -0,0 +1,275 @@
+import json
+from typing import List
+
+from pinecone import Pinecone, PodSpec, ServerlessSpec
+import pandas as pd
+from ..base import VannaBase
+from ..utils import deterministic_uuid
+
+from fastembed import TextEmbedding
+
+
+class PineconeDB_VectorStore(VannaBase):
+    """
+    Vectorstore using PineconeDB
+
+    Args:
+        config (dict): Configuration dictionary. Defaults to {}. You must provide either a Pinecone Client or an API key in the config.
+            - client (Pinecone, optional): Pinecone client. Defaults to None.
+            - api_key (str, optional): Pinecone API key. Defaults to None.
+            - n_results (int, optional): Number of results to return. Defaults to 10.
+            - dimensions (int, optional): Dimensions of the embeddings. Defaults to 384 which coresponds to the dimensions of BAAI/bge-small-en-v1.5.
+            - fastembed_model (str, optional): Fastembed model to use. Defaults to "BAAI/bge-small-en-v1.5".
+            - documentation_namespace (str, optional): Namespace for documentation. Defaults to "documentation".
+            - distance_metric (str, optional): Distance metric to use. Defaults to "cosine".
+            - ddl_namespace (str, optional): Namespace for DDL. Defaults to "ddl".
+            - sql_namespace (str, optional): Namespace for SQL. Defaults to "sql".
+            - index_name (str, optional): Name of the index. Defaults to "vanna-index".
+            - metadata_config (dict, optional): Metadata configuration if using a pinecone pod. Defaults to {}.
+            - server_type (str, optional): Type of Pinecone server to use. Defaults to "serverless". Options are "serverless" or "pod".
+            - podspec (PodSpec, optional): PodSpec configuration if using a pinecone pod. Defaults to PodSpec(environment="us-west-2", pod_type="p1.x1", metadata_config=self.metadata_config).
+            - serverless_spec (ServerlessSpec, optional): ServerlessSpec configuration if using a pinecone serverless index. Defaults to ServerlessSpec(cloud="aws", region="us-west-2").
+    Raises:
+        ValueError: If config is None, api_key is not provided OR client is not provided, client is not an instance of Pinecone, or server_type is not "serverless" or "pod".
+    """
+
+    def __init__(self, config=None):
+        VannaBase.__init__(self, config=config)
+        if config is None:
+            raise ValueError(
+                "config is required, pass either a Pinecone client or an API key in the config."
+            )
+        client = config.get("client")
+        api_key = config.get("api_key")
+        if not api_key and not client:
+            raise ValueError(
+                "api_key is required in config or pass a configured client"
+            )
+        if not client and api_key:
+            self._client = Pinecone(api_key=api_key)
+        elif not isinstance(client, Pinecone):
+            raise ValueError("client must be an instance of Pinecone")
+        else:
+            self._client = client
+
+        self.n_results = config.get("n_results", 10)
+        self.dimensions = config.get("dimensions", 384)
+        self.fastembed_model = config.get("fastembed_model", "BAAI/bge-small-en-v1.5")
+        self.documentation_namespace = config.get(
+            "documentation_namespace", "documentation"
+        )
+        self.distance_metric = config.get("distance_metric", "cosine")
+        self.ddl_namespace = config.get("ddl_namespace", "ddl")
+        self.sql_namespace = config.get("sql_namespace", "sql")
+        self.index_name = config.get("index_name", "vanna-index")
+        self.metadata_config = config.get("metadata_config", {})
+        self.server_type = config.get("server_type", "serverless")
+        if self.server_type not in ["serverless", "pod"]:
+            raise ValueError("server_type must be either 'serverless' or 'pod'")
+        self.podspec = config.get(
+            "podspec",
+            PodSpec(
+                environment="us-west-2",
+                pod_type="p1.x1",
+                metadata_config=self.metadata_config,
+            ),
+        )
+        self.serverless_spec = config.get(
+            "serverless_spec", ServerlessSpec(cloud="aws", region="us-west-2")
+        )
+        self._setup_index()
+
+    def _set_index_host(self, host: str) -> None:
+        self.Index = self._client.Index(host=host)
+
+    def _setup_index(self) -> None:
+        existing_indexes = self._get_indexes()
+        if self.index_name not in existing_indexes and self.server_type == "serverless":
+            self._client.create_index(
+                name=self.index_name,
+                dimension=self.dimensions,
+                metric=self.distance_metric,
+                spec=self.serverless_spec,
+            )
+            pinecone_index_host = self._client.describe_index(self.index_name)["host"]
+            self._set_index_host(pinecone_index_host)
+        elif self.index_name not in existing_indexes and self.server_type == "pod":
+            self._client.create_index(
+                name=self.index_name,
+                dimension=self.dimensions,
+                metric=self.distance_metric,
+                spec=self.podspec,
+            )
+            pinecone_index_host = self._client.describe_index(self.index_name)["host"]
+            self._set_index_host(pinecone_index_host)
+        else:
+            pinecone_index_host = self._client.describe_index(self.index_name)["host"]
+            self._set_index_host(pinecone_index_host)
+
+    def _get_indexes(self) -> list:
+        return [index["name"] for index in self._client.list_indexes()]
+
+    def _check_if_embedding_exists(self, id: str, namespace: str) -> bool:
+        fetch_response = self.Index.fetch(ids=[id], namespace=namespace)
+        if fetch_response["vectors"] == {}:
+            return False
+        return True
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        id = deterministic_uuid(ddl) + "-ddl"
+        if self._check_if_embedding_exists(id=id, namespace=self.ddl_namespace):
+            print(f"DDL with id: {id} already exists in the index. Skipping...")
+            return id
+        self.Index.upsert(
+            vectors=[(id, self.generate_embedding(ddl), {"ddl": ddl})],
+            namespace=self.ddl_namespace,
+        )
+        return id
+
+    def add_documentation(self, doc: str, **kwargs) -> str:
+        id = deterministic_uuid(doc) + "-doc"
+
+        if self._check_if_embedding_exists(
+            id=id, namespace=self.documentation_namespace
+        ):
+            print(
+                f"Documentation with id: {id} already exists in the index. Skipping..."
+            )
+            return id
+        self.Index.upsert(
+            vectors=[(id, self.generate_embedding(doc), {"documentation": doc})],
+            namespace=self.documentation_namespace,
+        )
+        return id
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        question_sql_json = json.dumps(
+            {
+                "question": question,
+                "sql": sql,
+            },
+            ensure_ascii=False,
+        )
+        id = deterministic_uuid(question_sql_json) + "-sql"
+        if self._check_if_embedding_exists(id=id, namespace=self.sql_namespace):
+            print(
+                f"Question-SQL with id: {id} already exists in the index. Skipping..."
+            )
+            return id
+        self.Index.upsert(
+            vectors=[
+                (
+                    id,
+                    self.generate_embedding(question_sql_json),
+                    {"sql": question_sql_json},
+                )
+            ],
+            namespace=self.sql_namespace,
+        )
+        return id
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        res = self.Index.query(
+            namespace=self.ddl_namespace,
+            vector=self.generate_embedding(question),
+            top_k=self.n_results,
+            include_values=True,
+            include_metadata=True,
+        )
+        return [match["metadata"]["ddl"] for match in res["matches"]] if res else []
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        res = self.Index.query(
+            namespace=self.documentation_namespace,
+            vector=self.generate_embedding(question),
+            top_k=self.n_results,
+            include_values=True,
+            include_metadata=True,
+        )
+        return (
+            [match["metadata"]["documentation"] for match in res["matches"]]
+            if res
+            else []
+        )
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> list:
+        res = self.Index.query(
+            namespace=self.sql_namespace,
+            vector=self.generate_embedding(question),
+            top_k=self.n_results,
+            include_values=True,
+            include_metadata=True,
+        )
+        return (
+            [
+                {
+                    key: value
+                    for key, value in json.loads(match["metadata"]["sql"]).items()
+                }
+                for match in res["matches"]
+            ]
+            if res
+            else []
+        )
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        # Pinecone does not support getting all vectors in a namespace, so we have to query for the top_k vectors with a dummy vector
+        df = pd.DataFrame()
+        namespaces = {
+            "sql": self.sql_namespace,
+            "ddl": self.ddl_namespace,
+            "documentation": self.documentation_namespace,
+        }
+
+        for data_type, namespace in namespaces.items():
+            data = self.Index.query(
+                top_k=10000,  # max results that pinecone allows
+                namespace=namespace,
+                include_values=True,
+                include_metadata=True,
+                vector=[0.0] * self.dimensions,
+            )
+
+            if data is not None:
+                id_list = [match["id"] for match in data["matches"]]
+                content_list = [
+                    match["metadata"][data_type] for match in data["matches"]
+                ]
+                question_list = [
+                    (
+                        json.loads(match["metadata"][data_type])["question"]
+                        if data_type == "sql"
+                        else None
+                    )
+                    for match in data["matches"]
+                ]
+
+                df_data = pd.DataFrame(
+                    {
+                        "id": id_list,
+                        "question": question_list,
+                        "content": content_list,
+                    }
+                )
+                df_data["training_data_type"] = data_type
+                df = pd.concat([df, df_data])
+
+        return df
+
+    def remove_training_data(self, id: str, **kwargs) -> bool:
+        if id.endswith("-sql"):
+            self.Index.delete(ids=[id], namespace=self.sql_namespace)
+            return True
+        elif id.endswith("-ddl"):
+            self.Index.delete(ids=[id], namespace=self.ddl_namespace)
+            return True
+        elif id.endswith("-doc"):
+            self.Index.delete(ids=[id], namespace=self.documentation_namespace)
+            return True
+        else:
+            return False
+
+    def generate_embedding(self, data: str, **kwargs) -> List[float]:
+        embedding_model = TextEmbedding(model_name=self.fastembed_model)
+        embedding = next(embedding_model.embed(data))
+        return embedding.tolist()