Merge pull request #830 from ai-yann/main

Add Cohere integration with Cohere_Chat and Cohere_Embeddings classes

Author: zainhoda

src/vanna/cohere/__init__.py

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

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) 

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) 

tests/test_imports.py

@@ -1,11 +1,11 @@
-
-
 def test_regular_imports():
     from vanna.anthropic.anthropic_chat import Anthropic_Chat
     from vanna.azuresearch.azuresearch_vector import AzureAISearch_VectorStore
     from vanna.base.base import VannaBase
     from vanna.bedrock.bedrock_converse import Bedrock_Converse
     from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore
+    from vanna.cohere.cohere_chat import Cohere_Chat
+    from vanna.cohere.cohere_embeddings import Cohere_Embeddings
     from vanna.faiss.faiss import FAISS
     from vanna.google.bigquery_vector import BigQuery_VectorStore
     from vanna.google.gemini_chat import GoogleGeminiChat
@@ -40,6 +40,7 @@ def test_shortcut_imports():
     from vanna.azuresearch import AzureAISearch_VectorStore
     from vanna.base import VannaBase
     from vanna.chromadb import ChromaDB_VectorStore
+    from vanna.cohere import Cohere_Chat, Cohere_Embeddings
     from vanna.faiss import FAISS
     from vanna.hf import Hf
     from vanna.marqo import Marqo_VectorStore

tests/test_vanna.py

@@ -1,6 +1,7 @@
 import os
 
 from vanna.anthropic.anthropic_chat import Anthropic_Chat
+from vanna.cohere.cohere_chat import Cohere_Chat
 from vanna.google import GoogleGeminiChat
 from vanna.mistral.mistral import Mistral
 from vanna.openai.openai_chat import OpenAI_Chat
@@ -227,6 +228,23 @@ def test_vn_gemini():
     df = vn_gemini.run_sql(sql)
     assert len(df) == 9
 
+class VannaCohere(VannaDB_VectorStore, Cohere_Chat):
+    def __init__(self, config=None):
+        VannaDB_VectorStore.__init__(self, vanna_model=MY_VANNA_MODEL, vanna_api_key=MY_VANNA_API_KEY, config=config)
+        Cohere_Chat.__init__(self, config=config)
+
+try:
+    COHERE_API_KEY = os.environ['COHERE_API_KEY']
+    vn_cohere = VannaCohere(config={'api_key': COHERE_API_KEY, 'model': 'command-a-03-2025'})
+    vn_cohere.connect_to_sqlite('https://vanna.ai/Chinook.sqlite')
+    
+    def test_vn_cohere():
+        sql = vn_cohere.generate_sql("What are the top 10 customers by sales?")
+        df = vn_cohere.run_sql(sql)
+        assert len(df) == 10
+except KeyError:
+    print("Skipping Cohere tests - COHERE_API_KEY not found in environment variables")
+
 def test_training_plan():
     vn_dummy = VannaDefault(model=MY_VANNA_MODEL, api_key=MY_VANNA_API_KEY)