charieswei
/
TradingAgents
mirror of https://github.com/TauricResearch/TradingAgents.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat: support custom provider and model

This commit is contained in:
mogita 2025-08-16 00:48:23 +08:00
parent 95572ece42
commit b9ad5adc78
No known key found for this signature in database
GPG Key ID: A0AA1B9C57A48ECF
7 changed files with 341 additions and 82 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

310
cli/utils.py
View File

@ -122,38 +122,165 @@ def select_research_depth() -> int:
return choice
# Centralized model definitions - single source of truth
SHALLOW_AGENT_OPTIONS = {
"openai": [
("GPT-4o-mini - Fast and efficient for quick tasks", "gpt-4o-mini"),
("GPT-4.1-nano - Ultra-lightweight model for basic operations", "gpt-4.1-nano"),
("GPT-4.1-mini - Compact model with good performance", "gpt-4.1-mini"),
("GPT-4o - Standard model with solid capabilities", "gpt-4o"),
],
"anthropic": [
("Claude Haiku 3.5 - Fast inference and standard capabilities", "claude-3-5-haiku-latest"),
("Claude Sonnet 3.5 - Highly capable standard model", "claude-3-5-sonnet-latest"),
("Claude Sonnet 3.7 - Exceptional hybrid reasoning and agentic capabilities", "claude-3-7-sonnet-latest"),
("Claude Sonnet 4 - High performance and excellent reasoning", "claude-sonnet-4-0"),
],
"google": [
("Gemini 2.0 Flash-Lite - Cost efficiency and low latency", "gemini-2.0-flash-lite"),
("Gemini 2.0 Flash - Next generation features, speed, and thinking", "gemini-2.0-flash"),
("Gemini 2.5 Flash - Adaptive thinking, cost efficiency", "gemini-2.5-flash-preview-05-20"),
],
"openrouter": [
("Meta: Llama 4 Scout", "meta-llama/llama-4-scout:free"),
("Meta: Llama 3.3 8B Instruct - A lightweight and ultra-fast variant of Llama 3.3 70B", "meta-llama/llama-3.3-8b-instruct:free"),
("google/gemini-2.0-flash-exp:free - Gemini Flash 2.0 offers a significantly faster time to first token", "google/gemini-2.0-flash-exp:free"),
],
"ollama": [
("llama3.1 local", "llama3.1"),
("llama3.2 local", "llama3.2"),
]
}
DEEP_AGENT_OPTIONS = {
"openai": [
("GPT-4.1-nano - Ultra-lightweight model for basic operations", "gpt-4.1-nano"),
("GPT-4.1-mini - Compact model with good performance", "gpt-4.1-mini"),
("GPT-4o - Standard model with solid capabilities", "gpt-4o"),
("o4-mini - Specialized reasoning model (compact)", "o4-mini"),
("o3-mini - Advanced reasoning model (lightweight)", "o3-mini"),
("o3 - Full advanced reasoning model", "o3"),
("o1 - Premier reasoning and problem-solving model", "o1"),
],
"anthropic": [
("Claude Haiku 3.5 - Fast inference and standard capabilities", "claude-3-5-haiku-latest"),
("Claude Sonnet 3.5 - Highly capable standard model", "claude-3-5-sonnet-latest"),
("Claude Sonnet 3.7 - Exceptional hybrid reasoning and agentic capabilities", "claude-3-7-sonnet-latest"),
("Claude Sonnet 4 - High performance and excellent reasoning", "claude-sonnet-4-0"),
("Claude Opus 4 - Most powerful Anthropic model", "claude-opus-4-0"),
],
"google": [
("Gemini 2.0 Flash-Lite - Cost efficiency and low latency", "gemini-2.0-flash-lite"),
("Gemini 2.0 Flash - Next generation features, speed, and thinking", "gemini-2.0-flash"),
("Gemini 2.5 Flash - Adaptive thinking, cost efficiency", "gemini-2.5-flash-preview-05-20"),
("Gemini 2.5 Pro", "gemini-2.5-pro-preview-06-05"),
],
"openrouter": [
("DeepSeek V3 - a 685B-parameter, mixture-of-experts model", "deepseek/deepseek-chat-v3-0324:free"),
("Deepseek - latest iteration of the flagship chat model family from the DeepSeek team.", "deepseek/deepseek-chat-v3-0324:free"),
],
"ollama": [
("llama3.1 local", "llama3.1"),
("qwen3", "qwen3"),
]
}
def _get_all_models_for_custom_provider(model_type: str) -> list:
"""Get unified model list for custom provider with all available models from all providers.
Args:
model_type: Either 'shallow' or 'deep' to get the appropriate model set
Returns:
List of (description, model_value) tuples
"""
# Use the centralized model definitions
if model_type == "shallow":
provider_models = SHALLOW_AGENT_OPTIONS
else: # deep
provider_models = DEEP_AGENT_OPTIONS
# Combine all models with provider labels
all_models = []
for provider_name, models in provider_models.items():
provider_display_name = provider_name.title()
for description, model_value in models:
labeled_description = f"{description} ({provider_display_name})"
all_models.append((labeled_description, model_value))
# Add custom model option at the end
all_models.append(("Custom Model - Enter your own model name", "__CUSTOM_MODEL__"))
return all_models
def _select_custom_provider_model(model_type: str, title: str, default_model: str) -> str:
"""Handle model selection for custom provider with unified model list.
Args:
model_type: Either 'shallow' or 'deep'
title: Title for the selection prompt
default_model: Default model name for custom input
Returns:
Selected model name
"""
all_models = _get_all_models_for_custom_provider(model_type)
choice = questionary.select(
title,
choices=[
questionary.Choice(display, value=value)
for display, value in all_models
],
instruction="\n- Use arrow keys to navigate\n- Press Enter to select\n- Your custom endpoint should support the selected model",
style=questionary.Style(
[
("selected", "fg:magenta noinherit"),
("highlighted", "fg:magenta noinherit"),
("pointer", "fg:magenta noinherit"),
]
),
).ask()
if choice is None:
from rich.console import Console
console = Console()
console.print(f"\n[red]No {model_type} thinking model selected. Exiting...[/red]")
exit(1)
# Handle custom model input
if choice == "__CUSTOM_MODEL__":
custom_model = questionary.text(
f"Enter your custom {model_type} thinking model name:",
default=default_model,
instruction="\n- Enter the exact model name as supported by your custom endpoint\n- Press Enter to confirm"
).ask()
if not custom_model:
from rich.console import Console
console = Console()
console.print(f"\n[red]No custom {model_type} model name entered. Exiting...[/red]")
exit(1)
return custom_model
return choice
def select_shallow_thinking_agent(provider) -> str:
"""Select shallow thinking llm engine using an interactive selection."""
# Define shallow thinking llm engine options with their corresponding model names
SHALLOW_AGENT_OPTIONS = {
"openai": [
("GPT-4o-mini - Fast and efficient for quick tasks", "gpt-4o-mini"),
("GPT-4.1-nano - Ultra-lightweight model for basic operations", "gpt-4.1-nano"),
("GPT-4.1-mini - Compact model with good performance", "gpt-4.1-mini"),
("GPT-4o - Standard model with solid capabilities", "gpt-4o"),
],
"anthropic": [
("Claude Haiku 3.5 - Fast inference and standard capabilities", "claude-3-5-haiku-latest"),
("Claude Sonnet 3.5 - Highly capable standard model", "claude-3-5-sonnet-latest"),
("Claude Sonnet 3.7 - Exceptional hybrid reasoning and agentic capabilities", "claude-3-7-sonnet-latest"),
("Claude Sonnet 4 - High performance and excellent reasoning", "claude-sonnet-4-0"),
],
"google": [
("Gemini 2.0 Flash-Lite - Cost efficiency and low latency", "gemini-2.0-flash-lite"),
("Gemini 2.0 Flash - Next generation features, speed, and thinking", "gemini-2.0-flash"),
("Gemini 2.5 Flash - Adaptive thinking, cost efficiency", "gemini-2.5-flash-preview-05-20"),
],
"openrouter": [
("Meta: Llama 4 Scout", "meta-llama/llama-4-scout:free"),
("Meta: Llama 3.3 8B Instruct - A lightweight and ultra-fast variant of Llama 3.3 70B", "meta-llama/llama-3.3-8b-instruct:free"),
("google/gemini-2.0-flash-exp:free - Gemini Flash 2.0 offers a significantly faster time to first token", "google/gemini-2.0-flash-exp:free"),
],
"ollama": [
("llama3.1 local", "llama3.1"),
("llama3.2 local", "llama3.2"),
]
}
# Handle custom provider - use unified model selection
if provider.lower().startswith("custom"):
return _select_custom_provider_model(
model_type="shallow",
title="Select Your [Quick-Thinking LLM Engine] (Custom Provider - All Models Available):",
default_model="gpt-4o-mini"
)
# Use centralized shallow thinking model definitions
choice = questionary.select(
"Select Your [Quick-Thinking LLM Engine]:",
@ -172,6 +299,8 @@ def select_shallow_thinking_agent(provider) -> str:
).ask()
if choice is None:
from rich.console import Console
console = Console()
console.print(
"\n[red]No shallow thinking llm engine selected. Exiting...[/red]"
)
@ -183,40 +312,16 @@ def select_shallow_thinking_agent(provider) -> str:
def select_deep_thinking_agent(provider) -> str:
"""Select deep thinking llm engine using an interactive selection."""
# Define deep thinking llm engine options with their corresponding model names
DEEP_AGENT_OPTIONS = {
"openai": [
("GPT-4.1-nano - Ultra-lightweight model for basic operations", "gpt-4.1-nano"),
("GPT-4.1-mini - Compact model with good performance", "gpt-4.1-mini"),
("GPT-4o - Standard model with solid capabilities", "gpt-4o"),
("o4-mini - Specialized reasoning model (compact)", "o4-mini"),
("o3-mini - Advanced reasoning model (lightweight)", "o3-mini"),
("o3 - Full advanced reasoning model", "o3"),
("o1 - Premier reasoning and problem-solving model", "o1"),
],
"anthropic": [
("Claude Haiku 3.5 - Fast inference and standard capabilities", "claude-3-5-haiku-latest"),
("Claude Sonnet 3.5 - Highly capable standard model", "claude-3-5-sonnet-latest"),
("Claude Sonnet 3.7 - Exceptional hybrid reasoning and agentic capabilities", "claude-3-7-sonnet-latest"),
("Claude Sonnet 4 - High performance and excellent reasoning", "claude-sonnet-4-0"),
("Claude Opus 4 - Most powerful Anthropic model", " claude-opus-4-0"),
],
"google": [
("Gemini 2.0 Flash-Lite - Cost efficiency and low latency", "gemini-2.0-flash-lite"),
("Gemini 2.0 Flash - Next generation features, speed, and thinking", "gemini-2.0-flash"),
("Gemini 2.5 Flash - Adaptive thinking, cost efficiency", "gemini-2.5-flash-preview-05-20"),
("Gemini 2.5 Pro", "gemini-2.5-pro-preview-06-05"),
],
"openrouter": [
("DeepSeek V3 - a 685B-parameter, mixture-of-experts model", "deepseek/deepseek-chat-v3-0324:free"),
("Deepseek - latest iteration of the flagship chat model family from the DeepSeek team.", "deepseek/deepseek-chat-v3-0324:free"),
],
"ollama": [
("llama3.1 local", "llama3.1"),
("qwen3", "qwen3"),
]
}
# Handle custom provider - use unified model selection
if provider.lower().startswith("custom"):
return _select_custom_provider_model(
model_type="deep",
title="Select Your [Deep-Thinking LLM Engine] (Custom Provider - All Models Available):",
default_model="o4-mini"
)
# Use centralized deep thinking model definitions
choice = questionary.select(
"Select Your [Deep-Thinking LLM Engine]:",
choices=[
@ -234,24 +339,83 @@ def select_deep_thinking_agent(provider) -> str:
).ask()
if choice is None:
from rich.console import Console
console = Console()
console.print("\n[red]No deep thinking llm engine selected. Exiting...[/red]")
exit(1)
return choice
def select_llm_provider() -> tuple[str, str]:
"""Select the OpenAI api url using interactive selection."""
def validate_custom_url(url: str) -> str:
"""Validate that a custom URL is properly formatted and has a valid hostname."""
import re
from urllib.parse import urlparse
from rich.console import Console
if not url:
return ""
console = Console()
# Basic URL format validation
url_pattern = re.compile(
r'^https?://' # http:// or https://
r'(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+[A-Z]{2,6}\.?|' # domain...
r'localhost|' # localhost...
r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})' # ...or ip
r'(?::\d+)?' # optional port
r'(?:/?|[/?]\S+)$', re.IGNORECASE)
if not url_pattern.match(url):
console.print(f"[red]Error: Invalid CUSTOM_BASE_URL format: {url}[/red]")
console.print(f"[red]Please provide a valid URL (e.g., https://api.example.com/v1)[/red]")
exit(1)
# Additional validation using urlparse
try:
parsed = urlparse(url)
if not parsed.netloc:
raise ValueError("No hostname found")
return url
except Exception as e:
console.print(f"[red]Error: Invalid CUSTOM_BASE_URL: {url}[/red]")
console.print(f"[red]URL parsing error: {e}[/red]")
exit(1)
def get_custom_provider_info() -> tuple[str, str] | None:
"""Get custom provider info if both URL and API key are provided."""
import os
# Define OpenAI api options with their corresponding endpoints
# Use custom URL from environment if available, otherwise use default
openai_url = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
from urllib.parse import urlparse
custom_url = os.getenv("CUSTOM_BASE_URL")
custom_api_key = os.getenv("CUSTOM_API_KEY")
if custom_url and custom_api_key:
validated_url = validate_custom_url(custom_url)
parsed = urlparse(validated_url)
hostname = parsed.netloc
return f"Custom ({hostname})", validated_url
return None
def select_llm_provider() -> tuple[str, str]:
"""Select the LLM provider with support for a custom OpenAI-compatible endpoint."""
# Define default providers
BASE_URLS = [
("OpenAI", openai_url),
("OpenAI", "https://api.openai.com/v1"),
("Anthropic", "https://api.anthropic.com/"),
("Google", "https://generativelanguage.googleapis.com/v1"),
("Openrouter", "https://openrouter.ai/api/v1"),
("Ollama", "http://localhost:11434/v1"),
]
# Add custom provider at the beginning if available
custom_info = get_custom_provider_info()
if custom_info:
BASE_URLS.insert(0, custom_info)
choice = questionary.select(
"Select your LLM Provider:",
@ -270,10 +434,12 @@ def select_llm_provider() -> tuple[str, str]:
).ask()
if choice is None:
console.print("\n[red]no OpenAI backend selected. Exiting...[/red]")
from rich.console import Console
console = Console()
console.print("\n[red]No LLM provider selected. Exiting...[/red]")
exit(1)
display_name, url = choice
print(f"You selected: {display_name}\tURL: {url}")
return display_name, url

