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

vanna/legacy/bedrock/__init__.py

@@ -0,0 +1 @@
+from .bedrock_converse import Bedrock_Converse

vanna/{bedrock → legacy/bedrock}/bedrock_converse.py

@@ -6,6 +6,7 @@ try:
 except ImportError:
     raise ImportError("Please install boto3 and botocore to use Amazon Bedrock models")
 
+
 class Bedrock_Converse(VannaBase):
     def __init__(self, client=None, config=None):
         VannaBase.__init__(self, config=config)
@@ -13,29 +14,27 @@ class Bedrock_Converse(VannaBase):
         # default parameters
         self.temperature = 0.0
         self.max_tokens = 500
-        
+
         if client is None:
             raise ValueError(
                 "A valid Bedrock runtime client must be provided to invoke Bedrock models"
             )
         else:
             self.client = client
-        
+
         if config is None:
             raise ValueError(
                 "Config is required with model_id and inference parameters"
             )
-        
+
         if "modelId" not in config:
-            raise ValueError(
-                "config must contain a modelId to invoke"
-            )
+            raise ValueError("config must contain a modelId to invoke")
         else:
             self.model = config["modelId"]
-        
+
         if "temperature" in config:
             self.temperature = config["temperature"]
-        
+
         if "max_tokens" in config:
             self.max_tokens = config["max_tokens"]
 
@@ -51,7 +50,7 @@ class Bedrock_Converse(VannaBase):
     def submit_prompt(self, prompt, **kwargs) -> str:
         inference_config = {
             "temperature": self.temperature,
-            "maxTokens": self.max_tokens
+            "maxTokens": self.max_tokens,
         }
         additional_model_fields = {
             "top_p": 1,  # setting top_p value for nucleus sampling
@@ -64,13 +63,15 @@ class Bedrock_Converse(VannaBase):
             if role == "system":
                 system_message = prompt_message["content"]
             else:
-                no_system_prompt.append({"role": role, "content":[{"text": prompt_message["content"]}]})
+                no_system_prompt.append(
+                    {"role": role, "content": [{"text": prompt_message["content"]}]}
+                )
 
         converse_api_params = {
             "modelId": self.model,
             "messages": no_system_prompt,
             "inferenceConfig": inference_config,
-            "additionalModelRequestFields": additional_model_fields
+            "additionalModelRequestFields": additional_model_fields,
         }
 
         if system_message:
@@ -82,4 +83,4 @@ class Bedrock_Converse(VannaBase):
             return text_content
         except ClientError as err:
             message = err.response["Error"]["Message"]
-            raise Exception(f"A Bedrock client error occurred: {message}")
+            raise Exception(f"A Bedrock client error occurred: {message}")

vanna/{chromadb → legacy/chromadb}/chromadb_vector.py

@@ -23,7 +23,9 @@ class ChromaDB_VectorStore(VannaBase):
         curr_client = config.get("client", "persistent")
         collection_metadata = config.get("collection_metadata", None)
         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_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 curr_client == "persistent":

vanna/legacy/cohere/__init__.py

@@ -0,0 +1,2 @@
+from .cohere_chat import Cohere_Chat
+from .cohere_embeddings import Cohere_Embeddings

vanna/{cohere → legacy/cohere}/cohere_chat.py

@@ -25,15 +25,17 @@ class Cohere_Chat(VannaBase):
 
         # Check for API key in environment variable
         api_key = os.getenv("COHERE_API_KEY")
-        
+
         # Check for API key in config
         if config is not None and "api_key" in config:
             api_key = config["api_key"]
-            
+
         # Validate API key
         if not api_key:
-            raise ValueError("Cohere API key is required. Please provide it via config or set the COHERE_API_KEY environment variable.")
-            
+            raise ValueError(
+                "Cohere API key is required. Please provide it via config or set the COHERE_API_KEY environment variable."
+            )
+
         # Initialize client with validated API key
         self.client = OpenAI(
             base_url="https://api.cohere.ai/compatibility/v1",
@@ -41,7 +43,10 @@ class Cohere_Chat(VannaBase):
         )
 
     def system_message(self, message: str) -> any:
-        return {"role": "developer", "content": message}  # Cohere uses 'developer' for system role
+        return {
+            "role": "developer",
+            "content": message,
+        }  # Cohere uses 'developer' for system role
 
     def user_message(self, message: str) -> any:
         return {"role": "user", "content": message}
@@ -74,21 +79,21 @@ class Cohere_Chat(VannaBase):
                 messages=prompt,
                 temperature=self.temperature,
             )
