vanna 0.7.9__py3-none-any.whl → 2.0.0rc1__py3-none-any.whl

vanna/{pgvector → legacy/pgvector}/pgvector.py

@@ -17,7 +17,8 @@ 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.")
+                "A valid 'config' dictionary with a 'connection_string' is required."
+            )
 
         VannaBase.__init__(self, config=config)
 
@@ -29,7 +30,10 @@ class PG_VectorStore(VannaBase):
             self.embedding_function = config.get("embedding_function")
         else:
             from langchain_huggingface import HuggingFaceEmbeddings
-            self.embedding_function = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+
+            self.embedding_function = HuggingFaceEmbeddings(
+                model_name="all-MiniLM-L6-v2"
+            )
 
         self.sql_collection = PGVector(
             embeddings=self.embedding_function,
@@ -95,15 +99,21 @@ class PG_VectorStore(VannaBase):
                 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)
+        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)
+        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)
+        documents = self.documentation_collection.similarity_search(
+            query=question, k=self.n_results
+        )
         return [document.page_content for document in documents]
 
     def train(
@@ -123,7 +133,9 @@ class PG_VectorStore(VannaBase):
             return self.add_documentation(documentation)
 
         if sql and question:
-            return self.add_question_sql(question=question, sql=sql, createdat=createdat)
+            return self.add_question_sql(
+                question=question, sql=sql, createdat=createdat
+            )
 
         if ddl:
             logging.info(f"Adding ddl: {ddl}")
@@ -135,7 +147,9 @@ class PG_VectorStore(VannaBase):
                     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:
+                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:
@@ -153,7 +167,9 @@ class PG_VectorStore(VannaBase):
         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:]
+            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
@@ -162,19 +178,28 @@ class PG_VectorStore(VannaBase):
                     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.")
+                    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.")
+                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}
+                {
+                    "id": custom_id,
+                    "question": question,
+                    "content": content,
+                    "training_data_type": training_data_type,
+                }
             )
 
         # Create a DataFrame from the list of processed rows
@@ -218,7 +243,9 @@ class PG_VectorStore(VannaBase):
         suffix = suffix_map.get(collection_name)
 
         if not suffix:
-            logging.info("Invalid collection name. Choose from 'ddl', 'sql', or 'documentation'.")
+            logging.info(
+                "Invalid collection name. Choose from 'ddl', 'sql', or 'documentation'."
+            )
             return False
 
         # SQL query to delete rows based on the condition
@@ -242,7 +269,9 @@ class PG_VectorStore(VannaBase):
                         )
                         return True
                     else:
-                        logging.info(f"No rows deleted for collection {collection_name}.")
+                        logging.info(
+                            f"No rows deleted for collection {collection_name}."
+                        )
                         return False
                 except Exception as e:
                     logging.error(f"An error occurred: {e}")

vanna/{qdrant → legacy/qdrant}/qdrant.py

@@ -71,12 +71,8 @@ class Qdrant_VectorStore(VannaBase):
         self.documentation_collection_name = config.get(
             "documentation_collection_name", "documentation"
         )