20
example.env Normal file
View File

@ -0,0 +1,20 @@
# Copy this to your .env file and modify the URLs and API keys as needed
# Custom OpenAI-Compatible Provider (optional)
# If provided, a "Custom" option will appear first in the provider list
# The custom endpoint must be OpenAI-compatible (REST API, not gRPC)
# CUSTOM_BASE_URL=https://www.example.com/v1
# CUSTOM_API_KEY=sk-your-custom-api-key-here
# Standard Provider API Keys, please replace with your own keys to use the corresponding provider
OPENAI_API_KEY=sk-your-openai-api-key-here
ANTHROPIC_API_KEY=sk-ant-your-anthropic-api-key-here
GOOGLE_API_KEY=your-google-api-key-here
OPENROUTER_API_KEY=sk-or-your-openrouter-api-key-here
# OLLAMA_API_KEY is usually not needed for local Ollama instances
# Other Configuration
FINNHUB_API_KEY=your-finnhub-api-key-here
# Optional, uncomment to modify
# TRADINGAGENTS_RESULTS_DIR=./results

11
tradingagents/agents/utils/memory.py
View File

@ -1,6 +1,7 @@
import chromadb
from chromadb.config import Settings
from openai import OpenAI
import os
class FinancialSituationMemory:
@ -9,7 +10,15 @@ class FinancialSituationMemory:
self.embedding = "nomic-embed-text"
else:
self.embedding = "text-embedding-3-small"
self.client = OpenAI(base_url=config["backend_url"])
# Use CUSTOM_API_KEY if provider is custom, otherwise use OPENAI_API_KEY
provider = config.get("llm_provider", "openai").lower()
if provider.startswith("custom"):
api_key = os.getenv("CUSTOM_API_KEY")
else:
api_key = os.getenv("OPENAI_API_KEY")
self.client = OpenAI(base_url=config["backend_url"], api_key=api_key)
self.chroma_client = chromadb.Client(Settings(allow_reset=True))
self.situation_collection = self.chroma_client.create_collection(name=name)