-            
+
             # Check if response has expected structure
-            if not response or not hasattr(response, 'choices') or not response.choices:
+            if not response or not hasattr(response, "choices") or not response.choices:
                 raise ValueError("Received empty or malformed response from API")
-                
-            if not response.choices[0] or not hasattr(response.choices[0], 'message'):
+
+            if not response.choices[0] or not hasattr(response.choices[0], "message"):
                 raise ValueError("Response is missing expected 'message' field")
-                
-            if not hasattr(response.choices[0].message, 'content'):
+
+            if not hasattr(response.choices[0].message, "content"):
                 raise ValueError("Response message is missing expected 'content' field")
-                
+
             return response.choices[0].message.content
-            
+
         except Exception as e:
             # Log the error and raise a more informative exception
             error_msg = f"Error processing Cohere chat response: {str(e)}"
             print(error_msg)
-            raise Exception(error_msg) 
+            raise Exception(error_msg)

vanna/{cohere → legacy/cohere}/cohere_embeddings.py

@@ -8,10 +8,10 @@ from ..base import VannaBase
 class Cohere_Embeddings(VannaBase):
     def __init__(self, client=None, config=None):
         VannaBase.__init__(self, config=config)
-        
+
         # Default embedding model
         self.model = "embed-multilingual-v3.0"
-        
+
         if config is not None and "model" in config:
             self.model = config["model"]
 
@@ -21,15 +21,17 @@ class Cohere_Embeddings(VannaBase):
 
         # Check for API key in environment variable
         api_key = os.getenv("COHERE_API_KEY")
-        
+
         # Check for API key in config
         if config is not None and "api_key" in config:
             api_key = config["api_key"]
-            
+
         # Validate API key
         if not api_key:
-            raise ValueError("Cohere API key is required. Please provide it via config or set the COHERE_API_KEY environment variable.")
-            
+            raise ValueError(
+                "Cohere API key is required. Please provide it via config or set the COHERE_API_KEY environment variable."
+            )
+
         # Initialize client with validated API key
         self.client = OpenAI(
             base_url="https://api.cohere.ai/compatibility/v1",
@@ -39,33 +41,37 @@ class Cohere_Embeddings(VannaBase):
     def generate_embedding(self, data: str, **kwargs) -> list[float]:
         if not data:
             raise ValueError("Cannot generate embedding for empty input data")
-            
+
         # Use model from kwargs, config, or default
         model = kwargs.get("model", self.model)
         if self.config is not None and "model" in self.config and model == self.model:
             model = self.config["model"]
-        
-        try:    
+
+        try:
             embedding = self.client.embeddings.create(
                 model=model,
                 input=data,
                 encoding_format="float",  # Ensure we get float values
             )
-            
+
             # Check if response has expected structure
-            if not embedding or not hasattr(embedding, 'data') or not embedding.data:
-                raise ValueError("Received empty or malformed embedding response from API")
-                
-            if not embedding.data[0] or not hasattr(embedding.data[0], 'embedding'):
-                raise ValueError("Embedding response is missing expected 'embedding' field")
-                
+            if not embedding or not hasattr(embedding, "data") or not embedding.data:
+                raise ValueError(
+                    "Received empty or malformed embedding response from API"
+                )
+
+            if not embedding.data[0] or not hasattr(embedding.data[0], "embedding"):
+                raise ValueError(
+                    "Embedding response is missing expected 'embedding' field"
+                )
+
             if not embedding.data[0].embedding:
                 raise ValueError("Received empty embedding vector")
-                
+
             return embedding.data[0].embedding
-            
+
         except Exception as e:
             # Log the error and raise a more informative exception
             error_msg = f"Error generating embedding with Cohere: {str(e)}"
             print(error_msg)
-            raise Exception(error_msg) 
+            raise Exception(error_msg)

vanna/{deepseek → legacy/deepseek}/deepseek_chat.py

@@ -5,7 +5,6 @@ from openai import OpenAI
 from ..base import VannaBase
 
 
-
 # from vanna.chromadb import ChromaDB_VectorStore
 
 # class DeepSeekVanna(ChromaDB_VectorStore, DeepSeekChat):
@@ -27,12 +26,12 @@ class DeepSeekChat(VannaBase):
 
         if "model" not in config:
             raise ValueError("config must contain a DeepSeek model")
-    
+
         api_key = config["api_key"]
         model = config["model"]
         self.model = model
         self.client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com/v1")
-        
+
     def system_message(self, message: str) -> any:
         return {"role": "system", "content": message}
 
@@ -45,10 +44,10 @@ class DeepSeekChat(VannaBase):
     def generate_sql(self, question: str, **kwargs) -> str:
         # 使用父类的 generate_sql
         sql = super().generate_sql(question, **kwargs)
-        
+
         # 替换 "\_" 为 "_"
         sql = sql.replace("\\_", "_")
-        
+
         return sql
 
     def submit_prompt(self, prompt, **kwargs) -> str:
@@ -56,5 +55,5 @@ class DeepSeekChat(VannaBase):
             model=self.model,
             messages=prompt,
         )
