vanna 0.7.5__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.5 → vanna-0.7.7}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: vanna
-Version: 0.7.5
+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"
@@ -52,6 +53,8 @@ Requires-Dist: boto3 ; extra == "all"
 Requires-Dist: botocore ; extra == "all"
 Requires-Dist: langchain_core ; extra == "all"
 Requires-Dist: langchain_postgres ; extra == "all"
+Requires-Dist: langchain-community ; extra == "all"
+Requires-Dist: langchain-huggingface ; extra == "all"
 Requires-Dist: xinference-client ; extra == "all"
 Requires-Dist: anthropic ; extra == "anthropic"
 Requires-Dist: azure-search-documents ; extra == "azuresearch"
@@ -79,8 +82,12 @@ Requires-Dist: httpx ; extra == "ollama"
 Requires-Dist: openai ; extra == "openai"
 Requires-Dist: opensearch-py ; extra == "opensearch"
 Requires-Dist: opensearch-dsl ; extra == "opensearch"
+Requires-Dist: langchain-community ; extra == "opensearch"
+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"
@@ -115,6 +122,7 @@ Provides-Extra: mysql
 Provides-Extra: ollama
 Provides-Extra: openai
 Provides-Extra: opensearch
+Provides-Extra: oracle
 Provides-Extra: pgvector
 Provides-Extra: pinecone
 Provides-Extra: postgres
@@ -129,9 +137,9 @@ Provides-Extra: zhipuai
 
 
 
