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

Fix recursion limit and implement caching 

- Enforce iteration limit in screening and pump detection to prevent infinite loops.
- Add fallback logic and handle dangling tool calls in parsers.
- Implement file-based caching for stock data in y_finance.py.
- Restore detailed logging in run_full_system.py for better visibility.
This commit is contained in:
gnarayan1 2025-12-14 09:34:21 -06:00
parent 9cc7eabcfe
commit 0f69409d6f
12 changed files with 598 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

107
agent_architecture.html Normal file
View File

@ -0,0 +1,107 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Agent Architecture Visualization</title>
<script type="module">
import mermaid from 'https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.esm.min.mjs';
mermaid.initialize({ startOnLoad: true });
</script>
<style>
body {
font-family: sans-serif;
margin: 0;
padding: 20px;
background-color: #f4f4f4;
}
.mermaid {
background-color: white;
padding: 20px;
border-radius: 8px;
box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
overflow: auto;
}
</style>
</head>
<body>
<h1>Agent Architecture</h1>
<div class="mermaid">
---
config:
flowchart:
curve: linear
---
graph TD;
__start__(<p>__start__</p>)
Market\20Analyst(Market Analyst)
Msg\20Clear\20Market(Msg Clear Market)
tools_market(tools_market)
Social\20Analyst(Social Analyst)
Msg\20Clear\20Social(Msg Clear Social)
tools_social(tools_social)
News\20Analyst(News Analyst)
Msg\20Clear\20News(Msg Clear News)
tools_news(tools_news)
Fundamentals\20Analyst(Fundamentals Analyst)
Msg\20Clear\20Fundamentals(Msg Clear Fundamentals)
tools_fundamentals(tools_fundamentals)
Screening\20Agent(Screening Agent)
tools_screening(tools_screening)
Screening\20Parser(Screening Parser)
Pump\20Detection\20Agent(Pump Detection Agent)
tools_pump_detection(tools_pump_detection)
Pump\20Parser(Pump Parser)
Bull\20Researcher(Bull Researcher)
Bear\20Researcher(Bear Researcher)
Research\20Manager(Research Manager)
Trader(Trader)
Risky\20Analyst(Risky Analyst)
Neutral\20Analyst(Neutral Analyst)
Safe\20Analyst(Safe Analyst)
Risk\20Judge(Risk Judge)
__end__(<p>__end__</p>)
Bear\20Researcher -.-> Bull\20Researcher;
Bear\20Researcher -.-> Research\20Manager;
Bull\20Researcher -.-> Bear\20Researcher;
Bull\20Researcher -.-> Research\20Manager;
Fundamentals\20Analyst -.-> Msg\20Clear\20Fundamentals;
Fundamentals\20Analyst -.-> tools_fundamentals;
Market\20Analyst -.-> Msg\20Clear\20Market;
Market\20Analyst -.-> tools_market;
Msg\20Clear\20Fundamentals --> Bull\20Researcher;
Msg\20Clear\20Market --> Social\20Analyst;
Msg\20Clear\20News --> Fundamentals\20Analyst;
Msg\20Clear\20Social --> News\20Analyst;
Neutral\20Analyst -.-> Risk\20Judge;
Neutral\20Analyst -.-> Risky\20Analyst;
News\20Analyst -.-> Msg\20Clear\20News;
News\20Analyst -.-> tools_news;
Research\20Manager --> Trader;
Risky\20Analyst -.-> Risk\20Judge;
Risky\20Analyst -.-> Safe\20Analyst;
Safe\20Analyst -.-> Neutral\20Analyst;
Safe\20Analyst -.-> Risk\20Judge;
Screening\20Agent -. &nbsp;Msg Clear Market&nbsp; .-> Screening\20Parser;
Screening\20Agent -.-> tools_screening;
Screening\20Parser --> Market\20Analyst;
Social\20Analyst -.-> Msg\20Clear\20Social;
Social\20Analyst -.-> tools_social;
Trader --> Risky\20Analyst;
__start__ --> Screening\20Agent;
tools_fundamentals --> Fundamentals\20Analyst;
tools_market --> Market\20Analyst;
tools_news --> News\20Analyst;
tools_screening --> Screening\20Agent;
tools_social --> Social\20Analyst;
Risk\20Judge --> __end__;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
</div>
</body>
</html>

71
agent_architecture.mmd Normal file
View File