-        self.ddl_collection_name = config.get(
-            "ddl_collection_name", "ddl"
-        )
-        self.sql_collection_name = config.get(
-            "sql_collection_name", "sql"
-        )
+        self.ddl_collection_name = config.get("ddl_collection_name", "ddl")
+        self.sql_collection_name = config.get("sql_collection_name", "sql")
 
         self.id_suffixes = {
             self.ddl_collection_name: "ddl",

vanna/legacy/qianfan/Qianfan_Chat.py

@@ -0,0 +1,170 @@
+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/legacy/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/legacy/qianwen/QianwenAI_chat.py

@@ -0,0 +1,132 @@
+import os
+
+from openai import OpenAI
+
+from ..base import VannaBase
+
+
+class QianWenAI_Chat(VannaBase):
+    def __init__(self, client=None, config=None):
+        VannaBase.__init__(self, config=config)
+
+        # default parameters - can be overrided using config
+        self.temperature = 0.7
+
+        if "temperature" in config:
+            self.temperature = config["temperature"]
+
+        if "api_type" in config:
+            raise Exception(
+                "Passing api_type is now deprecated. Please pass an OpenAI client instead."
+            )
+
+        if "api_base" in config:
+            raise Exception(
+                "Passing api_base is now deprecated. Please pass an OpenAI client instead."
+            )
+
+        if "api_version" in config:
+            raise Exception(
+                "Passing api_version is now deprecated. Please pass an OpenAI client instead."
+            )
+
+        if client is not None:
+            self.client = client
+            return
+
+        if config is None and client is None:
+            self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+            return
+
+        if "api_key" in config:
+            if "base_url" not in config:
+                self.client = OpenAI(
+                    api_key=config["api_key"],
+                    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
+                )
+            else:
+                self.client = OpenAI(
+                    api_key=config["api_key"], base_url=config["base_url"]
+                )
+
+    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:
+        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.chat.completions.create(
+                model=model,
+                messages=prompt,
+                stop=None,
+                temperature=self.temperature,
+            )
+        elif kwargs.get("engine", None) is not None:
+            engine = kwargs.get("engine", None)
+            print(f"Using model {engine} for {num_tokens} tokens (approx)")
+            response = self.client.chat.completions.create(
+                engine=engine,
+                messages=prompt,
+                stop=None,
+                temperature=self.temperature,
+            )
+        elif self.config is not None and "engine" in self.config:
+            print(
+                f"Using engine {self.config['engine']} for {num_tokens} tokens (approx)"
+            )
+            response = self.client.chat.completions.create(
+                engine=self.config["engine"],
+                messages=prompt,
+                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.chat.completions.create(
+                model=self.config["model"],
+                messages=prompt,
+                stop=None,
+                temperature=self.temperature,
+            )
+        else:
+            if num_tokens > 3500:
+                model = "qwen-long"
+            else:
+                model = "qwen-plus"
+
+            print(f"Using model {model} for {num_tokens} tokens (approx)")
+            response = self.client.chat.completions.create(
+                model=model,
+                messages=prompt,
+                stop=None,
+                temperature=self.temperature,
+            )
+
+        # Find the first response from the chatbot that has text in it (some responses may not have text)
+        for choice in response.choices:
+            if "text" in choice:
+                return choice.text
+
+        # If no response with text is found, return the first response's content (which may be empty)
+        return response.choices[0].message.content

vanna/{remote.py → legacy/remote.py}

@@ -8,31 +8,31 @@ import requests
 
 from .base import VannaBase
 from .types import (
-  AccuracyStats,
-  ApiKey,
-  DataFrameJSON,
-  DataResult,
-  Explanation,
-  FullQuestionDocument,
-  NewOrganization,
-  NewOrganizationMember,
-  Organization,
-  OrganizationList,
-  PlotlyResult,
-  Question,
-  QuestionCategory,
-  QuestionId,
-  QuestionList,
-  QuestionSQLPair,
-  QuestionStringList,
-  SQLAnswer,
-  Status,
-  StatusWithId,
-  StringData,
-  TrainingData,
-  UserEmail,
-  UserOTP,
-  Visibility,
+    AccuracyStats,
+    ApiKey,
+    DataFrameJSON,
+    DataResult,
+    Explanation,
+    FullQuestionDocument,
+    NewOrganization,
+    NewOrganizationMember,
+    Organization,
+    OrganizationList,
+    PlotlyResult,
+    Question,
+    QuestionCategory,
+    QuestionId,
+    QuestionList,
+    QuestionSQLPair,
+    QuestionStringList,
+    SQLAnswer,
+    Status,
+    StatusWithId,
+    StringData,
+    TrainingData,
+    UserEmail,
+    UserOTP,
+    Visibility,
 )
 from .vannadb import VannaDB_VectorStore
 
@@ -40,7 +40,9 @@ from .vannadb import VannaDB_VectorStore
 class VannaDefault(VannaDB_VectorStore):
     def __init__(self, model: str, api_key: str, config=None):
         VannaBase.__init__(self, config=config)
-        VannaDB_VectorStore.__init__(self, vanna_model=model, vanna_api_key=api_key, config=config)
+        VannaDB_VectorStore.__init__(
+            self, vanna_model=model, vanna_api_key=api_key, config=config
+        )
 
         self._model = model
         self._api_key = api_key

vanna/{utils.py → legacy/utils.py}

@@ -9,18 +9,14 @@ from .exceptions import ImproperlyConfigured, ValidationError
 
 def validate_config_path(path):
     if not os.path.exists(path):
-        raise ImproperlyConfigured(
-            f'No such configuration file: {path}'
-        )
+        raise ImproperlyConfigured(f"No such configuration file: {path}")
 
     if not os.path.isfile(path):
-        raise ImproperlyConfigured(
-            f'Config should be a file: {path}'
-        )
+        raise ImproperlyConfigured(f"Config should be a file: {path}")
 
     if not os.access(path, os.R_OK):
         raise ImproperlyConfigured(
-            f'Cannot read the config file. Please grant read privileges: {path}'
+            f"Cannot read the config file. Please grant read privileges: {path}"
         )
 
 
@@ -31,13 +27,12 @@ def sanitize_model_name(model_name):
         # Replace spaces with a hyphen
         model_name = model_name.replace(" ", "-")
 
-        if '-' in model_name:
-
+        if "-" in model_name:
             # remove double hyphones
             model_name = re.sub(r"-+", "-", model_name)
-            if '_' in model_name:
+            if "_" in model_name:
                 # If name contains both underscores and hyphen replace all underscores with hyphens
-                model_name = re.sub(r'_', '-', model_name)
+                model_name = re.sub(r"_", "-", model_name)
 
         # Remove special characters only allow underscore
         model_name = re.sub(r"[^a-zA-Z0-9-_]", "", model_name)