-        
+
         return chat_response.choices[0].message.content

vanna/legacy/faiss/__init__.py

@@ -0,0 +1 @@
+from .faiss import FAISS

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

@@ -1,4 +1,4 @@
-import os 
+import os
 import json
 import uuid
 from typing import List, Dict, Any
@@ -10,13 +10,14 @@ import pandas as pd
 from ..base import VannaBase
 from ..exceptions import DependencyError
 
+
 class FAISS(VannaBase):
     def __init__(self, config=None):
         if config is None:
             config = {}
-        
+
         VannaBase.__init__(self, config=config)
-        
+
         try:
             import faiss
         except ImportError:
@@ -30,34 +31,48 @@ class FAISS(VannaBase):
             raise DependencyError(
                 "SentenceTransformer is not installed. Please install it with 'pip install sentence-transformers'."
             )
-        
+
         self.path = config.get("path", ".")
-        self.embedding_dim = config.get('embedding_dim', 384)
-        self.n_results_sql = config.get('n_results_sql', config.get("n_results", 10))
-        self.n_results_ddl = config.get('n_results_ddl', config.get("n_results", 10))
-        self.n_results_documentation = config.get('n_results_documentation', config.get("n_results", 10))
+        self.embedding_dim = config.get("embedding_dim", 384)
+        self.n_results_sql = config.get("n_results_sql", config.get("n_results", 10))
+        self.n_results_ddl = config.get("n_results_ddl", config.get("n_results", 10))
+        self.n_results_documentation = config.get(
+            "n_results_documentation", config.get("n_results", 10)
+        )
         self.curr_client = config.get("client", "persistent")
 
-        if self.curr_client == 'persistent':
-            self.sql_index = self._load_or_create_index('sql_index.faiss')
-            self.ddl_index = self._load_or_create_index('ddl_index.faiss')
-            self.doc_index = self._load_or_create_index('doc_index.faiss')
-        elif self.curr_client == 'in-memory':
+        if self.curr_client == "persistent":
+            self.sql_index = self._load_or_create_index("sql_index.faiss")
+            self.ddl_index = self._load_or_create_index("ddl_index.faiss")
+            self.doc_index = self._load_or_create_index("doc_index.faiss")
+        elif self.curr_client == "in-memory":
             self.sql_index = faiss.IndexFlatL2(self.embedding_dim)
             self.ddl_index = faiss.IndexFlatL2(self.embedding_dim)
             self.doc_index = faiss.IndexFlatL2(self.embedding_dim)
-        elif isinstance(self.curr_client, list) and len(self.curr_client) == 3 and all(isinstance(idx, faiss.Index) for idx in self.curr_client):
+        elif (
+            isinstance(self.curr_client, list)
+            and len(self.curr_client) == 3
+            and all(isinstance(idx, faiss.Index) for idx in self.curr_client)
+        ):
             self.sql_index = self.curr_client[0]
             self.ddl_index = self.curr_client[1]
             self.doc_index = self.curr_client[2]
         else:
-            raise ValueError(f"Unsupported storage type was set in config: {self.curr_client}")
-
-        self.sql_metadata: List[Dict[str, Any]] = self._load_or_create_metadata('sql_metadata.json')
-        self.ddl_metadata: List[Dict[str, str]] = self._load_or_create_metadata('ddl_metadata.json')
-        self.doc_metadata: List[Dict[str, str]] = self._load_or_create_metadata('doc_metadata.json')
+            raise ValueError(
+                f"Unsupported storage type was set in config: {self.curr_client}"
+            )
 