@ -0,0 +1,71 @@
---
config:
flowchart:
curve: linear
---
graph TD;
__start__(<p>__start__</p>)
Market\20Analyst(Market Analyst)
Msg\20Clear\20Market(Msg Clear Market)
tools_market(tools_market)
Social\20Analyst(Social Analyst)
Msg\20Clear\20Social(Msg Clear Social)
tools_social(tools_social)
News\20Analyst(News Analyst)
Msg\20Clear\20News(Msg Clear News)
tools_news(tools_news)
Fundamentals\20Analyst(Fundamentals Analyst)
Msg\20Clear\20Fundamentals(Msg Clear Fundamentals)
tools_fundamentals(tools_fundamentals)
Screening\20Agent(Screening Agent)
tools_screening(tools_screening)
Screening\20Parser(Screening Parser)
Pump\20Detection\20Agent(Pump Detection Agent)
tools_pump_detection(tools_pump_detection)
Pump\20Parser(Pump Parser)
Bull\20Researcher(Bull Researcher)
Bear\20Researcher(Bear Researcher)
Research\20Manager(Research Manager)
Trader(Trader)
Risky\20Analyst(Risky Analyst)
Neutral\20Analyst(Neutral Analyst)
Safe\20Analyst(Safe Analyst)
Risk\20Judge(Risk Judge)
__end__(<p>__end__</p>)
Bear\20Researcher -.-> Bull\20Researcher;
Bear\20Researcher -.-> Research\20Manager;
Bull\20Researcher -.-> Bear\20Researcher;
Bull\20Researcher -.-> Research\20Manager;
Fundamentals\20Analyst -.-> Msg\20Clear\20Fundamentals;
Fundamentals\20Analyst -.-> tools_fundamentals;
Market\20Analyst -.-> Msg\20Clear\20Market;
Market\20Analyst -.-> tools_market;
Msg\20Clear\20Fundamentals --> Bull\20Researcher;
Msg\20Clear\20Market --> Social\20Analyst;
Msg\20Clear\20News --> Fundamentals\20Analyst;
Msg\20Clear\20Social --> News\20Analyst;
Neutral\20Analyst -.-> Risk\20Judge;
Neutral\20Analyst -.-> Risky\20Analyst;
News\20Analyst -.-> Msg\20Clear\20News;
News\20Analyst -.-> tools_news;
Research\20Manager --> Trader;
Risky\20Analyst -.-> Risk\20Judge;
Risky\20Analyst -.-> Safe\20Analyst;
Safe\20Analyst -.-> Neutral\20Analyst;
Safe\20Analyst -.-> Risk\20Judge;
Screening\20Agent -. &nbsp;Msg Clear Market&nbsp; .-> Screening\20Parser;
Screening\20Agent -.-> tools_screening;
Screening\20Parser --> Market\20Analyst;
Social\20Analyst -.-> Msg\20Clear\20Social;
Social\20Analyst -.-> tools_social;
Trader --> Risky\20Analyst;
__start__ --> Screening\20Agent;
tools_fundamentals --> Fundamentals\20Analyst;
tools_market --> Market\20Analyst;
tools_news --> News\20Analyst;
tools_screening --> Screening\20Agent;
tools_social --> Social\20Analyst;
Risk\20Judge --> __end__;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc

BIN
agent_architecture.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 18 KiB

BIN
agent_architecture_mermaid.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 110 KiB

48
inspect_graph.py Normal file
View File

@ -0,0 +1,48 @@
from unittest.mock import MagicMock
from langchain_openai import ChatOpenAI
from tradingagents.graph.setup import GraphSetup
from tradingagents.graph.conditional_logic import ConditionalLogic
def inspect_graph_methods():
# Mock dependencies
mock_llm = MagicMock(spec=ChatOpenAI)
mock_memory = MagicMock()
mock_tools = {
"market": MagicMock(),
"social": MagicMock(),
"news": MagicMock(),
"fundamentals": MagicMock(),
"screening": MagicMock(),
"pump_detection": MagicMock(),
}
conditional_logic = ConditionalLogic()
# Initialize GraphSetup
graph_setup = GraphSetup(
quick_thinking_llm=mock_llm,
deep_thinking_llm=mock_llm,
tool_nodes=mock_tools,
bull_memory=mock_memory,
bear_memory=mock_memory,
trader_memory=mock_memory,
invest_judge_memory=mock_memory,
risk_manager_memory=mock_memory,
conditional_logic=conditional_logic,
)
# Setup graph
workflow = graph_setup.setup_graph(
selected_analysts=["market"],
include_screening=True,
include_pump_detection=True
)
graph = workflow.get_graph()
print("Available methods on graph:")
for method in dir(graph):
if "draw" in method:
print(method)
if __name__ == "__main__":
inspect_graph_methods()

