vanna 0.6.3__py3-none-any.whl → 0.6.5__py3-none-any.whl

vanna/openai/openai_chat.py

@@ -11,14 +11,10 @@ class OpenAI_Chat(VannaBase):
 
         # default parameters - can be overrided using config
         self.temperature = 0.7
-        self.max_tokens = 500
 
         if "temperature" in config:
             self.temperature = config["temperature"]
 
-        if "max_tokens" in config:
-            self.max_tokens = config["max_tokens"]
-
         if "api_type" in config:
             raise Exception(
                 "Passing api_type is now deprecated. Please pass an OpenAI client instead."
@@ -75,7 +71,6 @@ class OpenAI_Chat(VannaBase):
             response = self.client.chat.completions.create(
                 model=model,
                 messages=prompt,
-                max_tokens=self.max_tokens,
                 stop=None,
                 temperature=self.temperature,
             )
@@ -87,7 +82,6 @@ class OpenAI_Chat(VannaBase):
             response = self.client.chat.completions.create(
                 engine=engine,
                 messages=prompt,
-                max_tokens=self.max_tokens,
                 stop=None,
                 temperature=self.temperature,
             )
@@ -98,7 +92,6 @@ class OpenAI_Chat(VannaBase):
             response = self.client.chat.completions.create(
                 engine=self.config["engine"],
                 messages=prompt,
-                max_tokens=self.max_tokens,
                 stop=None,
                 temperature=self.temperature,
             )
@@ -109,7 +102,6 @@ class OpenAI_Chat(VannaBase):
             response = self.client.chat.completions.create(
                 model=self.config["model"],
                 messages=prompt,
-                max_tokens=self.max_tokens,
                 stop=None,
                 temperature=self.temperature,
             )
@@ -123,7 +115,6 @@ class OpenAI_Chat(VannaBase):
             response = self.client.chat.completions.create(
                 model=model,
                 messages=prompt,
-                max_tokens=self.max_tokens,
                 stop=None,
                 temperature=self.temperature,
             )

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/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/qianfan/__init__.py

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

{vanna-0.6.3.dist-info → vanna-0.6.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.6.3
+Version: 0.6.5
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -16,6 +16,7 @@ Requires-Dist: sqlparse
 Requires-Dist: kaleido
 Requires-Dist: flask
 Requires-Dist: flask-sock
+Requires-Dist: flasgger
 Requires-Dist: sqlalchemy
 Requires-Dist: psycopg2-binary ; extra == "all"
 Requires-Dist: db-dtypes ; extra == "all"
@@ -24,6 +25,7 @@ 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: qianfan ; extra == "all"
 Requires-Dist: mistralai ; extra == "all"
 Requires-Dist: chromadb ; extra == "all"
 Requires-Dist: anthropic ; extra == "all"
@@ -67,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"
@@ -94,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.3.dist-info → vanna-0.6.5.dist-info}/RECORD

@@ -7,15 +7,15 @@ vanna/ZhipuAI/ZhipuAI_embeddings.py,sha256=lUqzJg9fOx7rVFhjdkFjXcDeVGV4aAB5Ss0oE
 vanna/ZhipuAI/__init__.py,sha256=NlsijtcZp5Tj9jkOe9fNcOQND_QsGgu7otODsCLBPr0,116
 vanna/advanced/__init__.py,sha256=oDj9g1JbrbCfp4WWdlr_bhgdMqNleyHgr6VXX6DcEbo,658
 vanna/anthropic/__init__.py,sha256=85s_2mAyyPxc0T_0JEvYeAkEKWJwkwqoyUwSC5dw9Gk,43
-vanna/anthropic/anthropic_chat.py,sha256=Wk0o-NMW1uvR2fhSWxrR_2FqNh-dLprNG4uuVqpqAkY,2615
+vanna/anthropic/anthropic_chat.py,sha256=7X3x8SYwDY28aGyBnt0YNRMG8YY1p_t-foMfKGj8_Oo,2627
 vanna/base/__init__.py,sha256=Sl-HM1RRYzAZoSqmL1CZQmF3ZF-byYTCFQP3JZ2A5MU,28
-vanna/base/base.py,sha256=Cz5fW0Odg5GaRg2svvl6YPeBULKXKMNODAb3Zc1Q8bU,70993
+vanna/base/base.py,sha256=3Du70NrXQMn_LOif2YFPRRVKo4wH5-f6eZcLlXEX0X8,71705
 vanna/bedrock/__init__.py,sha256=hRT2bgJbHEqViLdL-t9hfjSfFdIOkPU2ADBt-B1En-8,46
 vanna/bedrock/bedrock_converse.py,sha256=Nx5kYm-diAfYmsWAnTP5xnv7V84Og69-AP9b3seIe0E,2869
 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
-vanna/flask/__init__.py,sha256=urPrHUqM1mpx96VHiQWVXCy3NQwDh6OsSkm4V4wqccY,30211
+vanna/flask/__init__.py,sha256=r1ucQupb6wuTcjVVKpkdrg6R38FZe6KQoKw9AtcghDQ,42889
 vanna/flask/assets.py,sha256=_UoUr57sS0QL2BuTxAOe9k4yy8T7-fp2NpbRSVtW3IM,451769
 vanna/flask/auth.py,sha256=UpKxh7W5cd43W0LGch0VqhncKwB78L6dtOQkl1JY5T0,1246
 vanna/google/__init__.py,sha256=M-dCxCZcKL4bTQyMLj6r6VRs65YNX9Tl2aoPCuqGm-8,41
@@ -35,7 +35,7 @@ vanna/mock/vectordb.py,sha256=h45znfYMUnttE2BBC8v6TKeMaA58pFJL-5B3OGeRNFI,2681
 vanna/ollama/__init__.py,sha256=4xyu8aHPdnEHg5a-QAMwr5o0ns5wevsp_zkI-ndMO2k,27
 vanna/ollama/ollama.py,sha256=rXa7cfvdlO1E5SLysXIl3IZpIaA2r0RBvV5jX2-upiE,3794
 vanna/openai/__init__.py,sha256=tGkeQ7wTIPsando7QhoSHehtoQVdYLwFbKNlSmCmNeQ,86
-vanna/openai/openai_chat.py,sha256=lm-hUsQxu6Q1t06A2csC037zI4VkMk0wFbQ-_Lj74Wg,4764
+vanna/openai/openai_chat.py,sha256=KU6ynOQ5v7vwrQQ13phXoUXeQUrH6_vmhfiPvWddTrQ,4427
 vanna/openai/openai_embeddings.py,sha256=g4pNh9LVcYP9wOoO8ecaccDFWmCUYMInebfHucAa2Gc,1260
 vanna/opensearch/__init__.py,sha256=0unDevWOTs7o8S79TOHUKF1mSiuQbBUVm-7k9jV5WW4,54
 vanna/opensearch/opensearch_vector.py,sha256=VhIcrSyNzWR9ZrqrJnyGFOyuQZs3swfbhr8QyVGI0eI,12226
@@ -43,6 +43,9 @@ 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=qkTWhGrVSAngJZkrcRQ8YFVHcI9j_ZoOGbF6ZVUUdsU,12567
+vanna/qianfan/Qianfan_Chat.py,sha256=Z-s9MwH22T4KMR8AViAjms6qoj67pHeQkMsbK-aXf1M,5273
+vanna/qianfan/Qianfan_embeddings.py,sha256=TYynAJXlyuZfmoj49h8nU6bXu_GjlXREp3tgfQUca04,954
+vanna/qianfan/__init__.py,sha256=QpR43BjZQZcrcDRkyYcYiS-kyqtYmu23AHDzK0Wy1D0,90
 vanna/types/__init__.py,sha256=Qhn_YscKtJh7mFPCyCDLa2K8a4ORLMGVnPpTbv9uB2U,4957
 vanna/vannadb/__init__.py,sha256=C6UkYocmO6dmzfPKZaWojN0mI5YlZZ9VIbdcquBE58A,48
 vanna/vannadb/vannadb_vector.py,sha256=N8poMYvAojoaOF5gI4STD5pZWK9lBKPvyIjbh9dPBa0,14189
@@ -50,6 +53,6 @@ vanna/vllm/__init__.py,sha256=aNlUkF9tbURdeXAJ8ytuaaF1gYwcG3ny1MfNl_cwQYg,23
 vanna/vllm/vllm.py,sha256=oM_aA-1Chyl7T_Qc_yRKlL6oSX1etsijY9zQdjeMGMQ,2827
 vanna/weaviate/__init__.py,sha256=HL6PAl7ePBAkeG8uln-BmM7IUtWohyTPvDfcPzSGSCg,46
 vanna/weaviate/weaviate_vector.py,sha256=GEiu4Vd9w-7j10aB-zTxJ8gefqe_F-LUUGvttFs1vlg,7539
-vanna-0.6.3.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-vanna-0.6.3.dist-info/METADATA,sha256=-CFXJd2aKy2tq2Gxipp-52YCPHJBlI0wIOPzcOV5s80,11865
-vanna-0.6.3.dist-info/RECORD,,
+vanna-0.6.5.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+vanna-0.6.5.dist-info/METADATA,sha256=77ggtzQplTTDdsOI0_U4k-t5UwHVmI5O3TSkpm5OXzY,11997
+vanna-0.6.5.dist-info/RECORD,,