Add Cohere integration with Cohere_Chat and Cohere_Embeddings classes

[ai-yann]

• Introduced Cohere_Chat class for chat interactions with customizable parameters.
• Added Cohere_Embeddings class for generating embeddings with a default model.
• Updated imports in test files to include new classes and added tests for Cohere functionality.

[llamapreview]

Auto Pull Request Review from LlamaPReview

1. Overview

1.1 Core Changes

• Primary purpose and scope: This PR introduces Cohere integration to the Vanna project by adding Cohere_Chat and Cohere_Embeddings classes. The Cohere_Chat class handles chat interactions with customizable parameters, while the Cohere_Embeddings class generates embeddings using a default model.
• Key components modified: The PR adds new classes and updates imports in test files to include these new classes. It also adds tests for the Cohere functionality.
• Cross-component impacts: The integration of Cohere affects the overall system by expanding the supported LLM providers, increasing platform flexibility.
• Business value alignment: This integration aligns with the business goal of providing accurate Text-to-SQL Generation via LLMs using RAG, enhancing the platform's capabilities.

1.2 Technical Architecture

• System design modifications: The PR introduces new classes for handling chat interactions and generating embeddings, which are integrated into the existing architecture.
• Component interaction changes: The new classes interact with the Cohere API using an OpenAI-compatible endpoint, ensuring smooth integration with the existing system.
• Integration points impact: The integration points include the initialization of the Cohere client and the submission of prompts for chat interactions and embedding generation.
• Dependency changes and implications: The PR adds dependencies on the Cohere API and the OpenAI library, which are used for compatibility.

2. Critical Findings

2.1 Must Fix (P0🔴)

Issue: Insecure API Key Handling

• Analysis Confidence: High
• Impact: The current implementation silently fails if the API key is missing, which can lead to unauthorized access or failed API calls.
• Resolution: Add validation to ensure the API key is provided via config or environment variable.

Issue: Fragile Response Handling

• Analysis Confidence: High
• Impact: The current implementation does not handle empty or malformed API responses, which can lead to runtime errors.
• Resolution: Add checks to handle empty or malformed API responses.

2.2 Should Fix (P1🟡)

Issue: Inaccurate Token Counting

• Analysis Confidence: High
• Impact: The current token counting method is inaccurate, which can lead to incorrect token usage and potential API call failures.
• Suggested Solution: Use Cohere's actual tokenizer for accurate token counting.

Issue: Configuration Precedence

• Analysis Confidence: High
• Impact: The current model selection logic is unclear, which can lead to incorrect model usage.
• Suggested Solution: Clarify the precedence chain for model selection.

2.3 Consider (P2🟢)

Area: Shared Client Configuration

• Analysis Confidence: High
• Improvement Opportunity: Reducing code duplication between chat and embeddings classes by creating a shared client configuration.

Area: Rate Limiting Handling

• Analysis Confidence: High
• Improvement Opportunity: Adding resilience to API rate limits by implementing retry logic.

2.4 Summary of Action Items

• P0🔴: Address insecure API key handling and fragile response handling before merging.
• P1🟡: Implement accurate token counting and clarify configuration precedence.
• P2🟢: Consider shared client configuration and rate limiting handling for future improvements.

3. Technical Analysis

3.1 Code Logic Analysis

📁 src/vanna/cohere/cohere_chat.py - Cohere_Chat

• Submitted PR Code:

    class Cohere_Chat(VannaBase):
        def __init__(self, client=None, config=None):
            VannaBase.__init__(self, config=config)
            self.temperature = 0.2
            self.model = "command-a-03-2025"
            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
            if config is None and client is None:
                self.client = OpenAI(
                    base_url="https://api.cohere.ai/compatibility/v1",
                    api_key=os.getenv("COHERE_API_KEY"),
                )
                return
            if "api_key" in config:
                self.client = OpenAI(
                    base_url="https://api.cohere.ai/compatibility/v1",
                    api_key=config["api_key"],
                )
        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")
            num_tokens = 0
            for message in prompt:
                num_tokens += len(message["content"]) / 4
            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)")
            response = self.client.chat.completions.create(
                model=model,
                messages=prompt,
                temperature=self.temperature,
            )
            return response.choices[0].message.content

• Analysis:
  • Current logic and potential issues: The initialization of the Cohere client lacks validation for the API key, which can lead to silent failures. The token counting method is inaccurate, and the model selection logic is unclear.
  • Edge cases and error handling: The code does not handle empty or malformed API responses, which can lead to runtime errors.
  • Cross-component impact : The integration with the Cohere API affects the overall system by expanding the supported LLM providers.
  • Business logic considerations : The business logic for generating SQL queries using LLMs is enhanced by the Cohere integration.
