From 9254cde228358fd1ad9e887ce2fd5dcd5baa34e2 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Sat, 16 Aug 2025 00:36:13 +0000 Subject: [PATCH] feat: Migrate to Gemini and remove other AI providers --- API_KEY_MANAGEMENT.md | 45 ++++++++ README.md | 10 +- cli/main.py | 30 +----- cli/utils.py | 152 --------------------------- main.py | 7 +- tradingagents/agents/utils/memory.py | 16 +-- tradingagents/default_config.py | 7 +- tradingagents/graph/trading_graph.py | 15 +-- 8 files changed, 66 insertions(+), 216 deletions(-) create mode 100644 API_KEY_MANAGEMENT.md diff --git a/API_KEY_MANAGEMENT.md b/API_KEY_MANAGEMENT.md new file mode 100644 index 00000000..4c15abe5 --- /dev/null +++ b/API_KEY_MANAGEMENT.md @@ -0,0 +1,45 @@ +# API Key Management + +This document provides instructions for managing the API keys required to run the TradingAgents framework. + +## Google API Key + +The TradingAgents framework uses Google's Generative AI models. You will need a Google API key to use the service. + +### Obtaining a Google API Key + +1. Go to the [Google AI Studio](https://aistudio.google.com/). +2. Log in with your Google account. +3. Click on "Get API key" to create a new API key. + +### Setting the Google API Key + +To use the Google API key, you need to set it as an environment variable named `GOOGLE_API_KEY`. + +#### For Linux and macOS: + +You can set the environment variable in your shell's configuration file (e.g., `.bashrc`, `.zshrc`). + +```bash +export GOOGLE_API_KEY="YOUR_GOOGLE_API_KEY" +``` + +Replace `"YOUR_GOOGLE_API_KEY"` with the API key you obtained from the Google AI Studio. + +After adding the line to your configuration file, restart your shell or run the following command to apply the changes: + +```bash +source ~/.bashrc # or source ~/.zshrc +``` + +#### For Windows: + +You can set the environment variable through the system settings. + +1. Search for "Environment Variables" in the Start menu and select "Edit the system environment variables". +2. In the System Properties window, click on the "Environment Variables..." button. +3. In the Environment Variables window, click on "New..." under the "System variables" section. +4. Set the "Variable name" to `GOOGLE_API_KEY` and the "Variable value" to your Google API key. +5. Click "OK" to close all windows. + +You may need to restart your command prompt or IDE for the changes to take effect. diff --git a/README.md b/README.md index cac18691..48d924b0 100644 --- a/README.md +++ b/README.md @@ -119,9 +119,9 @@ You will also need the FinnHub API for financial data. All of our code is implem export FINNHUB_API_KEY=$YOUR_FINNHUB_API_KEY ``` -You will need the OpenAI API for all the agents. +You will need the Google API for all the agents. See `API_KEY_MANAGEMENT.md` for more details. ```bash -export OPENAI_API_KEY=$YOUR_OPENAI_API_KEY +export GOOGLE_API_KEY=$YOUR_GOOGLE_API_KEY ``` ### CLI Usage @@ -150,7 +150,7 @@ An interface will appear showing results as they load, letting you track the age ### Implementation Details -We built TradingAgents with LangGraph to ensure flexibility and modularity. We utilize `o1-preview` and `gpt-4o` as our deep thinking and fast thinking LLMs for our experiments. However, for testing purposes, we recommend you use `o4-mini` and `gpt-4.1-mini` to save on costs as our framework makes **lots of** API calls. +We built TradingAgents with LangGraph to ensure flexibility and modularity. We utilize `gemini-2.5-pro` and `gemini-2.5-flash` as our deep thinking and fast thinking LLMs. ### Python Usage @@ -175,8 +175,8 @@ from tradingagents.default_config import DEFAULT_CONFIG # Create a custom config config = DEFAULT_CONFIG.copy() -config["deep_think_llm"] = "gpt-4.1-nano" # Use a different model -config["quick_think_llm"] = "gpt-4.1-nano" # Use a different model +config["deep_think_llm"] = "gemini-2.5-pro" +config["quick_think_llm"] = "gemini-2.5-flash" config["max_debate_rounds"] = 1 # Increase debate rounds config["online_tools"] = True # Use online tools or cached data diff --git a/cli/main.py b/cli/main.py index 64616ee1..f28a461b 100644 --- a/cli/main.py +++ b/cli/main.py @@ -23,7 +23,7 @@ from rich.rule import Rule from tradingagents.graph.trading_graph import TradingAgentsGraph from tradingagents.default_config import DEFAULT_CONFIG from cli.models import AnalystType -from cli.utils import * +from cli.utils import get_ticker, get_analysis_date, select_analysts, select_research_depth console = Console() @@ -463,32 +463,11 @@ def get_user_selections(): ) selected_research_depth = select_research_depth() - # Step 5: OpenAI backend - console.print( - create_question_box( - "Step 5: OpenAI backend", "Select which service to talk to" - ) - ) - selected_llm_provider, backend_url = select_llm_provider() - - # Step 6: Thinking agents - console.print( - create_question_box( - "Step 6: Thinking Agents", "Select your thinking agents for analysis" - ) - ) - selected_shallow_thinker = select_shallow_thinking_agent(selected_llm_provider) - selected_deep_thinker = select_deep_thinking_agent(selected_llm_provider) - return { "ticker": selected_ticker, "analysis_date": analysis_date, "analysts": selected_analysts, "research_depth": selected_research_depth, - "llm_provider": selected_llm_provider.lower(), - "backend_url": backend_url, - "shallow_thinker": selected_shallow_thinker, - "deep_thinker": selected_deep_thinker, } @@ -739,10 +718,9 @@ def run_analysis(): config = DEFAULT_CONFIG.copy() config["max_debate_rounds"] = selections["research_depth"] config["max_risk_discuss_rounds"] = selections["research_depth"] - config["quick_think_llm"] = selections["shallow_thinker"] - config["deep_think_llm"] = selections["deep_thinker"] - config["backend_url"] = selections["backend_url"] - config["llm_provider"] = selections["llm_provider"].lower() + config["quick_think_llm"] = "gemini-2.5-flash" + config["deep_think_llm"] = "gemini-2.5-pro" + config["llm_provider"] = "google" # Initialize the graph graph = TradingAgentsGraph( diff --git a/cli/utils.py b/cli/utils.py index 7b9682a6..3aab88ca 100644 --- a/cli/utils.py +++ b/cli/utils.py @@ -122,155 +122,3 @@ def select_research_depth() -> int: 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"), - ] - } - - choice = questionary.select( - "Select Your [Quick-Thinking LLM Engine]:", - choices=[ - questionary.Choice(display, value=value) - for display, value in SHALLOW_AGENT_OPTIONS[provider.lower()] - ], - instruction="\n- Use arrow keys to navigate\n- Press Enter to select", - style=questionary.Style( - [ - ("selected", "fg:magenta noinherit"), - ("highlighted", "fg:magenta noinherit"), - ("pointer", "fg:magenta noinherit"), - ] - ), - ).ask() - - if choice is None: - console.print( - "\n[red]No shallow thinking llm engine selected. Exiting...[/red]" - ) - exit(1) - - return choice - - -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"), - ] - } - - choice = questionary.select( - "Select Your [Deep-Thinking LLM Engine]:", - choices=[ - questionary.Choice(display, value=value) - for display, value in DEEP_AGENT_OPTIONS[provider.lower()] - ], - instruction="\n- Use arrow keys to navigate\n- Press Enter to select", - style=questionary.Style( - [ - ("selected", "fg:magenta noinherit"), - ("highlighted", "fg:magenta noinherit"), - ("pointer", "fg:magenta noinherit"), - ] - ), - ).ask() - - if choice is None: - 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.""" - # Define OpenAI api options with their corresponding endpoints - BASE_URLS = [ - ("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"), - ] - - choice = questionary.select( - "Select your LLM Provider:", - choices=[ - questionary.Choice(display, value=(display, value)) - for display, value in BASE_URLS - ], - instruction="\n- Use arrow keys to navigate\n- Press Enter to select", - style=questionary.Style( - [ - ("selected", "fg:magenta noinherit"), - ("highlighted", "fg:magenta noinherit"), - ("pointer", "fg:magenta noinherit"), - ] - ), - ).ask() - - if choice is None: - console.print("\n[red]no OpenAI backend selected. Exiting...[/red]") - exit(1) - - display_name, url = choice - print(f"You selected: {display_name}\tURL: {url}") - - return display_name, url diff --git a/main.py b/main.py index 6c8ae3d9..d3764bb8 100644 --- a/main.py +++ b/main.py @@ -3,10 +3,9 @@ from tradingagents.default_config import DEFAULT_CONFIG # Create a custom config config = DEFAULT_CONFIG.copy() -config["llm_provider"] = "google" # Use a different model -config["backend_url"] = "https://generativelanguage.googleapis.com/v1" # Use a different backend -config["deep_think_llm"] = "gemini-2.0-flash" # Use a different model -config["quick_think_llm"] = "gemini-2.0-flash" # Use a different model +config["llm_provider"] = "google" +config["deep_think_llm"] = "gemini-2.5-pro" +config["quick_think_llm"] = "gemini-2.5-flash" config["max_debate_rounds"] = 1 # Increase debate rounds config["online_tools"] = True # Increase debate rounds diff --git a/tradingagents/agents/utils/memory.py b/tradingagents/agents/utils/memory.py index 69b8ab8c..46f804d6 100644 --- a/tradingagents/agents/utils/memory.py +++ b/tradingagents/agents/utils/memory.py @@ -1,25 +1,17 @@ import chromadb from chromadb.config import Settings -from openai import OpenAI +from langchain_google_genai import GoogleGenerativeAIEmbeddings class FinancialSituationMemory: def __init__(self, name, config): - if config["backend_url"] == "http://localhost:11434/v1": - self.embedding = "nomic-embed-text" - else: - self.embedding = "text-embedding-3-small" - self.client = OpenAI(base_url=config["backend_url"]) + self.embedding = GoogleGenerativeAIEmbeddings(model="models/embedding-001") self.chroma_client = chromadb.Client(Settings(allow_reset=True)) self.situation_collection = self.chroma_client.create_collection(name=name) def get_embedding(self, text): - """Get OpenAI embedding for a text""" - - response = self.client.embeddings.create( - model=self.embedding, input=text - ) - return response.data[0].embedding + """Get Google embedding for a text""" + return self.embedding.embed_query(text) def add_situations(self, situations_and_advice): """Add financial situations and their corresponding advice. Parameter is a list of tuples (situation, rec)""" diff --git a/tradingagents/default_config.py b/tradingagents/default_config.py index 089e9c24..2b1c8222 100644 --- a/tradingagents/default_config.py +++ b/tradingagents/default_config.py @@ -9,10 +9,9 @@ DEFAULT_CONFIG = { "dataflows/data_cache", ), # LLM settings - "llm_provider": "openai", - "deep_think_llm": "o4-mini", - "quick_think_llm": "gpt-4o-mini", - "backend_url": "https://api.openai.com/v1", + "llm_provider": "google", + "deep_think_llm": "gemini-2.5-pro", + "quick_think_llm": "gemini-2.5-flash", # Debate and discussion settings "max_debate_rounds": 1, "max_risk_discuss_rounds": 1, diff --git a/tradingagents/graph/trading_graph.py b/tradingagents/graph/trading_graph.py index 80a29e53..0961173a 100644 --- a/tradingagents/graph/trading_graph.py +++ b/tradingagents/graph/trading_graph.py @@ -6,8 +6,6 @@ import json from datetime import date from typing import Dict, Any, Tuple, List, Optional -from langchain_openai import ChatOpenAI -from langchain_anthropic import ChatAnthropic from langchain_google_genai import ChatGoogleGenerativeAI from langgraph.prebuilt import ToolNode @@ -58,17 +56,8 @@ 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"], base_url=self.config["backend_url"]) - self.quick_thinking_llm = ChatOpenAI(model=self.config["quick_think_llm"], base_url=self.config["backend_url"]) - elif self.config["llm_provider"].lower() == "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": - self.deep_thinking_llm = ChatGoogleGenerativeAI(model=self.config["deep_think_llm"]) - self.quick_thinking_llm = ChatGoogleGenerativeAI(model=self.config["quick_think_llm"]) - else: - raise ValueError(f"Unsupported LLM provider: {self.config['llm_provider']}") + self.deep_thinking_llm = ChatGoogleGenerativeAI(model=self.config["deep_think_llm"]) + self.quick_thinking_llm = ChatGoogleGenerativeAI(model=self.config["quick_think_llm"]) self.toolkit = Toolkit(config=self.config)