vanna 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

vanna/ZhipuAI/ZhipuAI_Chat.py

@@ -220,19 +220,14 @@ class ZhipuAI_Chat(VannaBase):
         if len(prompt) == 0:
             raise Exception("Prompt is empty")
 
-        client = ZhipuAI(api_key=self.api_key)  # 填写您自己的APIKey
+        client = ZhipuAI(api_key=self.api_key)
         response = client.chat.completions.create(
-            model="glm-4",  # 填写需要调用的模型名称
+            model="glm-4",
             max_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
             stop=stop,
             messages=prompt,
         )
-        # print(prompt)
-
-        # print(response)
-
-        # print(f"Cost {response.usage.total_tokens} token")
 
         return response.choices[0].message.content

vanna/base/base.py

@@ -147,19 +147,21 @@ class VannaBase(ABC):
             return False
 
     def generate_followup_questions(
-        self, question: str, sql: str, df: pd.DataFrame, **kwargs
+        self, question: str, sql: str, df: pd.DataFrame, n_questions: int = 5, **kwargs
     ) -> list:
         """
         **Example:**
         ```python
-        vn.generate_followup_questions("What are the top 10 customers by sales?", df)
+        vn.generate_followup_questions("What are the top 10 customers by sales?", sql, df)
         ```
 
         Generate a list of followup questions that you can ask Vanna.AI.
 
         Args:
             question (str): The question that was asked.
+            sql (str): The LLM-generated SQL query.
             df (pd.DataFrame): The results of the SQL query.
+            n_questions (int): Number of follow-up questions to generate.
 
         Returns:
             list: A list of followup questions that you can ask Vanna.AI.
@@ -170,7 +172,7 @@ class VannaBase(ABC):
                 f"You are a helpful data assistant. The user asked the question: '{question}'\n\nThe SQL query for this question was: {sql}\n\nThe following is a pandas DataFrame with the results of the query: \n{df.to_markdown()}\n\n"
             ),
             self.user_message(
-                "Generate a list of followup questions that the user might ask about this data. Respond with a list of questions, one per line. Do not answer with any explanations -- just the questions. Remember that there should be an unambiguous SQL query that can be generated from the question. Prefer questions that are answerable outside of the context of this conversation. Prefer questions that are slight modifications of the SQL query that was generated that allow digging deeper into the data. Each question will be turned into a button that the user can click to generate a new SQL query so don't use 'example' type questions. Each question must have a one-to-one correspondence with an instantiated SQL query."
+                f"Generate a list of {n_questions} followup questions that the user might ask about this data. Respond with a list of questions, one per line. Do not answer with any explanations -- just the questions. Remember that there should be an unambiguous SQL query that can be generated from the question. Prefer questions that are answerable outside of the context of this conversation. Prefer questions that are slight modifications of the SQL query that was generated that allow digging deeper into the data. Each question will be turned into a button that the user can click to generate a new SQL query so don't use 'example' type questions. Each question must have a one-to-one correspondence with an instantiated SQL query."
             ),
         ]
 
@@ -1304,12 +1306,14 @@ class VannaBase(ABC):
         table_column = df.columns[
             df.columns.str.lower().str.contains("table_name")
         ].to_list()[0]
-        column_column = df.columns[
-            df.columns.str.lower().str.contains("column_name")
-        ].to_list()[0]
-        data_type_column = df.columns[
-            df.columns.str.lower().str.contains("data_type")
-        ].to_list()[0]
+        columns = [database_column,
+                    schema_column,
+                    table_column]
+        candidates = ["column_name",
+                      "data_type",
+                      "comment"]
+        matches = df.columns.str.lower().str.contains("|".join(candidates), regex=True)
+        columns += df.columns[matches].to_list()
 
         plan = TrainingPlan([])
 
@@ -1330,15 +1334,7 @@ class VannaBase(ABC):
                         f'{database_column} == "{database}" and {schema_column} == "{schema}" and {table_column} == "{table}"'
                     )
                     doc = f"The following columns are in the {table} table in the {database} database:\n\n"
-                    doc += df_columns_filtered_to_table[
-                        [
-                            database_column,
-                            schema_column,
-                            table_column,
-                            column_column,
-                            data_type_column,
-                        ]
-                    ].to_markdown()
+                    doc += df_columns_filtered_to_table[columns].to_markdown()
 
                     plan._plan.append(
                         TrainingPlanItem(

vanna/chromadb/chromadb_vector.py

@@ -8,6 +8,7 @@ from chromadb.config import Settings
 from chromadb.utils import embedding_functions
 
 from ..base import VannaBase
+from ..utils import deterministic_uuid
 
 default_ef = embedding_functions.DefaultEmbeddingFunction()
 
@@ -65,7 +66,7 @@ class ChromaDB_VectorStore(VannaBase):
             },
             ensure_ascii=False,
         )
-        id = str(uuid.uuid4()) + "-sql"
+        id = deterministic_uuid(question_sql_json) + "-sql"
         self.sql_collection.add(
             documents=question_sql_json,
             embeddings=self.generate_embedding(question_sql_json),
@@ -75,7 +76,7 @@ class ChromaDB_VectorStore(VannaBase):
         return id
 
     def add_ddl(self, ddl: str, **kwargs) -> str:
-        id = str(uuid.uuid4()) + "-ddl"
+        id = deterministic_uuid(ddl) + "-ddl"
         self.ddl_collection.add(
             documents=ddl,
             embeddings=self.generate_embedding(ddl),
@@ -84,7 +85,7 @@ class ChromaDB_VectorStore(VannaBase):
         return id
 
     def add_documentation(self, documentation: str, **kwargs) -> str:
-        id = str(uuid.uuid4()) + "-doc"
+        id = deterministic_uuid(documentation) + "-doc"
         self.documentation_collection.add(
             documents=documentation,
             embeddings=self.generate_embedding(documentation),