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

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.4.dist-info → vanna-0.6.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.6.4
+Version: 0.6.5
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -25,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"
@@ -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.dist-info → vanna-0.6.5.dist-info}/RECORD

@@ -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.4.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-vanna-0.6.4.dist-info/METADATA,sha256=LqIi4Hg1y_aTEH79PX48nnY1TM-u6ese9K8Os9Cqkg0,11889
-vanna-0.6.4.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,,