# How to Add Support for Additional LLM Providers in AI Hedge Fund

> Easily add LLM provider support to AI Hedge Fund. Extend the ModelProvider enum and implement a new factory function branch. Learn how to integrate new LLM clients.

- Repository: [Virat Singh/ai-hedge-fund](https://github.com/virattt/ai-hedge-fund)
- Tags: how-to-guide
- Published: 2026-03-09

---

**To add support for additional LLM providers in the AI Hedge Fund framework, extend the `ModelProvider` enum in [`src/llm/models.py`](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/models.py), import the appropriate LangChain chat wrapper, and implement a new conditional branch in the `get_model` factory function that handles API key validation and client instantiation.**

The virattt/ai-hedge-fund repository abstracts LLM integrations behind a unified factory pattern, allowing you to **add support for additional LLM providers** without modifying the high-level inference logic. By updating the provider enum and model factory in [`src/llm/models.py`](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/models.py), you can integrate any LangChain-compatible chat model while maintaining the existing `call_llm` interface used throughout the application.

## Step-by-Step Implementation Guide

This guide walks through integrating a new provider (using Cohere as an example) into the codebase.

### Step 1: Extend the ModelProvider Enum

Open [`src/llm/models.py`](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/models.py) and add a new member to the `ModelProvider` enumeration (around line 18). This string-based enum identifies the provider throughout the application.

```python

# src/llm/models.py

class ModelProvider(str, Enum):
    """Enum for supported LLM providers"""
    OPENAI = "OpenAI"
    ANTHROPIC = "Anthropic"
    GROQ = "Groq"
    XAI = "xAI"
    # Add your new provider below

    COHERE = "Cohere"

```

*Reference:* [ModelProvider enum definition](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/models.py#L18)

### Step 2: Install and Import the LangChain Wrapper

Install the provider's LangChain integration package. For example, to add Cohere:

```bash
pip install langchain-cohere

```

Then import the chat wrapper at the top of [`src/llm/models.py`](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/models.py) (around lines 3-12):

```python

# src/llm/models.py

from langchain_cohere import ChatCohere  # New import

```

*Reference:* [Existing imports in models.py](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/models.py#L3)

### Step 3: Add the Provider Case to get_model

Locate the `get_model` factory function (starting around line 138) and add a new `elif` block that instantiates the provider's client. The function receives `model_name`, `model_provider`, and an optional `api_keys` dictionary.

```python

# src/llm/models.py

def get_model(model_name: str, model_provider: ModelProvider, api_keys: dict = None):
    ...
    elif model_provider == ModelProvider.COHERE:
        api_key = (api_keys or {}).get("COHERE_API_KEY") or os.getenv("COHERE_API_KEY")
        if not api_key:
            print("API Key Error: Please set COHERE_API_KEY in your .env or provide via API keys.")
            raise ValueError("Cohere API key not found.")
        return ChatCohere(model=model_name, api_key=api_key)
    ...

```

*Reference:* [get_model factory implementation](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/models.py#L138)

### Step 4: Register Models in JSON (Optional)

To make the new provider's models selectable in the UI, add entries to [`src/llm/api_models.json`](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/api_models.json) (or [`ollama_models.json`](https://github.com/virattt/ai-hedge-fund/blob/main/ollama_models.json) for local models). Each entry requires the provider name exactly as defined in the enum.

```json
{
  "display_name": "Cohere Command R+",
  "model_name": "command-r-plus",
  "provider": "Cohere"
}

```

*Reference:* [api_models.json structure](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/api_models.json#L2)

### Step 5: Verify UI Integration

The application builds its selection dropdowns automatically from the `LLM_ORDER` constant (defined around line 108 in [`src/llm/models.py`](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/models.py)), which derives from `AVAILABLE_MODELS` populated via the JSON files. No additional UI code changes are required—restart the application to see the new provider in the dropdown menus.

```python

# src/llm/models.py

LLM_ORDER = [model.to_choice_tuple() for model in AVAILABLE_MODELS]

```

*Reference:* [LLM_ORDER construction](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/models.py#L108)

### Step 6: Test the Integration

The high-level `call_llm` function in [`src/utils/llm.py`](https://github.com/virattt/ai-hedge-fund/blob/main/src/utils/llm.py) automatically routes to your new provider via `get_model`. Verify end-to-end functionality by invoking the model:

```python

# src/utils/llm.py usage pattern

model_info = get_model_info(model_name, model_provider)
llm = get_model(model_name, model_provider, api_keys)  # Now supports your new provider

```

*Reference:* [call_llm implementation](https://github.com/virattt/ai-hedge-fund/blob/main/src/utils/llm.py#L48)

## Complete Implementation Example

Below is a unified diff showing the complete changes required to add Cohere support:

```diff
--- a/src/llm/models.py
+++ b/src/llm/models.py
@@ -1,6 +1,7 @@
 from langchain_openai import ChatOpenAI, AzureChatOpenAI
 from langchain_gigachat import GigaChat
 from langchain_ollama import ChatOllama
+from langchain_cohere import ChatCohere
 
 class ModelProvider(str, Enum):
     """Enum for supported LLM providers"""
@@ -9,6 +10,7 @@
     GROQ = "Groq"
     XAI = "xAI"
     GIGACHAT = "GigaChat"
+    COHERE = "Cohere"
 
 def get_model(model_name: str, model_provider: ModelProvider, api_keys: dict = None):
     ...
@@ -20,6 +22,12 @@
         if not api_key:
             raise ValueError("xAI API key not found.")
         return ChatXAI(model=model_name, api_key=api_key)
+    elif model_provider == ModelProvider.COHERE:
+        api_key = (api_keys or {}).get("COHERE_API_KEY") or os.getenv("COHERE_API_KEY")
+        if not api_key:
+            print("API Key Error: Please set COHERE_API_KEY in your .env or provide via API keys.")
+            raise ValueError("Cohere API key not found.")
+        return ChatCohere(model=model_name, api_key=api_key)
     ...

```

Add the corresponding entry to [`src/llm/api_models.json`](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/api_models.json):

```json
{
  "display_name": "Cohere Command R+",
  "model_name": "command-r-plus",
  "provider": "Cohere"
}

```

## Summary

- **Extend the enum:** Add the provider identifier to `ModelProvider` in [`src/llm/models.py`](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/models.py).

- **Import the wrapper:** Install and import the LangChain chat class for your provider.

- **Implement the factory:** Add an `elif` branch in `get_model` that handles API key retrieval and client instantiation.

- **Register models:** Add model metadata to [`src/llm/api_models.json`](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/api_models.json) to populate the UI dropdowns.

- **Verify integration:** The `call_llm` function in [`src/utils/llm.py`](https://github.com/virattt/ai-hedge-fund/blob/main/src/utils/llm.py) automatically utilizes the new provider without additional changes.

## Frequently Asked Questions

### Do I need to modify the UI code to add a new LLM provider?

No. The UI generates provider and model dropdowns automatically from the `LLM_ORDER` constant in [`src/llm/models.py`](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/models.py), which is constructed from `AVAILABLE_MODELS`. As long as you register your models in [`src/llm/api_models.json`](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/api_models.json) and extend the `ModelProvider` enum, the new options appear after restarting the application.

### What if my LLM provider doesn't have a LangChain integration?

You must implement a custom LangChain chat model class that conforms to the `BaseChatModel` interface. The `get_model` factory expects objects implementing LangChain's standard methods like `invoke()` and `stream()`. Without this compatibility layer, you cannot integrate with the existing `call_llm` abstraction.

### How do I handle API key validation for new providers?

Follow the established pattern in [`src/llm/models.py`](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/models.py): retrieve the key from the optional `api_keys` dictionary parameter first, then fall back to environment variables using `os.getenv()`. Always validate the key exists before instantiating the client, and raise a descriptive `ValueError` if missing to match the framework's error handling conventions.

### Where should I store model metadata for local Ollama models?

Register local Ollama models in [`src/llm/ollama_models.json`](https://github.com/virattt/ai-hedge-fund/blob/main/src/llm/ollama_models.json) rather than [`api_models.json`](https://github.com/virattt/ai-hedge-fund/blob/main/api_models.json). The framework maintains separate JSON files for API-based providers and local Ollama instances, though both use the same `ModelProvider.OLLAMA` enum value and `ChatOllama` wrapper in the factory function.