vanna 0.6.4__tar.gz → 0.6.6__tar.gz

{vanna-0.6.4 → vanna-0.6.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.6.4
+Version: 0.6.6
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -25,7 +25,8 @@ Requires-Dist: google-cloud-bigquery ; extra == "all"
 Requires-Dist: snowflake-connector-python ; extra == "all"
 Requires-Dist: duckdb ; extra == "all"
 Requires-Dist: openai ; extra == "all"
-Requires-Dist: mistralai ; extra == "all"
+Requires-Dist: qianfan ; extra == "all"
+Requires-Dist: mistralai>=1.0.0 ; extra == "all"
 Requires-Dist: chromadb ; extra == "all"
 Requires-Dist: anthropic ; extra == "all"
 Requires-Dist: zhipuai ; extra == "all"
@@ -55,7 +56,7 @@ 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: mistralai>=1.0.0 ; extra == "mistralai"
 Requires-Dist: PyMySQL ; extra == "mysql"
 Requires-Dist: ollama ; extra == "ollama"
 Requires-Dist: httpx ; extra == "ollama"
@@ -68,6 +69,7 @@ Requires-Dist: psycopg2-binary ; extra == "postgres"
 Requires-Dist: db-dtypes ; extra == "postgres"
 Requires-Dist: qdrant-client ; extra == "qdrant"
 Requires-Dist: fastembed ; extra == "qdrant"
+Requires-Dist: qianfan ; extra == "qianfan"
 Requires-Dist: snowflake-connector-python ; extra == "snowflake"
 Requires-Dist: tox ; extra == "test"
 Requires-Dist: vllm ; extra == "vllm"
@@ -95,6 +97,7 @@ Provides-Extra: opensearch
 Provides-Extra: pinecone
 Provides-Extra: postgres
 Provides-Extra: qdrant
+Provides-Extra: qianfan
 Provides-Extra: snowflake
 Provides-Extra: test
 Provides-Extra: vllm

{vanna-0.6.4 → vanna-0.6.6}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "vanna"
-version = "0.6.4"
+version = "0.6.6"
 authors = [
   { name="Zain Hoda", email="zain@vanna.ai" },
 ]
@@ -33,11 +33,12 @@ 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", "pinecone-client", "pymilvus[model]","weaviate-client"]
+all = ["psycopg2-binary", "db-dtypes", "PyMySQL", "google-cloud-bigquery", "snowflake-connector-python", "duckdb", "openai", "qianfan", "mistralai>=1.0.0", "chromadb", "anthropic", "zhipuai", "marqo", "google-generativeai", "google-cloud-aiplatform", "qdrant-client", "fastembed", "ollama", "httpx", "opensearch-py", "opensearch-dsl", "transformers", "pinecone-client", "pymilvus[model]","weaviate-client"]
 test = ["tox"]
 chromadb = ["chromadb"]
 openai = ["openai"]
-mistralai = ["mistralai"]
+qianfan = ["qianfan"]
+mistralai = ["mistralai>=1.0.0"]
 anthropic = ["anthropic"]
 gemini = ["google-generativeai"]
 marqo = ["marqo"]

{vanna-0.6.4 → vanna-0.6.6}/src/vanna/base/base.py

@@ -437,7 +437,7 @@ class VannaBase(ABC):
         pass
 
     @abstractmethod
-    def remove_training_data(id: str, **kwargs) -> bool:
+    def remove_training_data(self, id: str, **kwargs) -> bool:
         """
         Example:
         ```python
@@ -1276,15 +1276,10 @@ class VannaBase(ABC):
 
         def run_sql_bigquery(sql: str) -> Union[pd.DataFrame, None]:
             if conn:
-                try:
-                    job = conn.query(sql)
-                    df = job.result().to_dataframe()
-                    return df
-                except GoogleAPIError as error:
-                    errors = []
-                    for error in error.errors:
-                        errors.append(error["message"])
-                    raise errors
+                job = conn.query(sql)
+                df = job.result().to_dataframe()
+                return df
+
             return None
 
         self.dialect = "BigQuery SQL"

{vanna-0.6.4 → vanna-0.6.6}/src/vanna/flask/__init__.py

@@ -12,9 +12,9 @@ from flasgger import Swagger
 from flask import Flask, Response, jsonify, request, send_from_directory
 from flask_sock import Sock
 
+from ..base import VannaBase
 from .assets import css_content, html_content, js_content
 from .auth import AuthInterface, NoAuth
-from ..base import VannaBase
 
 
 class Cache(ABC):
@@ -1211,7 +1211,7 @@ class VannaFlaskApp(VannaFlaskAPI):
         self.config["ask_results_correct"] = ask_results_correct
         self.config["followup_questions"] = followup_questions
         self.config["summarization"] = summarization
-        self.config["function_generation"] = function_generation
+        self.config["function_generation"] = function_generation and hasattr(vn, "get_function")
 
         self.index_html_path = index_html_path
         self.assets_folder = assets_folder

vanna-0.6.6/src/vanna/google/__init__.py

@@ -0,0 +1,2 @@
+from .bigquery_vector import BigQuery_VectorStore
+from .gemini_chat import GoogleGeminiChat

vanna-0.6.6/src/vanna/google/bigquery_vector.py

@@ -0,0 +1,230 @@
+import datetime
+import os
+import uuid
+from typing import List, Optional
+
+import pandas as pd
+from google.cloud import bigquery
+
+from ..base import VannaBase
+
+
+class BigQuery_VectorStore(VannaBase):
+    def __init__(self, config: dict, **kwargs):
+        self.config = config
+
+        self.n_results_sql = config.get("n_results_sql", config.get("n_results", 10))
+        self.n_results_documentation = config.get("n_results_documentation", config.get("n_results", 10))
+        self.n_results_ddl = config.get("n_results_ddl", config.get("n_results", 10))
+
+        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.
+            """
+            print("Configuring genai")
+            import google.generativeai as genai
+
+            genai.configure(api_key=config["api_key"])
+
+            self.genai = genai
+        else:
+            # Authenticate using VertexAI
+            from vertexai.language_models import (
+              TextEmbeddingInput,
+              TextEmbeddingModel,
+            )
+
+        if self.config.get("project_id"):
+            self.project_id = self.config.get("project_id")
+        else:
+            self.project_id = os.getenv("GOOGLE_CLOUD_PROJECT")
+
+        if self.project_id is None:
+            raise ValueError("Project ID is not set")
+
+        self.conn = bigquery.Client(project=self.project_id)
+
+        dataset_name = self.config.get('bigquery_dataset_name', 'vanna_managed')
+        self.dataset_id = f"{self.project_id}.{dataset_name}"
+        dataset = bigquery.Dataset(self.dataset_id)
+
+        try:
+            self.conn.get_dataset(self.dataset_id)  # Make an API request.
+            print(f"Dataset {self.dataset_id} already exists")
+        except Exception:
+            # Dataset does not exist, create it
+            dataset.location = "US"
+            self.conn.create_dataset(dataset, timeout=30)  # Make an API request.
+            print(f"Created dataset {self.dataset_id}")
+
+        # Create a table called training_data in the dataset that contains the columns:
+        # id, training_data_type, question, content, embedding, created_at
+
+        self.table_id = f"{self.dataset_id}.training_data"
+        schema = [
+            bigquery.SchemaField("id", "STRING", mode="REQUIRED"),
+            bigquery.SchemaField("training_data_type", "STRING", mode="REQUIRED"),
+            bigquery.SchemaField("question", "STRING", mode="REQUIRED"),
+            bigquery.SchemaField("content", "STRING", mode="REQUIRED"),
+            bigquery.SchemaField("embedding", "FLOAT64", mode="REPEATED"),
+            bigquery.SchemaField("created_at", "TIMESTAMP", mode="REQUIRED"),
+        ]
+
+        table = bigquery.Table(self.table_id, schema=schema)
+
+        try:
+            self.conn.get_table(self.table_id)  # Make an API request.
+            print(f"Table {self.table_id} already exists")
+        except Exception:
+            # Table does not exist, create it
+            self.conn.create_table(table, timeout=30)  # Make an API request.
+            print(f"Created table {self.table_id}")
+
+        # Create VECTOR INDEX IF NOT EXISTS
+        # TODO: This requires 5000 rows before it can be created
+        # vector_index_query = f"""
+        # CREATE VECTOR INDEX IF NOT EXISTS my_index
+        # ON `{self.table_id}`(embedding)
+        # OPTIONS(
+        #     distance_type='COSINE',
+        #     index_type='IVF',
+        #     ivf_options='{{"num_lists": 1000}}'
+        # )
+        # """
+
+        # try:
+        #     self.conn.query(vector_index_query).result()  # Make an API request.
+        #     print(f"Vector index on {self.table_id} created or already exists")
+        # except Exception as e:
+        #     print(f"Failed to create vector index: {e}")
+
+    def store_training_data(self, training_data_type: str, question: str, content: str, embedding: List[float], **kwargs) -> str:
+        id = str(uuid.uuid4())
+        created_at = datetime.datetime.now()
+        self.conn.insert_rows_json(self.table_id, [{
+            "id": id,
+            "training_data_type": training_data_type,
+            "question": question,
+            "content": content,
+            "embedding": embedding,
+            "created_at": created_at.isoformat()
+        }])
+
+        return id
+
+    def fetch_similar_training_data(self, training_data_type: str, question: str, n_results, **kwargs) -> pd.DataFrame:
+        question_embedding = self.generate_question_embedding(question)
+
+        query = f"""
+        SELECT
+            base.id as id,
+            base.question as question,
+            base.training_data_type as training_data_type,
+            base.content as content,
+            distance
+        FROM
+            VECTOR_SEARCH(
+                TABLE `{self.table_id}`,
+                'embedding',
+                (SELECT * FROM UNNEST([STRUCT({question_embedding})])),
+                top_k => 5,
+                distance_type => 'COSINE',
+                options => '{{"use_brute_force":true}}'
+            )
+        WHERE
+            base.training_data_type = '{training_data_type}'
+        """
+
+        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(
+            model="models/text-embedding-004",
+            content=data,
+            task_type="retrieval_query")
+
+        if 'embedding' in result:
+            return result['embedding']
+        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")
+
+        if 'embedding' in result:
+            return result['embedding']
+        else:
+            raise ValueError("No embeddings returned")
+
+        # task = "RETRIEVAL_DOCUMENT"
+        # inputs = [TextEmbeddingInput(data, task)]
+        # embeddings = self.vertex_embedding_model.get_embeddings(inputs)
+
+        # if len(embeddings) == 0:
+        #     raise ValueError("No embeddings returned")
+
+        # return embeddings[0].values
+
+        return result
+
+    def generate_embedding(self, data: str, **kwargs) -> List[float]:
+        return self.generate_storage_embedding(data, **kwargs)
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> list:
+        df = self.fetch_similar_training_data(training_data_type="sql", question=question, n_results=self.n_results_sql)
+
+        # Return a list of dictionaries with only question, sql fields. The content field needs to be renamed to sql
+        return df.rename(columns={"content": "sql"})[["question", "sql"]].to_dict(orient="records")
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        df = self.fetch_similar_training_data(training_data_type="ddl", question=question, n_results=self.n_results_ddl)
+
+        # Return a list of strings of the content
+        return df["content"].tolist()
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        df = self.fetch_similar_training_data(training_data_type="documentation", question=question, n_results=self.n_results_documentation)
+
+        # Return a list of strings of the content
+        return df["content"].tolist()
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        doc = {
+            "question": question,
+            "sql": sql
+        }
+
+        embedding = self.generate_embedding(str(doc))
+
+        return self.store_training_data(training_data_type="sql", question=question, content=sql, embedding=embedding)
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        embedding = self.generate_embedding(ddl)
+
+        return self.store_training_data(training_data_type="ddl", question="", content=ddl, embedding=embedding)
+
+    def add_documentation(self, documentation: str, **kwargs) -> str:
+        embedding = self.generate_embedding(documentation)
+
+        return self.store_training_data(training_data_type="documentation", question="", content=documentation, embedding=embedding)
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        query = f"SELECT id, training_data_type, question, content FROM `{self.table_id}`"
+
+        return self.conn.query(query).result().to_dataframe()
+
+    def remove_training_data(self, id: str, **kwargs) -> bool:
+        query = f"DELETE FROM `{self.table_id}` WHERE id = '{id}'"
+
+        try:
+            self.conn.query(query).result()
+            return True
+
+        except Exception as e:
+            print(f"Failed to remove training data: {e}")
+            return False

{vanna-0.6.4 → vanna-0.6.6}/src/vanna/mistral/mistral.py

@@ -1,5 +1,7 @@
-from mistralai.client import MistralClient
-from mistralai.models.chat_completion import ChatMessage
+import os
+
+from mistralai import Mistral as MistralClient
+from mistralai import UserMessage
 
 from ..base import VannaBase
 
@@ -23,13 +25,13 @@ class Mistral(VannaBase):
         self.model = model
 
     def system_message(self, message: str) -> any:
-        return ChatMessage(role="system", content=message)
+        return {"role": "system", "content": message}
 
     def user_message(self, message: str) -> any:
-        return ChatMessage(role="user", content=message)
+        return {"role": "user", "content": message}
 
     def assistant_message(self, message: str) -> any:
-        return ChatMessage(role="assistant", content=message)
+        return {"role": "assistant", "content": message}
 
     def generate_sql(self, question: str, **kwargs) -> str:
         # Use the super generate_sql
@@ -41,7 +43,7 @@ class Mistral(VannaBase):
         return sql
 
     def submit_prompt(self, prompt, **kwargs) -> str:
-        chat_response = self.client.chat(
+        chat_response = self.client.chat.complete(
             model=self.model,
             messages=prompt,
         )

vanna-0.6.6/src/vanna/qianfan/Qianfan_Chat.py

@@ -0,0 +1,165 @@
+import qianfan
+
+from ..base import VannaBase
+
+
+class Qianfan_Chat(VannaBase):
+  def __init__(self, client=None, config=None):
+    VannaBase.__init__(self, config=config)
+
+    if "api_key" not in config:
+      raise Exception("Missing api_key in config")
+    self.api_key = config["api_key"]
+
+    if "secret_key" not in config:
+      raise Exception("Missing secret_key in config")
+    self.secret_key = config["secret_key"]
+
+    # default parameters - can be overrided using config
+    self.temperature = 0.9
+    self.max_tokens = 1024
+
+    if "temperature" in config:
+      self.temperature = config["temperature"]
+
+    if "max_tokens" in config:
+      self.max_tokens = config["max_tokens"]
+
+    self.model = config["model"] if "model" in config else "ERNIE-Speed"
+
+    if client is not None:
+      self.client = client
+      return
+
+    self.client = qianfan.ChatCompletion(ak=self.api_key,
+                                         sk=self.secret_key)
+
+  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 get_sql_prompt(
+    self,
+    initial_prompt: str,
+    question: str,
+    question_sql_list: list,
+    ddl_list: list,
+    doc_list: list,
+    **kwargs,
+  ):
+    """
+    Example:
+    ```python
+    vn.get_sql_prompt(
+        question="What are the top 10 customers by sales?",
+        question_sql_list=[{"question": "What are the top 10 customers by sales?", "sql": "SELECT * FROM customers ORDER BY sales DESC LIMIT 10"}],
+        ddl_list=["CREATE TABLE customers (id INT, name TEXT, sales DECIMAL)"],
+        doc_list=["The customers table contains information about customers and their sales."],
+    )
+
+    ```
+
+    This method is used to generate a prompt for the LLM to generate SQL.
+
+    Args:
+        question (str): The question to generate SQL for.
+        question_sql_list (list): A list of questions and their corresponding SQL statements.
+        ddl_list (list): A list of DDL statements.
+        doc_list (list): A list of documentation.
+
+    Returns:
+        any: The prompt for the LLM to generate SQL.
+    """
+
+    if initial_prompt is None:
+      initial_prompt = f"You are a {self.dialect} expert. " + \
+                       "Please help to generate a SQL to answer the question based on some context.Please don't give any explanation for your answer. Just only generate a SQL \n"
+
+    initial_prompt = self.add_ddl_to_prompt(
+      initial_prompt, ddl_list, max_tokens=self.max_tokens
+    )
+
+    if self.static_documentation != "":
+      doc_list.append(self.static_documentation)
+
+    initial_prompt = self.add_documentation_to_prompt(
+      initial_prompt, doc_list, max_tokens=self.max_tokens
+    )
+    message_log = []
+
+    if question_sql_list is None or len(question_sql_list) == 0:
+      initial_prompt = initial_prompt + f"question: {question}"
+      message_log.append(self.user_message(initial_prompt))
+    else:
+      for i, example in question_sql_list:
+        if example is None:
+          print("example is None")
+        else:
+          if example is not None and "question" in example and "sql" in example:
+            if i == 0:
+              initial_prompt = initial_prompt + f"question: {example['question']}"
+              message_log.append(self.user_message(initial_prompt))
+            else:
+              message_log.append(self.user_message(example["question"]))
+            message_log.append(self.assistant_message(example["sql"]))
+
+      message_log.append(self.user_message(question))
+    return message_log
+
+  def submit_prompt(self, prompt, **kwargs) -> str:
+    if prompt is None:
+      raise Exception("Prompt is None")
+
+    if len(prompt) == 0:
+      raise Exception("Prompt is empty")
+
+    # Count the number of tokens in the message log
+    # Use 4 as an approximation for the number of characters per token
+    num_tokens = 0
+    for message in prompt:
+      num_tokens += len(message["content"]) / 4
+
+    if kwargs.get("model", None) is not None:
+      model = kwargs.get("model", None)
+      print(
+        f"Using model {model} for {num_tokens} tokens (approx)"
+      )
+      response = self.client.do(
+        model=self.model,
+        messages=prompt,
+        max_output_tokens=self.max_tokens,
+        stop=None,
+        temperature=self.temperature,
+      )
+    elif self.config is not None and "model" in self.config:
+      print(
+        f"Using model {self.config['model']} for {num_tokens} tokens (approx)"
+      )
+      response = self.client.do(
+        model=self.config.get("model"),
+        messages=prompt,
+        max_output_tokens=self.max_tokens,
+        stop=None,
+        temperature=self.temperature,
+      )
+    else:
+      if num_tokens > 3500:
+        model = "ERNIE-Speed-128K"
+      else:
+        model = "ERNIE-Speed-8K"
+
+      print(f"Using model {model} for {num_tokens} tokens (approx)")
+      response = self.client.do(
+        model=model,
+        messages=prompt,
+        max_output_tokens=self.max_tokens,
+        stop=None,
+        temperature=self.temperature,
+      )
+
+    return response.body.get("result")

vanna-0.6.6/src/vanna/qianfan/Qianfan_embeddings.py

@@ -0,0 +1,36 @@
+import qianfan
+
+from ..base import VannaBase
+
+
+class Qianfan_Embeddings(VannaBase):
+  def __init__(self, client=None, config=None):
+    VannaBase.__init__(self, config=config)
+
+    if client is not None:
+      self.client = client
+      return
+
+    if "api_key" not in config:
+      raise Exception("Missing api_key in config")
+    self.api_key = config["api_key"]
+
+    if "secret_key" not in config:
+      raise Exception("Missing secret_key in config")
+    self.secret_key = config["secret_key"]
+
+    self.client = qianfan.Embedding(ak=self.api_key, sk=self.secret_key)
+
+  def generate_embedding(self, data: str, **kwargs) -> list[float]:
+    if self.config is not None and "model" in self.config:
+      embedding = self.client.do(
+        model=self.config["model"],
+        input=[data],
+      )
+    else:
+      embedding = self.client.do(
+        model="bge-large-zh",
+        input=[data],
+      )
+
+    return embedding.get("data")[0]["embedding"]

vanna-0.6.6/src/vanna/qianfan/__init__.py

@@ -0,0 +1,2 @@
+from .Qianfan_Chat import Qianfan_Chat
+from .Qianfan_embeddings import Qianfan_Embeddings

vanna-0.6.4/src/vanna/google/__init__.py

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