vanna 0.7.6__tar.gz → 0.7.7__tar.gz

vanna-0.7.7/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Vanna.AI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

{vanna-0.7.6 → vanna-0.7.7}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: vanna
-Version: 0.7.6
+Version: 0.7.7
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -8,6 +8,7 @@ Description-Content-Type: text/markdown
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
+License-File: LICENSE
 Requires-Dist: requests
 Requires-Dist: tabulate
 Requires-Dist: plotly
@@ -40,7 +41,7 @@ Requires-Dist: httpx ; extra == "all"
 Requires-Dist: opensearch-py ; extra == "all"
 Requires-Dist: opensearch-dsl ; extra == "all"
 Requires-Dist: transformers ; extra == "all"
-Requires-Dist: pinecone-client ; extra == "all"
+Requires-Dist: pinecone ; extra == "all"
 Requires-Dist: pymilvus[model] ; extra == "all"
 Requires-Dist: weaviate-client ; extra == "all"
 Requires-Dist: azure-search-documents ; extra == "all"
@@ -86,7 +87,7 @@ Requires-Dist: langchain-huggingface ; extra == "opensearch"
 Requires-Dist: oracledb ; extra == "oracle"
 Requires-Dist: chromadb ; extra == "oracle"
 Requires-Dist: langchain-postgres>=0.0.12 ; extra == "pgvector"
-Requires-Dist: pinecone-client ; extra == "pinecone"
+Requires-Dist: pinecone ; extra == "pinecone"
 Requires-Dist: fastembed ; extra == "pinecone"
 Requires-Dist: psycopg2-binary ; extra == "postgres"
 Requires-Dist: db-dtypes ; extra == "postgres"

{vanna-0.7.6 → vanna-0.7.7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "vanna"
-version = "0.7.6"
+version = "0.7.7"
 authors = [
   { name="Zain Hoda", email="zain@vanna.ai" },
 ]
@@ -33,7 +33,7 @@ bigquery = ["google-cloud-bigquery"]
 snowflake = ["snowflake-connector-python"]
 duckdb = ["duckdb"]
 google = ["google-generativeai", "google-cloud-aiplatform"]
-all = ["psycopg2-binary", "db-dtypes", "PyMySQL", "google-cloud-bigquery", "snowflake-connector-python", "duckdb", "openai", "qianfan", "mistralai>=1.0.0", "chromadb", "anthropic", "zhipuai", "marqo", "google-generativeai", "google-cloud-aiplatform", "qdrant-client", "fastembed", "ollama", "httpx", "opensearch-py", "opensearch-dsl", "transformers", "pinecone-client", "pymilvus[model]","weaviate-client", "azure-search-documents", "azure-identity", "azure-common", "faiss-cpu", "boto", "boto3", "botocore", "langchain_core", "langchain_postgres", "langchain-community", "langchain-huggingface", "xinference-client"]
+all = ["psycopg2-binary", "db-dtypes", "PyMySQL", "google-cloud-bigquery", "snowflake-connector-python", "duckdb", "openai", "qianfan", "mistralai>=1.0.0", "chromadb", "anthropic", "zhipuai", "marqo", "google-generativeai", "google-cloud-aiplatform", "qdrant-client", "fastembed", "ollama", "httpx", "opensearch-py", "opensearch-dsl", "transformers", "pinecone", "pymilvus[model]","weaviate-client", "azure-search-documents", "azure-identity", "azure-common", "faiss-cpu", "boto", "boto3", "botocore", "langchain_core", "langchain_postgres", "langchain-community", "langchain-huggingface", "xinference-client"]
 test = ["tox"]
 chromadb = ["chromadb"]
 openai = ["openai"]
@@ -46,7 +46,7 @@ zhipuai = ["zhipuai"]
 ollama = ["ollama", "httpx"]
 qdrant = ["qdrant-client", "fastembed"]
 vllm = ["vllm"]
-pinecone = ["pinecone-client", "fastembed"]
+pinecone = ["pinecone", "fastembed"]
 opensearch = ["opensearch-py", "opensearch-dsl", "langchain-community", "langchain-huggingface"]
 hf = ["transformers"]
 milvus = ["pymilvus[model]"]

vanna-0.7.7/src/vanna/cohere/__init__.py

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

vanna-0.7.7/src/vanna/cohere/cohere_chat.py

@@ -0,0 +1,94 @@
+import os
+
+from openai import OpenAI
+
+from ..base import VannaBase
+
+
+class Cohere_Chat(VannaBase):
+    def __init__(self, client=None, config=None):
+        VannaBase.__init__(self, config=config)
+
+        # default parameters - can be overridden using config
+        self.temperature = 0.2  # Lower temperature for more precise SQL generation
+        self.model = "command-a-03-2025"  # Cohere's default model
+
+        if config is not None:
+            if "temperature" in config:
+                self.temperature = config["temperature"]
+            if "model" in config:
+                self.model = config["model"]
+
+        if client is not None:
+            self.client = client
+            return
+
+        # 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.")
+            
+        # Initialize client with validated API key
+        self.client = OpenAI(
+            base_url="https://api.cohere.ai/compatibility/v1",
+            api_key=api_key,
+        )
+
+    def system_message(self, message: str) -> any:
+        return {"role": "developer", "content": message}  # Cohere uses 'developer' for system role
+
+    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
+
+        # 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"]
+
+        print(f"Using model {model} for {num_tokens} tokens (approx)")
+        try:
+            response = self.client.chat.completions.create(
+                model=model,
+                messages=prompt,
+                temperature=self.temperature,
+            )
+            
+            # Check if response has expected structure
+            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'):
+                raise ValueError("Response is missing expected 'message' field")
+                
+            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) 

vanna-0.7.7/src/vanna/cohere/cohere_embeddings.py

@@ -0,0 +1,71 @@
+import os
+
+from openai import OpenAI
+
+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"]
+
+        if client is not None:
+            self.client = client
+            return
+
+        # 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.")
+            
+        # Initialize client with validated API key
+        self.client = OpenAI(
+            base_url="https://api.cohere.ai/compatibility/v1",
+            api_key=api_key,
+        )
+
+    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:    
+            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.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)