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

vanna/google/bigquery_vector.py

@@ -2,6 +2,10 @@ import datetime
 import os
 import uuid
 from typing import List, Optional
+from vertexai.language_models import (
+  TextEmbeddingInput,
+  TextEmbeddingModel
+)
 
 import pandas as pd
 from google.cloud import bigquery
@@ -23,17 +27,15 @@ class BigQuery_VectorStore(VannaBase):
             or set as an environment variable, assign it.
             """
             print("Configuring genai")
+            self.type = "GEMINI"
             import google.generativeai as genai
 
             genai.configure(api_key=config["api_key"])
 
             self.genai = genai
         else:
+            self.type = "VERTEX_AI"
             # Authenticate using VertexAI
-            from vertexai.language_models import (
-              TextEmbeddingInput,
-              TextEmbeddingModel,
-            )
 
         if self.config.get("project_id"):
             self.project_id = self.config.get("project_id")
@@ -139,25 +141,42 @@ class BigQuery_VectorStore(VannaBase):
         results = self.conn.query(query).result().to_dataframe()
         return results
 
-    def generate_question_embedding(self, data: str, **kwargs) -> List[float]:
-        result = self.genai.embed_content(
+    def get_embeddings(self, data: str, task: str) -> List[float]:
+        embeddings = None
+
+        if self.type == "VERTEX_AI":
+          input = [TextEmbeddingInput(data, task)]
+          model = TextEmbeddingModel.from_pretrained("text-embedding-004")
+
+          result = model.get_embeddings(input)
+
+          if len(result) > 0:
+              embeddings = result[0].values
+        else:
+          # Use Gemini Consumer API
+          result = self.genai.embed_content(
             model="models/text-embedding-004",
             content=data,
-            task_type="retrieval_query")
+            task_type=task)
 
-        if 'embedding' in result:
-            return result['embedding']
+          if 'embedding' in result:
+            embeddings = result['embedding']
+
+        return embeddings
+
+    def generate_question_embedding(self, data: str, **kwargs) -> List[float]:
+        result = self.get_embeddings(data, "RETRIEVAL_QUERY")
+
+        if result != None:
+            return result
         else:
             raise ValueError("No embeddings returned")
 
     def generate_storage_embedding(self, data: str, **kwargs) -> List[float]:
-        result = self.genai.embed_content(
-            model="models/text-embedding-004",
-            content=data,
-            task_type="retrieval_document")
+        result = self.get_embeddings(data, "RETRIEVAL_DOCUMENT")
 
-        if 'embedding' in result:
-            return result['embedding']
+        if result != None:
+            return result
         else:
             raise ValueError("No embeddings returned")
 

vanna/google/gemini_chat.py

@@ -15,7 +15,7 @@ class GoogleGeminiChat(VannaBase):
         if "model_name" in config:
             model_name = config["model_name"]
         else:
-            model_name = "gemini-1.0-pro"
+            model_name = "gemini-1.5-pro"
 
         self.google_api_key = None
 
@@ -30,7 +30,7 @@ class GoogleGeminiChat(VannaBase):
             self.chat_model = genai.GenerativeModel(model_name)
         else:
             # Authenticate using VertexAI
-            from vertexai.preview.generative_models import GenerativeModel
+            from vertexai.generative_models import GenerativeModel
             self.chat_model = GenerativeModel(model_name)
 
     def system_message(self, message: str) -> any:

vanna/pgvector/__init__.py

@@ -0,0 +1 @@
+from .pgvector import PG_VectorStore

vanna/pgvector/pgvector.py

@@ -0,0 +1,253 @@
+import ast
+import json
+import logging
+import uuid
+
+import pandas as pd
+from langchain_core.documents import Document
+from langchain_postgres.vectorstores import PGVector
+from sqlalchemy import create_engine, text
+
+from .. import ValidationError
+from ..base import VannaBase
+from ..types import TrainingPlan, TrainingPlanItem
+
+
+class PG_VectorStore(VannaBase):
+    def __init__(self, config=None):
+        if not config or "connection_string" not in config:
+            raise ValueError(
+                "A valid 'config' dictionary with a 'connection_string' is required.")
+
+        VannaBase.__init__(self, config=config)
+
+        if config and "connection_string" in config:
+            self.connection_string = config.get("connection_string")
+            self.n_results = config.get("n_results", 10)
+
+        if config and "embedding_function" in config:
+            self.embedding_function = config.get("embedding_function")
+        else:
+            from langchain_huggingface import HuggingFaceEmbeddings
+            self.embedding_function = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+
+        self.sql_collection = PGVector(
+            embeddings=self.embedding_function,
+            collection_name="sql",
+            connection=self.connection_string,
+        )
+        self.ddl_collection = PGVector(
+            embeddings=self.embedding_function,
+            collection_name="ddl",
+            connection=self.connection_string,
+        )
+        self.documentation_collection = PGVector(
+            embeddings=self.embedding_function,
+            collection_name="documentation",
+            connection=self.connection_string,
+        )
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        question_sql_json = json.dumps(
+            {
+                "question": question,
+                "sql": sql,
+            },
+            ensure_ascii=False,
+        )
+        id = str(uuid.uuid4()) + "-sql"
+        createdat = kwargs.get("createdat")
+        doc = Document(
+            page_content=question_sql_json,
+            metadata={"id": id, "createdat": createdat},
+        )
+        self.sql_collection.add_documents([doc], ids=[doc.metadata["id"]])
+
+        return id
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        _id = str(uuid.uuid4()) + "-ddl"
+        doc = Document(
+            page_content=ddl,
+            metadata={"id": _id},
+        )
+        self.ddl_collection.add_documents([doc], ids=[doc.metadata["id"]])
+        return _id
+
+    def add_documentation(self, documentation: str, **kwargs) -> str:
+        _id = str(uuid.uuid4()) + "-doc"
+        doc = Document(
+            page_content=documentation,
+            metadata={"id": _id},
+        )
+        self.documentation_collection.add_documents([doc], ids=[doc.metadata["id"]])
+        return _id
+
+    def get_collection(self, collection_name):
+        match collection_name:
+            case "sql":
+                return self.sql_collection
+            case "ddl":
+                return self.ddl_collection
+            case "documentation":
+                return self.documentation_collection
+            case _:
+                raise ValueError("Specified collection does not exist.")
+
+    def get_similar_question_sql(self, question: str) -> list:
+        documents = self.sql_collection.similarity_search(query=question, k=self.n_results)
+        return [ast.literal_eval(document.page_content) for document in documents]
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        documents = self.ddl_collection.similarity_search(query=question, k=self.n_results)
+        return [document.page_content for document in documents]
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        documents = self.documentation_collection.similarity_search(query=question, k=self.n_results)
+        return [document.page_content for document in documents]
+
+    def train(
+        self,
+        question: str | None = None,
+        sql: str | None = None,
+        ddl: str | None = None,
+        documentation: str | None = None,
+        plan: TrainingPlan | None = None,
+        createdat: str | None = None,
+    ):
+        if question and not sql:
+            raise ValidationError("Please provide a SQL query.")
+
+        if documentation:
+            logging.info(f"Adding documentation: {documentation}")
+            return self.add_documentation(documentation)
+
+        if sql and question:
+            return self.add_question_sql(question=question, sql=sql, createdat=createdat)
+
+        if ddl:
+            logging.info(f"Adding ddl: {ddl}")
+            return self.add_ddl(ddl)
+
+        if plan:
+            for item in plan._plan:
+                if item.item_type == TrainingPlanItem.ITEM_TYPE_DDL:
+                    self.add_ddl(item.item_value)
+                elif item.item_type == TrainingPlanItem.ITEM_TYPE_IS:
+                    self.add_documentation(item.item_value)
+                elif item.item_type == TrainingPlanItem.ITEM_TYPE_SQL and item.item_name:
+                    self.add_question_sql(question=item.item_name, sql=item.item_value)
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        # Establishing the connection
+        engine = create_engine(self.connection_string)
+
+        # Querying the 'langchain_pg_embedding' table
+        query_embedding = "SELECT cmetadata, document FROM langchain_pg_embedding"
+        df_embedding = pd.read_sql(query_embedding, engine)
+
+        # List to accumulate the processed rows
+        processed_rows = []
+
+        # Process each row in the DataFrame
+        for _, row in df_embedding.iterrows():
+            custom_id = row["cmetadata"]["id"]
+            document = row["document"]
+            training_data_type = "documentation" if custom_id[-3:] == "doc" else custom_id[-3:]
+
+            if training_data_type == "sql":
+                # Convert the document string to a dictionary
+                try:
+                    doc_dict = ast.literal_eval(document)
+                    question = doc_dict.get("question")
+                    content = doc_dict.get("sql")
+                except (ValueError, SyntaxError):
+                    logging.info(f"Skipping row with custom_id {custom_id} due to parsing error.")
+                    continue
+            elif training_data_type in ["documentation", "ddl"]:
+                question = None  # Default value for question
+                content = document
+            else:
+                # If the suffix is not recognized, skip this row
+                logging.info(f"Skipping row with custom_id {custom_id} due to unrecognized training data type.")
+                continue
+
+            # Append the processed data to the list
+            processed_rows.append(
+                {"id": custom_id, "question": question, "content": content, "training_data_type": training_data_type}
+            )
+
+        # Create a DataFrame from the list of processed rows
+        df_processed = pd.DataFrame(processed_rows)
+
+        return df_processed
+
+    def remove_training_data(self, id: str, **kwargs) -> bool:
+        # Create the database engine
+        engine = create_engine(self.connection_string)
+
+        # SQL DELETE statement
+        delete_statement = text(
+            """
+            DELETE FROM langchain_pg_embedding
+            WHERE cmetadata ->> 'id' = :id
+        """
+        )
+
+        # Connect to the database and execute the delete statement
+        with engine.connect() as connection:
+            # Start a transaction
+            with connection.begin() as transaction:
+                try:
+                    result = connection.execute(delete_statement, {"id": id})
+                    # Commit the transaction if the delete was successful
+                    transaction.commit()
+                    # Check if any row was deleted and return True or False accordingly
+                    return result.rowcount > 0
+                except Exception as e:
+                    # Rollback the transaction in case of error
+                    logging.error(f"An error occurred: {e}")
+                    transaction.rollback()
+                    return False
+
+    def remove_collection(self, collection_name: str) -> bool:
+        engine = create_engine(self.connection_string)
+
+        # Determine the suffix to look for based on the collection name
+        suffix_map = {"ddl": "ddl", "sql": "sql", "documentation": "doc"}
+        suffix = suffix_map.get(collection_name)
+
+        if not suffix:
+            logging.info("Invalid collection name. Choose from 'ddl', 'sql', or 'documentation'.")
+            return False
+
+        # SQL query to delete rows based on the condition
+        query = text(
+            f"""
+            DELETE FROM langchain_pg_embedding
+            WHERE cmetadata->>'id' LIKE '%{suffix}'
+        """
+        )
+
+        # Execute the deletion within a transaction block
+        with engine.connect() as connection:
+            with connection.begin() as transaction:
+                try:
+                    result = connection.execute(query)
+                    transaction.commit()  # Explicitly commit the transaction
+                    if result.rowcount > 0:
+                        logging.info(
+                            f"Deleted {result.rowcount} rows from "
+                            f"langchain_pg_embedding where collection is {collection_name}."
+                        )
+                        return True
+                    else:
+                        logging.info(f"No rows deleted for collection {collection_name}.")
+                        return False
+                except Exception as e:
+                    logging.error(f"An error occurred: {e}")
+                    transaction.rollback()  # Rollback in case of error
+                    return False
+
+    def generate_embedding(self, *args, **kwargs):
+        pass

vanna/vllm/vllm.py

@@ -22,6 +22,12 @@ class Vllm(VannaBase):
         else:
             self.auth_key = None
 
+        if "temperature" in config:
+            self.temperature = config["temperature"]
+        else:
+            # default temperature - can be overrided using config
+            self.temperature = 0.7
+
     def system_message(self, message: str) -> any:
         return {"role": "system", "content": message}
 
@@ -68,6 +74,7 @@ class Vllm(VannaBase):
         url = f"{self.host}/v1/chat/completions"
         data = {
             "model": self.model,
+            "temperature": self.temperature,
             "stream": False,
             "messages": prompt,
         }

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.7.3
+Version: 0.7.4
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -46,6 +46,12 @@ Requires-Dist: weaviate-client ; extra == "all"
 Requires-Dist: azure-search-documents ; extra == "all"
 Requires-Dist: azure-identity ; extra == "all"
 Requires-Dist: azure-common ; extra == "all"
+Requires-Dist: faiss-cpu ; extra == "all"
+Requires-Dist: boto ; extra == "all"
+Requires-Dist: boto3 ; extra == "all"
+Requires-Dist: botocore ; extra == "all"
+Requires-Dist: langchain_core ; extra == "all"
+Requires-Dist: langchain_postgres ; extra == "all"
 Requires-Dist: anthropic ; extra == "anthropic"
 Requires-Dist: azure-search-documents ; extra == "azuresearch"
 Requires-Dist: azure-identity ; extra == "azuresearch"
@@ -57,6 +63,8 @@ Requires-Dist: google-cloud-bigquery ; extra == "bigquery"
 Requires-Dist: chromadb ; extra == "chromadb"
 Requires-Dist: clickhouse_connect ; extra == "clickhouse"
 Requires-Dist: duckdb ; extra == "duckdb"
+Requires-Dist: faiss-cpu ; extra == "faiss-cpu"
+Requires-Dist: faiss-gpu ; extra == "faiss-gpu"
 Requires-Dist: google-generativeai ; extra == "gemini"
 Requires-Dist: google-generativeai ; extra == "google"
 Requires-Dist: google-cloud-aiplatform ; extra == "google"
@@ -70,6 +78,7 @@ Requires-Dist: httpx ; extra == "ollama"
 Requires-Dist: openai ; extra == "openai"
 Requires-Dist: opensearch-py ; extra == "opensearch"
 Requires-Dist: opensearch-dsl ; extra == "opensearch"
+Requires-Dist: langchain-postgres>=0.0.12 ; extra == "pgvector"
 Requires-Dist: pinecone-client ; extra == "pinecone"
 Requires-Dist: fastembed ; extra == "pinecone"
 Requires-Dist: psycopg2-binary ; extra == "postgres"
@@ -92,6 +101,8 @@ Provides-Extra: bigquery
 Provides-Extra: chromadb
 Provides-Extra: clickhouse
 Provides-Extra: duckdb
+Provides-Extra: faiss-cpu
+Provides-Extra: faiss-gpu
 Provides-Extra: gemini
 Provides-Extra: google
 Provides-Extra: hf
@@ -102,6 +113,7 @@ Provides-Extra: mysql
 Provides-Extra: ollama
 Provides-Extra: openai
 Provides-Extra: opensearch
+Provides-Extra: pgvector
 Provides-Extra: pinecone
 Provides-Extra: postgres
 Provides-Extra: qdrant

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

@@ -23,8 +23,8 @@ vanna/flask/__init__.py,sha256=jcdaau1tQ142nL1ZsDklk0ilMkEyRxgQZdmsl1IN4LQ,43866
 vanna/flask/assets.py,sha256=af-vact_5HSftltugBpPxzLkAI14Z0lVWcObyVe6eKE,453462
 vanna/flask/auth.py,sha256=UpKxh7W5cd43W0LGch0VqhncKwB78L6dtOQkl1JY5T0,1246
 vanna/google/__init__.py,sha256=6D8rDBjKJJm_jpVn9b4Vc2NR-R779ed_bnHhWmxCJXE,92
-vanna/google/bigquery_vector.py,sha256=rkP94Xd1lNYjU1x3MDLvqmGSPUYtDfQwvlqVmX44jyM,8839
-vanna/google/gemini_chat.py,sha256=j1szC2PamMLFrs0Z4lYPS69i017FYICe-mNObNYFBPQ,1576
+vanna/google/bigquery_vector.py,sha256=mHggjvCsWMt4HK6Y4dAZUPgHi1uytxp2AEQ696TSsJA,9315
+vanna/google/gemini_chat.py,sha256=9xHvxArxHr7OWXHnDRz7wX7KTGbuy6xXxHMLkhOMkis,1568
 vanna/hf/__init__.py,sha256=vD0bIhfLkA1UsvVSF4MAz3Da8aQunkQo3wlDztmMuj0,19
 vanna/hf/hf.py,sha256=N8N5g3xvKDBt3dez2r_U0qATxbl2pN8SVLTZK9CSRA0,3020
 vanna/marqo/__init__.py,sha256=GaAWtJ0B-H5rTY607iLCCrLD7T0zMYM5qWIomEB9gLk,37
@@ -44,6 +44,8 @@ vanna/openai/openai_chat.py,sha256=KU6ynOQ5v7vwrQQ13phXoUXeQUrH6_vmhfiPvWddTrQ,4
 vanna/openai/openai_embeddings.py,sha256=g4pNh9LVcYP9wOoO8ecaccDFWmCUYMInebfHucAa2Gc,1260
 vanna/opensearch/__init__.py,sha256=0unDevWOTs7o8S79TOHUKF1mSiuQbBUVm-7k9jV5WW4,54
 vanna/opensearch/opensearch_vector.py,sha256=VhIcrSyNzWR9ZrqrJnyGFOyuQZs3swfbhr8QyVGI0eI,12226
+vanna/pgvector/__init__.py,sha256=7Wvu9qcNdNvZu26Dn53jhO9YXELm0_YsrwBab4BdgVM,37
+vanna/pgvector/pgvector.py,sha256=dJfm8rswYZvbaIbnjmyRjL071iw4siE0INibsZtaLXY,9919
 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
@@ -58,9 +60,9 @@ 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/vllm/vllm.py,sha256=oCdEjT2KP7gbZk-N7G9bfxB15OTtbwJHvNAceXx_r8g,3077
 vanna/weaviate/__init__.py,sha256=HL6PAl7ePBAkeG8uln-BmM7IUtWohyTPvDfcPzSGSCg,46
 vanna/weaviate/weaviate_vector.py,sha256=tUJIZjEy2mda8CB6C8zeN2SKkEO-UJdLsIqy69skuF0,7584
-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,,
+vanna-0.7.4.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+vanna-0.7.4.dist-info/METADATA,sha256=rGFflMIIcAqHlmt6g0gvHkgl6nBT-w9Y8jhwy3pRaYo,12900
+vanna-0.7.4.dist-info/RECORD,,