-| GitHub | PyPI | Documentation |
-| ------ | ---- | ------------- |
-| [![GitHub](https://img.shields.io/badge/GitHub-vanna-blue?logo=github)](https://github.com/vanna-ai/vanna) | [![PyPI](https://img.shields.io/pypi/v/vanna?logo=pypi)](https://pypi.org/project/vanna/) | [![Documentation](https://img.shields.io/badge/Documentation-vanna-blue?logo=read-the-docs)](https://vanna.ai/docs/) |
+| GitHub | PyPI | Documentation | Gurubase |
+| ------ | ---- | ------------- | -------- |
+| [![GitHub](https://img.shields.io/badge/GitHub-vanna-blue?logo=github)](https://github.com/vanna-ai/vanna) | [![PyPI](https://img.shields.io/pypi/v/vanna?logo=pypi)](https://pypi.org/project/vanna/) | [![Documentation](https://img.shields.io/badge/Documentation-vanna-blue?logo=read-the-docs)](https://vanna.ai/docs/) | [![Gurubase](https://img.shields.io/badge/Gurubase-Ask%20Vanna%20Guru-006BFF)](https://gurubase.io/g/vanna) |
 
 # Vanna
 Vanna is an MIT-licensed open-source Python RAG (Retrieval-Augmented Generation) framework for SQL generation and related functionality.
@@ -164,6 +172,46 @@ These are some of the user interfaces that we've built using Vanna. You can use
 - [vanna-ai/vanna-flask](https://github.com/vanna-ai/vanna-flask)
 - [vanna-ai/vanna-slack](https://github.com/vanna-ai/vanna-slack)
 
+## Supported LLMs
+
+- [OpenAI](https://github.com/vanna-ai/vanna/tree/main/src/vanna/openai)
+- [Anthropic](https://github.com/vanna-ai/vanna/tree/main/src/vanna/anthropic)
+- [Gemini](https://github.com/vanna-ai/vanna/blob/main/src/vanna/google/gemini_chat.py)
+- [HuggingFace](https://github.com/vanna-ai/vanna/blob/main/src/vanna/hf/hf.py)
+- [AWS Bedrock](https://github.com/vanna-ai/vanna/tree/main/src/vanna/bedrock)
+- [Ollama](https://github.com/vanna-ai/vanna/tree/main/src/vanna/ollama)
+- [Qianwen](https://github.com/vanna-ai/vanna/tree/main/src/vanna/qianwen)
+- [Qianfan](https://github.com/vanna-ai/vanna/tree/main/src/vanna/qianfan)
+- [Zhipu](https://github.com/vanna-ai/vanna/tree/main/src/vanna/ZhipuAI)
+
+## Supported VectorStores
+
+- [AzureSearch](https://github.com/vanna-ai/vanna/tree/main/src/vanna/azuresearch)
+- [Opensearch](https://github.com/vanna-ai/vanna/tree/main/src/vanna/opensearch)
+- [PgVector](https://github.com/vanna-ai/vanna/tree/main/src/vanna/pgvector)
+- [PineCone](https://github.com/vanna-ai/vanna/tree/main/src/vanna/pinecone)
+- [ChromaDB](https://github.com/vanna-ai/vanna/tree/main/src/vanna/chromadb)
+- [FAISS](https://github.com/vanna-ai/vanna/tree/main/src/vanna/faiss)
+- [Marqo](https://github.com/vanna-ai/vanna/tree/main/src/vanna/marqo)
+- [Milvus](https://github.com/vanna-ai/vanna/tree/main/src/vanna/milvus)
+- [Qdrant](https://github.com/vanna-ai/vanna/tree/main/src/vanna/qdrant)
+- [Weaviate](https://github.com/vanna-ai/vanna/tree/main/src/vanna/weaviate)
+- [Oracle](https://github.com/vanna-ai/vanna/tree/main/src/vanna/oracle)
+
+## Supported Databases
+
+- [PostgreSQL](https://www.postgresql.org/)
+- [MySQL](https://www.mysql.com/)
+- [PrestoDB](https://prestodb.io/)
+- [Apache Hive](https://hive.apache.org/)
+- [ClickHouse](https://clickhouse.com/)
+- [Snowflake](https://www.snowflake.com/en/)
+- [Oracle](https://www.oracle.com/)
+- [Microsoft SQL Server](https://www.microsoft.com/en-us/sql-server/sql-server-downloads)
+- [BigQuery](https://cloud.google.com/bigquery)
+- [SQLite](https://www.sqlite.org/)
+- [DuckDB](https://duckdb.org/)
+
 
 ## Getting started
 See the [documentation](https://vanna.ai/docs/) for specifics on your desired database, LLM, etc.

{vanna-0.7.5 → vanna-0.7.7}/README.md

@@ -1,8 +1,8 @@
 
 
-| GitHub | PyPI | Documentation |
-| ------ | ---- | ------------- |
-| [![GitHub](https://img.shields.io/badge/GitHub-vanna-blue?logo=github)](https://github.com/vanna-ai/vanna) | [![PyPI](https://img.shields.io/pypi/v/vanna?logo=pypi)](https://pypi.org/project/vanna/) | [![Documentation](https://img.shields.io/badge/Documentation-vanna-blue?logo=read-the-docs)](https://vanna.ai/docs/) |
+| GitHub | PyPI | Documentation | Gurubase |
+| ------ | ---- | ------------- | -------- |
+| [![GitHub](https://img.shields.io/badge/GitHub-vanna-blue?logo=github)](https://github.com/vanna-ai/vanna) | [![PyPI](https://img.shields.io/pypi/v/vanna?logo=pypi)](https://pypi.org/project/vanna/) | [![Documentation](https://img.shields.io/badge/Documentation-vanna-blue?logo=read-the-docs)](https://vanna.ai/docs/) | [![Gurubase](https://img.shields.io/badge/Gurubase-Ask%20Vanna%20Guru-006BFF)](https://gurubase.io/g/vanna) |
 
 # Vanna
 Vanna is an MIT-licensed open-source Python RAG (Retrieval-Augmented Generation) framework for SQL generation and related functionality.
@@ -35,6 +35,46 @@ These are some of the user interfaces that we've built using Vanna. You can use
 - [vanna-ai/vanna-flask](https://github.com/vanna-ai/vanna-flask)
 - [vanna-ai/vanna-slack](https://github.com/vanna-ai/vanna-slack)
 
+## Supported LLMs
+
+- [OpenAI](https://github.com/vanna-ai/vanna/tree/main/src/vanna/openai)
+- [Anthropic](https://github.com/vanna-ai/vanna/tree/main/src/vanna/anthropic)
+- [Gemini](https://github.com/vanna-ai/vanna/blob/main/src/vanna/google/gemini_chat.py)
+- [HuggingFace](https://github.com/vanna-ai/vanna/blob/main/src/vanna/hf/hf.py)
+- [AWS Bedrock](https://github.com/vanna-ai/vanna/tree/main/src/vanna/bedrock)
+- [Ollama](https://github.com/vanna-ai/vanna/tree/main/src/vanna/ollama)
+- [Qianwen](https://github.com/vanna-ai/vanna/tree/main/src/vanna/qianwen)
+- [Qianfan](https://github.com/vanna-ai/vanna/tree/main/src/vanna/qianfan)
+- [Zhipu](https://github.com/vanna-ai/vanna/tree/main/src/vanna/ZhipuAI)
+
+## Supported VectorStores
+
+- [AzureSearch](https://github.com/vanna-ai/vanna/tree/main/src/vanna/azuresearch)
+- [Opensearch](https://github.com/vanna-ai/vanna/tree/main/src/vanna/opensearch)
+- [PgVector](https://github.com/vanna-ai/vanna/tree/main/src/vanna/pgvector)
+- [PineCone](https://github.com/vanna-ai/vanna/tree/main/src/vanna/pinecone)
+- [ChromaDB](https://github.com/vanna-ai/vanna/tree/main/src/vanna/chromadb)
+- [FAISS](https://github.com/vanna-ai/vanna/tree/main/src/vanna/faiss)
+- [Marqo](https://github.com/vanna-ai/vanna/tree/main/src/vanna/marqo)
+- [Milvus](https://github.com/vanna-ai/vanna/tree/main/src/vanna/milvus)
+- [Qdrant](https://github.com/vanna-ai/vanna/tree/main/src/vanna/qdrant)
+- [Weaviate](https://github.com/vanna-ai/vanna/tree/main/src/vanna/weaviate)
+- [Oracle](https://github.com/vanna-ai/vanna/tree/main/src/vanna/oracle)
+
+## Supported Databases
+
+- [PostgreSQL](https://www.postgresql.org/)
+- [MySQL](https://www.mysql.com/)
+- [PrestoDB](https://prestodb.io/)
+- [Apache Hive](https://hive.apache.org/)
+- [ClickHouse](https://clickhouse.com/)
+- [Snowflake](https://www.snowflake.com/en/)
+- [Oracle](https://www.oracle.com/)
+- [Microsoft SQL Server](https://www.microsoft.com/en-us/sql-server/sql-server-downloads)
+- [BigQuery](https://cloud.google.com/bigquery)
+- [SQLite](https://www.sqlite.org/)
+- [DuckDB](https://duckdb.org/)
+
 
 ## Getting started
 See the [documentation](https://vanna.ai/docs/) for specifics on your desired database, LLM, etc.

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

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "vanna"
-version = "0.7.5"
+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", "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,8 +46,8 @@ zhipuai = ["zhipuai"]
 ollama = ["ollama", "httpx"]
 qdrant = ["qdrant-client", "fastembed"]
 vllm = ["vllm"]
-pinecone = ["pinecone-client", "fastembed"]
-opensearch = ["opensearch-py", "opensearch-dsl"]
+pinecone = ["pinecone", "fastembed"]
+opensearch = ["opensearch-py", "opensearch-dsl", "langchain-community", "langchain-huggingface"]
 hf = ["transformers"]
 milvus = ["pymilvus[model]"]
 bedrock = ["boto3", "botocore"]
@@ -57,3 +57,4 @@ pgvector = ["langchain-postgres>=0.0.12"]
 faiss-cpu = ["faiss-cpu"]
 faiss-gpu = ["faiss-gpu"]
 xinference-client = ["xinference-client"]
+oracle = ["oracledb", "chromadb"]

{vanna-0.7.5 → vanna-0.7.7}/src/vanna/base/base.py

@@ -306,7 +306,7 @@ class VannaBase(ABC):
 
         message_log = [
             self.system_message(
-                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"
+                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.head(25).to_markdown()}\n\n"
             ),
             self.user_message(
                 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." +
@@ -689,6 +689,9 @@ class VannaBase(ABC):
         return response
 
     def _extract_python_code(self, markdown_string: str) -> str:
+        # Strip whitespace to avoid indentation errors in LLM-generated code
+        markdown_string = markdown_string.strip()
+
         # Regex pattern to match Python code blocks
         pattern = r"```[\w\s]*python\n([\s\S]*?)```|```([\s\S]*?)```"
 
@@ -1167,7 +1170,7 @@ class VannaBase(ABC):
         vn.connect_to_oracle(
         user="username",
         password="password",
-        dns="host:port/sid",
+        dsn="host:port/sid",
         )
         ```
         Args:

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) 

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

@@ -0,0 +1 @@
+from .deepseek_chat import DeepSeekChat

vanna-0.7.7/src/vanna/deepseek/deepseek_chat.py

@@ -0,0 +1,60 @@
+import os
+
+from openai import OpenAI
+
+from ..base import VannaBase
+
+
+
+# from vanna.chromadb import ChromaDB_VectorStore
+
+# class DeepSeekVanna(ChromaDB_VectorStore, DeepSeekChat):
+#     def __init__(self, config=None):
+#         ChromaDB_VectorStore.__init__(self, config=config)
+#         DeepSeekChat.__init__(self, config=config)
+
+# vn = DeepSeekVanna(config={"api_key": "sk-************", "model": "deepseek-chat"})
+
+
+class DeepSeekChat(VannaBase):
+    def __init__(self, config=None):
+        if config is None:
+            raise ValueError(
+                "For DeepSeek, config must be provided with an api_key and model"
+            )
+        if "api_key" not in config:
+            raise ValueError("config must contain a DeepSeek api_key")
+
+        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}
+
+    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 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:
+        chat_response = self.client.chat.completions.create(
+            model=self.model,
+            messages=prompt,
+        )
+        
+        return chat_response.choices[0].message.content

{vanna-0.7.5 → vanna-0.7.7}/src/vanna/google/gemini_chat.py

@@ -1,4 +1,5 @@
 import os
+
 from ..base import VannaBase
 
 
@@ -30,8 +31,29 @@ class GoogleGeminiChat(VannaBase):
             self.chat_model = genai.GenerativeModel(model_name)
         else:
             # Authenticate using VertexAI
+            import google.auth
+            import vertexai
             from vertexai.generative_models import GenerativeModel
-            self.chat_model = GenerativeModel(model_name)
+
+            json_file_path = config.get("google_credentials")  # Assuming the JSON file path is provided in the config
+
+            if not json_file_path or not os.path.exists(json_file_path):
+                raise FileNotFoundError(f"JSON credentials file not found at: {json_file_path}")
+
+            try:
+                # Validate and set the JSON file path for GOOGLE_APPLICATION_CREDENTIALS
+                os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = json_file_path
+
+                # Initialize VertexAI with the credentials
+                credentials, _ = google.auth.default()
+                vertexai.init(credentials=credentials)
+                self.chat_model = GenerativeModel(model_name)
+            except google.auth.exceptions.DefaultCredentialsError as e:
+                raise RuntimeError(f"Default credentials error: {e}")
+            except google.auth.exceptions.TransportError as e:
+                raise RuntimeError(f"Transport error during authentication: {e}")
+            except Exception as e:
+                raise RuntimeError(f"Failed to authenticate using JSON file: {e}")
 
     def system_message(self, message: str) -> any:
         return message

{vanna-0.7.5 → vanna-0.7.7}/src/vanna/ollama/ollama.py

@@ -91,7 +91,7 @@ class Ollama(VannaBase):
       f"model={self.model},\n"
       f"options={self.ollama_options},\n"
       f"keep_alive={self.keep_alive}")
-    self.log(f"Prompt Content:\n{json.dumps(prompt)}")
+    self.log(f"Prompt Content:\n{json.dumps(prompt, ensure_ascii=False)}")
     response_dict = self.ollama_client.chat(model=self.model,
                                             messages=prompt,
                                             stream=False,

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

@@ -0,0 +1,2 @@
+from .opensearch_vector import OpenSearch_VectorStore
+from .opensearch_vector_semantic import OpenSearch_Semantic_VectorStore