-        model_name = config.get('embedding_model', 'all-MiniLM-L6-v2')
+        self.sql_metadata: List[Dict[str, Any]] = self._load_or_create_metadata(
+            "sql_metadata.json"
+        )
+        self.ddl_metadata: List[Dict[str, str]] = self._load_or_create_metadata(
+            "ddl_metadata.json"
+        )
+        self.doc_metadata: List[Dict[str, str]] = self._load_or_create_metadata(
+            "doc_metadata.json"
+        )
+
+        model_name = config.get("embedding_model", "all-MiniLM-L6-v2")
         self.embedding_model = SentenceTransformer(model_name)
 
     def _load_or_create_index(self, filename):
@@ -69,25 +84,26 @@ class FAISS(VannaBase):
     def _load_or_create_metadata(self, filename):
         filepath = os.path.join(self.path, filename)
         if os.path.exists(filepath):
-            with open(filepath, 'r') as f:
+            with open(filepath, "r") as f:
                 return json.load(f)
         return []
 
     def _save_index(self, index, filename):
-        if self.curr_client == 'persistent':
+        if self.curr_client == "persistent":
             filepath = os.path.join(self.path, filename)
             faiss.write_index(index, filepath)
 
     def _save_metadata(self, metadata, filename):
-        if self.curr_client == 'persistent':
+        if self.curr_client == "persistent":
             filepath = os.path.join(self.path, filename)
-            with open(filepath, 'w') as f:
+            with open(filepath, "w") as f:
                 json.dump(metadata, f)
 
     def generate_embedding(self, data: str, **kwargs) -> List[float]:
         embedding = self.embedding_model.encode(data)
-        assert embedding.shape[0] == self.embedding_dim, \
+        assert embedding.shape[0] == self.embedding_dim, (
             f"Embedding dimension mismatch: expected {self.embedding_dim}, got {embedding.shape[0]}"
+        )
         return embedding.tolist()
 
     def _add_to_index(self, index, metadata_list, text, extra_metadata=None) -> str:
@@ -96,81 +112,119 @@ class FAISS(VannaBase):
         entry_id = str(uuid.uuid4())
         metadata_list.append({"id": entry_id, **(extra_metadata or {})})
         return entry_id
-    
+
     def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
-        entry_id = self._add_to_index(self.sql_index, self.sql_metadata, question + " " + sql, {"question": question, "sql": sql})
-        self._save_index(self.sql_index, 'sql_index.faiss')
-        self._save_metadata(self.sql_metadata, 'sql_metadata.json')
+        entry_id = self._add_to_index(
+            self.sql_index,
+            self.sql_metadata,
+            question + " " + sql,
+            {"question": question, "sql": sql},
+        )
+        self._save_index(self.sql_index, "sql_index.faiss")
+        self._save_metadata(self.sql_metadata, "sql_metadata.json")
         return entry_id
 
     def add_ddl(self, ddl: str, **kwargs) -> str:
-        entry_id = self._add_to_index(self.ddl_index, self.ddl_metadata, ddl, {"ddl": ddl})
-        self._save_index(self.ddl_index, 'ddl_index.faiss')
-        self._save_metadata(self.ddl_metadata, 'ddl_metadata.json')
+        entry_id = self._add_to_index(
+            self.ddl_index, self.ddl_metadata, ddl, {"ddl": ddl}
+        )
+        self._save_index(self.ddl_index, "ddl_index.faiss")
+        self._save_metadata(self.ddl_metadata, "ddl_metadata.json")
         return entry_id
 
     def add_documentation(self, documentation: str, **kwargs) -> str:
-        entry_id = self._add_to_index(self.doc_index, self.doc_metadata, documentation, {"documentation": documentation})
-        self._save_index(self.doc_index, 'doc_index.faiss')
-        self._save_metadata(self.doc_metadata, 'doc_metadata.json')
+        entry_id = self._add_to_index(
+            self.doc_index,
+            self.doc_metadata,
+            documentation,
+            {"documentation": documentation},
+        )
+        self._save_index(self.doc_index, "doc_index.faiss")
+        self._save_metadata(self.doc_metadata, "doc_metadata.json")
         return entry_id
 
     def _get_similar(self, index, metadata_list, text, n_results) -> list:
         embedding = self.generate_embedding(text)
         D, I = index.search(np.array([embedding], dtype=np.float32), k=n_results)
