From 811d57fb00bff014c1e13fd0da195b43bf898416 Mon Sep 17 00:00:00 2001 From: mogita Date: Sat, 16 Aug 2025 16:21:03 +0800 Subject: [PATCH] refactor: clean up code as per PEP 8 and review suggestions --- tradingagents/graph/trading_graph.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/tradingagents/graph/trading_graph.py b/tradingagents/graph/trading_graph.py index dfb662b9..3bab4c22 100644 --- a/tradingagents/graph/trading_graph.py +++ b/tradingagents/graph/trading_graph.py @@ -13,6 +13,7 @@ from langchain_google_genai import ChatGoogleGenerativeAI from langgraph.prebuilt import ToolNode from tradingagents.agents import * +from tradingagents.utils.provider_utils import get_api_key_for_provider from tradingagents.default_config import DEFAULT_CONFIG from tradingagents.agents.utils.memory import FinancialSituationMemory from tradingagents.agents.utils.agent_states import ( @@ -60,16 +61,10 @@ class TradingAgentsGraph: # Initialize LLMs 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 + if provider in ("openai", "ollama", "openrouter") or provider.startswith("custom"): 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) + self.deep_thinking_llm = ChatOpenAI(model=self.config["deep_think_llm"], base_url=self.config["backend_url"], api_key=api_key) + self.quick_thinking_llm = ChatOpenAI(model=self.config["quick_think_llm"], base_url=self.config["backend_url"], api_key=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"])