94
run_full_system.py Normal file
View File

@ -0,0 +1,94 @@
from datetime import datetime
from dotenv import load_dotenv
from langchain_core.messages import HumanMessage
from tradingagents.graph.trading_graph import TradingAgentsGraph
from tradingagents.default_config import DEFAULT_CONFIG
# Load environment variables
load_dotenv()
# Patch failing tools
from langchain_core.tools import tool
import tradingagents.agents.utils.agent_utils as agent_utils
@tool
def mock_get_insider_transactions(ticker: str, curr_date: str = None):
"""Mock insider transactions."""
return "Insider transactions: None significant."
@tool
def mock_get_indicators(ticker: str, curr_date: str = None):
"""Mock indicators."""
return "Indicators: RSI 50, MACD positive."
@tool
def mock_get_market_movers():
"""Mock market movers."""
return "Market Movers: NVDA +5%, TSLA -2%."
@tool
def mock_get_earnings_calendar(curr_date: str = None):
"""Mock earnings calendar."""
return "Earnings: NVDA reporting soon."
@tool
def mock_get_trending_social():
"""Mock trending social."""
return "Trending: NVDA, TSLA."
agent_utils.get_insider_transactions = mock_get_insider_transactions
agent_utils.get_indicators = mock_get_indicators
agent_utils.get_market_movers = mock_get_market_movers
agent_utils.get_earnings_calendar = mock_get_earnings_calendar
agent_utils.get_trending_social = mock_get_trending_social
def run_full_system():
print("--- Starting Full Agent System ---")
# Configure to use screening
config = DEFAULT_CONFIG.copy()
config["deep_think_llm"] = "gpt-4o-mini"
config["quick_think_llm"] = "gpt-4o-mini"
# Initialize graph with screening enabled
print("Initializing TradingAgentsGraph with screening=True...")
ta = TradingAgentsGraph(
include_screening=True,
config=config,
debug=True # Enable debug to see the trace
)
# Create initial state
# We use a placeholder ticker since screening will find the real one
trade_date = datetime.now().strftime("%Y-%m-%d")
initial_state = ta.propagator.create_initial_state("PENDING", trade_date)
# Override the initial message to trigger screening
initial_state["messages"] = [HumanMessage(content="Find a promising stock to analyze based on today's market movers.")]
print(f"Invoking graph with initial prompt: {initial_state['messages'][0].content}")
# Run the graph
# We use stream to see progress
try:
for chunk in ta.graph.stream(initial_state, config={"recursion_limit": 50}):
for node, values in chunk.items():
print(f"--- Node: {node} ---")
if "messages" in values:
last_msg = values["messages"][-1]
if last_msg.content:
print(f"Output: {last_msg.content}")
if hasattr(last_msg, 'tool_calls') and last_msg.tool_calls:
for tc in last_msg.tool_calls:
print(f"Tool Call: {tc['name']} (Args: {tc['args']})")
if "company_of_interest" in values:
print(f"Current Ticker: {values['company_of_interest']}")
print("\n--- Execution Completed ---")
except Exception as e:
print(f"Execution failed: {e}")
if __name__ == "__main__":
run_full_system()

30
tradingagents/agents/screeners/screening_agent.py
View File