-        return [] if len(I[0]) == 0 or I[0][0] == -1 else [metadata_list[i] for i in I[0]]
+        return (
+            [] if len(I[0]) == 0 or I[0][0] == -1 else [metadata_list[i] for i in I[0]]
+        )
 
     def get_similar_question_sql(self, question: str, **kwargs) -> list:
-        return self._get_similar(self.sql_index, self.sql_metadata, question, self.n_results_sql)
-    
+        return self._get_similar(
+            self.sql_index, self.sql_metadata, question, self.n_results_sql
+        )
+
     def get_related_ddl(self, question: str, **kwargs) -> list:
-        return [metadata["ddl"] for metadata in self._get_similar(self.ddl_index, self.ddl_metadata, question, self.n_results_ddl)]
+        return [
+            metadata["ddl"]
+            for metadata in self._get_similar(
+                self.ddl_index, self.ddl_metadata, question, self.n_results_ddl
+            )
+        ]
 
     def get_related_documentation(self, question: str, **kwargs) -> list:
-        return [metadata["documentation"] for metadata in self._get_similar(self.doc_index, self.doc_metadata, question, self.n_results_documentation)]
+        return [
+            metadata["documentation"]
+            for metadata in self._get_similar(
+                self.doc_index,
+                self.doc_metadata,
+                question,
+                self.n_results_documentation,
+            )
+        ]
 
     def get_training_data(self, **kwargs) -> pd.DataFrame:
         sql_data = pd.DataFrame(self.sql_metadata)
-        sql_data['training_data_type'] = 'sql'
+        sql_data["training_data_type"] = "sql"
 
         ddl_data = pd.DataFrame(self.ddl_metadata)
-        ddl_data['training_data_type'] = 'ddl'
+        ddl_data["training_data_type"] = "ddl"
 
         doc_data = pd.DataFrame(self.doc_metadata)
-        doc_data['training_data_type'] = 'documentation'
+        doc_data["training_data_type"] = "documentation"
 
         return pd.concat([sql_data, ddl_data, doc_data], ignore_index=True)
 
     def remove_training_data(self, id: str, **kwargs) -> bool:
         for metadata_list, index, index_name in [
-            (self.sql_metadata, self.sql_index, 'sql_index.faiss'),
-            (self.ddl_metadata, self.ddl_index, 'ddl_index.faiss'),
-            (self.doc_metadata, self.doc_index, 'doc_index.faiss')
+            (self.sql_metadata, self.sql_index, "sql_index.faiss"),
+            (self.ddl_metadata, self.ddl_index, "ddl_index.faiss"),
+            (self.doc_metadata, self.doc_index, "doc_index.faiss"),
         ]:
             for i, item in enumerate(metadata_list):
-                if item['id'] == id:
+                if item["id"] == id:
                     del metadata_list[i]
                     new_index = faiss.IndexFlatL2(self.embedding_dim)
-                    embeddings = [self.generate_embedding(json.dumps(m)) for m in metadata_list]
+                    embeddings = [
+                        self.generate_embedding(json.dumps(m)) for m in metadata_list
+                    ]
                     if embeddings:
                         new_index.add(np.array(embeddings, dtype=np.float32))
-                    setattr(self, index_name.split('.')[0], new_index)
-                    
-                    if self.curr_client == 'persistent':
+                    setattr(self, index_name.split(".")[0], new_index)
+
+                    if self.curr_client == "persistent":
                         self._save_index(new_index, index_name)
-                        self._save_metadata(metadata_list, f"{index_name.split('.')[0]}_metadata.json")
-                    
+                        self._save_metadata(
+                            metadata_list, f"{index_name.split('.')[0]}_metadata.json"
+                        )
+
                     return True
         return False
 
     def remove_collection(self, collection_name: str) -> bool:
         if collection_name in ["sql", "ddl", "documentation"]:
-            setattr(self, f"{collection_name}_index", faiss.IndexFlatL2(self.embedding_dim))
+            setattr(
+                self, f"{collection_name}_index", faiss.IndexFlatL2(self.embedding_dim)
+            )
             setattr(self, f"{collection_name}_metadata", [])
-            
-            if self.curr_client == 'persistent':
-                self._save_index(getattr(self, f"{collection_name}_index"), f"{collection_name}_index.faiss")
+
+            if self.curr_client == "persistent":
+                self._save_index(
+                    getattr(self, f"{collection_name}_index"),
+                    f"{collection_name}_index.faiss",
+                )
                 self._save_metadata([], f"{collection_name}_metadata.json")
-            
+
             return True
-        return False
+        return False