• LlamaPReview Suggested Improvements:

    class Cohere_Chat(VannaBase):
        def __init__(self, client=None, config=None):
            VannaBase.__init__(self, config=config)
            self.temperature = 0.2
            self.model = "command-a-03-2025"
            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
            api_key = config.get("api_key") or os.environ.get("COHERE_API_KEY") or ""
            if not api_key:
                raise ValueError("Cohere API key must be provided via config or environment variable")
            self.client = OpenAI(
                base_url="https://api.cohere.ai/compatibility/v1",
                api_key=api_key,
            )
        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")
            from cohere import Client as CohereClient
            co = CohereClient(api_key=self.config.get("api_key"))
            num_tokens = 0
            for message in prompt:
                tokenized = co.tokenize(text=message["content"])
                num_tokens += len(tokenized.tokens)
            model = (
                kwargs.get("model") or
                (self.config.get("model") if self.config else None) or
                self.model
            )
            print(f"Using model {model} for {num_tokens} tokens (approx)")
            response = self.client.chat.completions.create(
                model=model,
                messages=prompt,
                temperature=self.temperature,
            )
            if not response.choices:
                raise ValueError("Empty response from Cohere API")
            return response.choices[0].message.content

• Improvement rationale:
  • Technical benefits: Adding API key validation prevents silent failures. Using Cohere's actual tokenizer ensures accurate token counting. Clarifying the model selection logic ensures the correct model is used. Handling empty or malformed API responses prevents runtime errors.
  • Business value: Enhances the reliability and accuracy of the Cohere integration, aligning with the business goal of providing accurate Text-to-SQL Generation via LLMs.
  • Risk assessment: Addressing these issues reduces the risk of unauthorized access, API call failures, and runtime errors.

📁 src/vanna/cohere/cohere_embeddings.py - Cohere_Embeddings

• Submitted PR Code:

    class Cohere_Embeddings(VannaBase):
        def __init__(self, client=None, config=None):
            VannaBase.__init__(self, config=config)
            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
            if config is None and client is None:
                self.client = OpenAI(
                    base_url="https://api.cohere.ai/compatibility/v1",
                    api_key=os.getenv("COHERE_API_KEY"),
                )
                return
            if "api_key" in config:
                self.client = OpenAI(
                    base_url="https://api.cohere.ai/compatibility/v1",
                    api_key=config["api_key"],
                )
        def generate_embedding(self, data: str, **kwargs) -> list[float]:
            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"]
            embedding = self.client.embeddings.create(
                model=model,
                input=data,
                encoding_format="float",
            )
            return embedding.data[0].embedding

• Analysis:
  • Current logic and potential issues: The initialization of the Cohere client lacks validation for the API key, which can lead to silent failures. The embedding generation does not handle empty input, which can lead to runtime errors.
  • Edge cases and error handling: The code does not handle empty input, which can lead to runtime errors.
  • Cross-component impact : The integration with the Cohere API affects the overall system by expanding the supported LLM providers.
  • Business logic considerations : The business logic for generating embeddings is enhanced by the Cohere integration.
• LlamaPReview Suggested Improvements:

    class Cohere_Embeddings(VannaBase):
        def __init__(self, client=None, config=None):
            VannaBase.__init__(self, config=config)
            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
            api_key = config.get("api_key") or os.environ.get("COHERE_API_KEY") or ""
            if not api_key:
                raise ValueError("Cohere API key must be provided via config or environment variable")
            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.strip():
                raise ValueError("Input text cannot be empty")
            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"]
            embedding = self.client.embeddings.create(
                model=model,
                input=data,
                encoding_format="float",
            )
            return embedding.data[0].embedding

• Improvement rationale:
  • Technical benefits: Adding API key validation prevents silent failures. Handling empty input prevents runtime errors.
  • Business value: Enhances the reliability of the Cohere integration, aligning with the business goal of providing accurate Text-to-SQL Generation via LLMs.
  • Risk assessment: Addressing these issues reduces the risk of unauthorized access and runtime errors.

3.2 Key Quality Aspects

• System scalability considerations: The PR introduces new classes for handling chat interactions and generating embeddings, which can be scaled as needed.
• Performance bottlenecks and optimizations: The PR does not introduce any significant performance bottlenecks. However, implementing batch embedding support and async API calls can further optimize performance.
• Testing strategy and coverage: The PR adds tests for the Cohere functionality, ensuring adequate test coverage. Adding negative tests for invalid API keys and mock API responses can further enhance the testing strategy.
• Documentation needs: The PR should include usage examples in class docstrings, document configuration options, and note the use of the OpenAI compatibility layer.

4. Overall Evaluation

• Technical assessment: The PR introduces Cohere integration with well-structured classes and adequate test coverage. However, it requires improvements in API key handling, response validation, and token counting.
• Business impact: The Cohere integration enhances the platform's capabilities by expanding the supported LLM providers, aligning with the business goal of providing accurate Text-to-SQL Generation via LLMs.
• Risk evaluation: The PR introduces risks related to unauthorized access, API call failures, and runtime errors, which can be mitigated by addressing the identified issues.
• Notable positive aspects and good practices: The PR follows existing patterns and integrates well with the current system. It also includes adequate test coverage for the new functionality.
• Implementation quality: The implementation quality is generally good, with well-structured classes and adequate test coverage. However, it requires improvements in error handling and documentation.
• Final recommendation (Approve/Request Changes/Needs Discussion): Request Changes. The PR requires addressing the identified issues before merging.

---

💡 LlamaPReview Community
Have feedback on this AI Code review tool? Join our GitHub Discussions to share your thoughts and help shape the future of LlamaPReview.

[ai-yann]

I've made the requested changes and added new commits to this pull request. However, it seems that the repository is not configured to re-trigger LlamaPReview when new commits are pushed. Could you please modify this setting, or should I just close the PR and re-open a new one @zainhoda?