27
tradingagents/dataflows/interface.py
View File

@ -704,7 +704,14 @@ def get_YFin_data(
def get_stock_news_openai(ticker, curr_date):
config = get_config()
client = OpenAI(base_url=config["backend_url"], api_key=os.getenv("OPENAI_API_KEY"))
# Use CUSTOM_API_KEY if provider is custom, otherwise use OPENAI_API_KEY
provider = config.get("llm_provider", "openai").lower()
if provider.startswith("custom"):
api_key = os.getenv("CUSTOM_API_KEY")
else:
api_key = os.getenv("OPENAI_API_KEY")
client = OpenAI(base_url=config["backend_url"], api_key=api_key)
response = client.responses.create(
model=config["quick_think_llm"],
@ -739,7 +746,14 @@ def get_stock_news_openai(ticker, curr_date):
def get_global_news_openai(curr_date):
config = get_config()
client = OpenAI(base_url=config["backend_url"], api_key=os.getenv("OPENAI_API_KEY"))
# Use CUSTOM_API_KEY if provider is custom, otherwise use OPENAI_API_KEY
provider = config.get("llm_provider", "openai").lower()
if provider.startswith("custom"):
api_key = os.getenv("CUSTOM_API_KEY")
else:
api_key = os.getenv("OPENAI_API_KEY")
client = OpenAI(base_url=config["backend_url"], api_key=api_key)
response = client.responses.create(
model=config["quick_think_llm"],
@ -774,7 +788,14 @@ def get_global_news_openai(curr_date):
def get_fundamentals_openai(ticker, curr_date):
config = get_config()
client = OpenAI(base_url=config["backend_url"], api_key=os.getenv("OPENAI_API_KEY"))
# Use CUSTOM_API_KEY if provider is custom, otherwise use OPENAI_API_KEY
provider = config.get("llm_provider", "openai").lower()
if provider.startswith("custom"):
api_key = os.getenv("CUSTOM_API_KEY")
else:
api_key = os.getenv("OPENAI_API_KEY")
client = OpenAI(base_url=config["backend_url"], api_key=api_key)
response = client.responses.create(
model=config["quick_think_llm"],

2
tradingagents/default_config.py
View File

@ -12,7 +12,7 @@ DEFAULT_CONFIG = {
"llm_provider": "openai",
"deep_think_llm": "o4-mini",
"quick_think_llm": "gpt-4o-mini",
"backend_url": os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1"),
"backend_url": "https://api.openai.com/v1", # Will be updated based on selected provider
# Debate and discussion settings
"max_debate_rounds": 1,
"max_risk_discuss_rounds": 1,

19
tradingagents/graph/trading_graph.py
View File

@ -58,13 +58,22 @@ class TradingAgentsGraph:
)
# Initialize LLMs
if self.config["llm_provider"].lower() == "openai" or self.config["llm_provider"] == "ollama" or self.config["llm_provider"] == "openrouter":
self.deep_thinking_llm = ChatOpenAI(model=self.config["deep_think_llm"], openai_api_base=self.config["backend_url"])
self.quick_thinking_llm = ChatOpenAI(model=self.config["quick_think_llm"], openai_api_base=self.config["backend_url"])
elif self.config["llm_provider"].lower() == "anthropic":
provider = self.config["llm_provider"].lower()
if provider == "openai" or provider == "ollama" or provider == "openrouter":
from tradingagents.utils.provider_utils import get_api_key_for_provider
api_key = get_api_key_for_provider(self.config)
self.deep_thinking_llm = ChatOpenAI(model=self.config["deep_think_llm"], openai_api_base=self.config["backend_url"], openai_api_key=api_key)
self.quick_thinking_llm = ChatOpenAI(model=self.config["quick_think_llm"], openai_api_base=self.config["backend_url"], openai_api_key=api_key)
elif provider.startswith("custom"):
# Custom provider uses OpenAI-compatible interface
custom_api_key = os.getenv("CUSTOM_API_KEY")
self.deep_thinking_llm = ChatOpenAI(model=self.config["deep_think_llm"], openai_api_base=self.config["backend_url"], openai_api_key=custom_api_key)
self.quick_thinking_llm = ChatOpenAI(model=self.config["quick_think_llm"], openai_api_base=self.config["backend_url"], openai_api_key=custom_api_key)
elif provider == "anthropic":
self.deep_thinking_llm = ChatAnthropic(model=self.config["deep_think_llm"], base_url=self.config["backend_url"])
self.quick_thinking_llm = ChatAnthropic(model=self.config["quick_think_llm"], base_url=self.config["backend_url"])
elif self.config["llm_provider"].lower() == "google":
elif provider == "google":
self.deep_thinking_llm = ChatGoogleGenerativeAI(model=self.config["deep_think_llm"])
self.quick_thinking_llm = ChatGoogleGenerativeAI(model=self.config["quick_think_llm"])
else:

34
tradingagents/utils/provider_utils.py Normal file
View File

@ -0,0 +1,34 @@
"""
Utility functions for LLM provider configuration and API key management.
"""
import os
def get_api_key_for_provider(config):
"""Get the appropriate API key based on the provider.
Args:
config (dict): Configuration dictionary containing llm_provider
Returns:
str: The API key for the provider, or None if not found
"""
provider = config.get("llm_provider", "openai").lower()
# Map providers to their environment variables
api_key_mapping = {
"openai": "OPENAI_API_KEY",
"anthropic": "ANTHROPIC_API_KEY",
"google": "GOOGLE_API_KEY",
"openrouter": "OPENROUTER_API_KEY",
"ollama": "OLLAMA_API_KEY",
}
env_var = api_key_mapping.get(provider, "OPENAI_API_KEY")
api_key = os.getenv(env_var)
if not api_key and provider != "ollama": # Ollama typically doesn't need API keys
print(f"Warning: {env_var} not found in environment variables")
return api_key