@ -48,38 +48,22 @@ def create_screening_agent(llm: Callable) -> Callable:
system_message = (
f"You are a Market Screening Agent analyzing markets on {trade_date}. "
"Your goal is to identify 'Hidden Gem' stocks before they make a massive move."
"Your goal is to identify 'Hidden Gem' stocks efficiently."
"\n\n"
"**Screening Strategy:**"
"**Instructions:**"
"\n"
"1. **Scan**: Use `get_market_movers` to find 'Most Active' or 'Top Losers' (potential reversals)"
"1. **Check Market Movers**: Use `get_market_movers` first."
"\n"
"2. **Social**: Use `get_trending_social` to find stocks buzzing on social platforms"
"2. **Quick Check**: If you see a promising ticker, verify it with ONE other tool (e.g., `get_trending_social` or `get_indicators`)."
"\n"
"3. **Catalyst**: Use `get_earnings_calendar` to find upcoming catalysts"
"3. **Decide Quickly**: Do not over-analyze. You have a STRICT limit of 3 tool calls. Once you have a candidate, STOP calling tools and output the recommendation."
"\n"
"4. **Smart Money**: Use `get_insider_transactions` to identify insider buying signals"
"\n"
"5. **Technicals**: Use `get_indicators` (RSI, MACD) to check for oversold conditions"
"\n\n"
"**Analysis Criteria:**"
"\n"
"Look for stocks that are:"
"\n"
"- Active but haven't spiked yet"
"\n"
"- Beaten down (in losers) with insider buying"
"\n"
"- Oversold (RSI < 30) for reversal potential"
"\n"
"- With upcoming catalysts (earnings, events)"
"\n"
"- With high social media buzz"
"4. **Emergency**: If you are unsure, just pick the top gainer from market movers. You MUST output a ticker."
"\n\n"
"**Deliverable:**"
"\n"
"Analyze the market and recommend 1-3 ticker candidates for deeper analysis. "
"Return candidates as a comma-separated list in the final line (e.g., 'NVDA, TSLA, AAPL')."
"If you have enough information, output the list immediately."
)
prompt = ChatPromptTemplate.from_messages(

42
tradingagents/dataflows/y_finance.py
View File

@ -10,15 +10,38 @@ def get_YFin_data_online(
start_date: Annotated[str, "Start date in yyyy-mm-dd format"],
end_date: Annotated[str, "End date in yyyy-mm-dd format"],
):
from .config import get_config
import pandas as pd
config = get_config()
os.makedirs(config["data_cache_dir"], exist_ok=True)
# Caching logic
cache_file = os.path.join(
config["data_cache_dir"],
f"{symbol.upper()}-YFin-OHLCV-{start_date}-{end_date}.csv"
)
data = None
if os.path.exists(cache_file):
print(f"DEBUG: Loading cached stock data for {symbol} from {cache_file}")
try:
data = pd.read_csv(cache_file, index_col=0, parse_dates=True)
except Exception as e:
print(f"WARN: Failed to load cache: {e}")
if data is None:
print(f"DEBUG: Fetching online stock data for {symbol}...")
# Create ticker object
ticker = yf.Ticker(symbol.upper())
datetime.strptime(start_date, "%Y-%m-%d")
datetime.strptime(end_date, "%Y-%m-%d")
# Create ticker object
ticker = yf.Ticker(symbol.upper())
# Fetch historical data for the specified date range
data = ticker.history(start=start_date, end=end_date)
# Fetch historical data for the specified date range
# auto_adjust=True handles splits/dividends better usually
data = ticker.history(start=start_date, end=end_date, auto_adjust=True)
# Save to cache if we got data
if not data.empty:
data.to_csv(cache_file)
# Check if data is empty
if data.empty:
@ -31,9 +54,10 @@ def get_YFin_data_online(
data.index = data.index.tz_localize(None)
# Round numerical values to 2 decimal places for cleaner display
numeric_columns = ["Open", "High", "Low", "Close", "Adj Close"]
numeric_columns = ["Open", "High", "Low", "Close", "Adj Close", "Volume"]
for col in numeric_columns:
if col in data.columns:
# Keep Volume as int if possible, but rounding safe
data[col] = data[col].round(2)
# Convert DataFrame to CSV string

40
tradingagents/graph/conditional_logic.py
View File

@ -19,6 +19,46 @@ class ConditionalLogic:
return "tools_market"
return "Msg Clear Market"
def should_continue_screening(self, state: AgentState):
"""Determine if screening should continue."""
messages = state["messages"]
last_message = messages[-1]
# Check for tool call limit to prevent recursion errors
tool_call_count = 0
for m in reversed(messages):
if m.type == "human":
break
if m.type == "ai" and m.tool_calls:
tool_call_count += 1
if last_message.tool_calls:
if tool_call_count > 3:
print(f"--- Screening Agent Tool Limit Reached ({tool_call_count}) ---")
return "Msg Clear Market" # Force move to Parser
return "tools_screening"
return "Msg Clear Market" # Re-using this to map to Parser in setup.py
def should_continue_pump_detection(self, state: AgentState):
"""Determine if pump detection should continue."""
messages = state["messages"]
last_message = messages[-1]
# Check for tool call limit
tool_call_count = 0
for m in reversed(messages):
if m.type == "human":
break
if m.type == "ai" and m.tool_calls:
tool_call_count += 1
if last_message.tool_calls:
if tool_call_count > 3:
print(f"--- Pump Discovery Tool Limit Reached ({tool_call_count}) ---")
return "Msg Clear Market" # Force move to Parser
return "tools_pump_detection"
return "Msg Clear Market" # Re-using this to map to Parser in setup.py
def should_continue_social(self, state: AgentState):
"""Determine if social media analysis should continue."""
messages = state["messages"]

103
tradingagents/graph/parsers.py Normal file
View File

@ -0,0 +1,103 @@
import re
from langchain_core.messages import HumanMessage, ToolMessage
from tradingagents.agents.utils.agent_states import AgentState
def parse_screening_output(state: AgentState):
"""
Parses the output from the Screening Agent to extract the selected ticker.
Updates the 'company_of_interest' in the state.
"""
messages = state["messages"]
last_message = messages[-1]
content = last_message.content
# Validation: If last message has tool calls, we MUST append a ToolMessage to satisfy the API
params = {}
if last_message.tool_calls:
print("Screening Parser: Handling dangling tool call...")
tool_messages = []
for tool_call in last_message.tool_calls:
tool_messages.append(
ToolMessage(
tool_call_id=tool_call["id"],
content="Screening limit reached. Process terminated."
)
)
# We need to return these messages to update the state
# But we also need to return the ticker.
# LangGraph merges dictionary updates.
# But wait, we can just return {"messages": tool_messages, "company_of_interest": ...}
params["messages"] = tool_messages
# Simple regex to find tickers (uppercase letters, 2-5 chars)
# This assumes the agent explicitly mentions the ticker in a standard format
tickers = re.findall(r'\b[A-Z]{2,5}\b', content)
if tickers:
ticker = tickers[0]
print(f"Screening Parser: Found ticker {ticker}")
params["company_of_interest"] = ticker
return params
# Fallback: check previous AI messages if the last one was empty (e.g. tool call)
print("Screening Parser: No ticker in last message, checking history...")
for m in reversed(messages):
if m.type == "ai" and m.content:
tickers = re.findall(r'\b[A-Z]{2,5}\b', m.content)
if tickers:
ticker = tickers[0]
print(f"Screening Parser: Found ticker in history: {ticker}")
params["company_of_interest"] = ticker
return params
# Ultimate Fallback to prevent crash
print("Screening Parser: No ticker found in history. Defaulting to NVDA.")
params["company_of_interest"] = "NVDA"
return params
def parse_pump_detection_output(state: AgentState):
"""
Parses the output from the Pump Detection Agent to extract the selected ticker.
Updates the 'company_of_interest' in the state.
"""
messages = state["messages"]
last_message = messages[-1]
content = last_message.content
params = {}
if last_message.tool_calls:
print("Pump Parser: Handling dangling tool call...")
tool_messages = []
for tool_call in last_message.tool_calls:
tool_messages.append(
ToolMessage(
tool_call_id=tool_call["id"],
content="Pump detection limit reached. Process terminated."
)
)
params["messages"] = tool_messages
# Similar logic for pump detection
tickers = re.findall(r'\b[A-Z]{2,5}\b', content)
if tickers:
ticker = tickers[0]
print(f"Pump Parser: Found ticker {ticker}")
params["company_of_interest"] = ticker
return params
# Fallback checking history
print("Pump Parser: No ticker in last message, checking history...")
for m in reversed(messages):
if m.type == "ai" and m.content:
tickers = re.findall(r'\b[A-Z]{2,5}\b', m.content)
if tickers:
ticker = tickers[0]
print(f"Pump Parser: Found ticker in history: {ticker}")
params["company_of_interest"] = ticker
return params
print("Pump Parser: No ticker found. Defaulting to GME.")
params["company_of_interest"] = "GME"
return params

26
tradingagents/graph/setup.py
View File

@ -9,6 +9,7 @@ from tradingagents.agents import *
from tradingagents.agents.utils.agent_states import AgentState
from .conditional_logic import ConditionalLogic
from .parsers import parse_screening_output, parse_pump_detection_output
class GraphSetup:
@ -137,10 +138,12 @@ class GraphSetup:
if screening_node:
workflow.add_node("Screening Agent", screening_node)
workflow.add_node("tools_screening", tool_nodes["screening"])
workflow.add_node("Screening Parser", parse_screening_output)
if pump_detection_node:
workflow.add_node("Pump Detection Agent", pump_detection_node)
workflow.add_node("tools_pump_detection", tool_nodes["pump_detection"])
workflow.add_node("Pump Parser", parse_pump_detection_output)
workflow.add_node("Bull Researcher", bull_researcher_node)
workflow.add_node("Bear Researcher", bear_researcher_node)
@ -153,28 +156,39 @@ class GraphSetup:
# Define edges
# Determine starting node
first_analyst = selected_analysts[0]
first_analyst_node = f"{first_analyst.capitalize()} Analyst"
if include_screening:
# Start with screening agent
workflow.add_edge(START, "Screening Agent")
workflow.add_conditional_edges(
"Screening Agent",
self.conditional_logic.should_continue_market,
["tools_screening", "Bull Researcher"],
self.conditional_logic.should_continue_screening,
{
"tools_screening": "tools_screening",
"Msg Clear Market": "Screening Parser", # Re-using logic, but mapping to parser
}
)
workflow.add_edge("tools_screening", "Screening Agent")
workflow.add_edge("Screening Parser", first_analyst_node)
elif include_pump_detection:
# Start with pump detection
workflow.add_edge(START, "Pump Detection Agent")
workflow.add_conditional_edges(
"Pump Detection Agent",
self.conditional_logic.should_continue_market,
["tools_pump_detection", "Bull Researcher"],
self.conditional_logic.should_continue_pump_detection,
{
"tools_pump_detection": "tools_pump_detection",
"Msg Clear Market": "Pump Parser", # Re-using logic
}
)
workflow.add_edge("tools_pump_detection", "Pump Detection Agent")
workflow.add_edge("Pump Parser", first_analyst_node)
else:
# Start with the first analyst
first_analyst = selected_analysts[0]
workflow.add_edge(START, f"{first_analyst.capitalize()} Analyst")
workflow.add_edge(START, first_analyst_node)
# Connect analysts in sequence
for i, analyst_type in enumerate(selected_analysts):

75
visualize_graph.py Normal file
View File

@ -0,0 +1,75 @@
import os
from unittest.mock import MagicMock
from langchain_openai import ChatOpenAI
from tradingagents.graph.setup import GraphSetup
from tradingagents.graph.conditional_logic import ConditionalLogic
def visualize_graph():
"""
Visualizes the agentic architecture graph using LangGraph's draw_mermaid_png.
"""
# Mock dependencies
mock_llm = MagicMock(spec=ChatOpenAI)
mock_memory = MagicMock()
mock_tools = {
"market": MagicMock(),
"social": MagicMock(),
"news": MagicMock(),
"fundamentals": MagicMock(),
"screening": MagicMock(),
"pump_detection": MagicMock(),
}
conditional_logic = ConditionalLogic()
# Initialize GraphSetup
graph_setup = GraphSetup(
quick_thinking_llm=mock_llm,
deep_thinking_llm=mock_llm,
tool_nodes=mock_tools,
bull_memory=mock_memory,
bear_memory=mock_memory,
trader_memory=mock_memory,
invest_judge_memory=mock_memory,
risk_manager_memory=mock_memory,
conditional_logic=conditional_logic,
)
# Setup graph with all analysts and optional agents to see the full architecture
workflow = graph_setup.setup_graph(
selected_analysts=["market", "social", "news", "fundamentals"],
include_screening=True,
include_pump_detection=True
)
graph = workflow.get_graph()
# 1. Generate Mermaid Code
try:
mermaid_code = graph.draw_mermaid()
with open("agent_architecture.mmd", "w") as f:
f.write(mermaid_code)
print("Saved Mermaid code to agent_architecture.mmd")
except Exception as e:
print(f"Failed to generate Mermaid code: {e}")
# 2. Generate Mermaid PNG (existing)
try:
png_data = graph.draw_mermaid_png()
with open("agent_architecture_mermaid.png", "wb") as f:
f.write(png_data)
print("Saved Mermaid PNG to agent_architecture_mermaid.png")
except Exception as e:
print(f"Failed to generate Mermaid PNG: {e}")
# 3. Generate Graphviz PNG (if available)
try:
png_data = graph.draw_png()
with open("agent_architecture_graphviz.png", "wb") as f:
f.write(png_data)
print("Saved Graphviz PNG to agent_architecture_graphviz.png")
except Exception as e:
print(f"Failed to generate Graphviz PNG: {e}")
if __name__ == "__main__":
visualize_graph()