115 lines
4.9 KiB
Python
115 lines
4.9 KiB
Python
import os
|
|
from dotenv import load_dotenv
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_core.messages import HumanMessage
|
|
|
|
from tradingagents.agents.screening_agent import create_screening_agent
|
|
from tradingagents.graph.trading_graph import TradingAgentsGraph
|
|
from tradingagents.default_config import DEFAULT_CONFIG
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
def main():
|
|
print("--- Starting Market Screening ---")
|
|
|
|
# 1. Initialize LLM for screening
|
|
config = DEFAULT_CONFIG.copy()
|
|
llm = ChatOpenAI(model=config["quick_think_llm"])
|
|
|
|
# 2. Create and run Screening Agent
|
|
screener = create_screening_agent(llm)
|
|
|
|
# Initial state for screening
|
|
state = {"messages": [HumanMessage(content="Find me the top gainers today and pick the most interesting one to analyze.")]}
|
|
|
|
print("Running Screening Agent...")
|
|
# In a real graph we would use langgraph, but here we can just invoke the node function for simplicity
|
|
# or build a mini-graph. Let's just invoke the node loop manually for this demo.
|
|
|
|
# Loop to handle multiple rounds of tool calls
|
|
max_iterations = 5
|
|
iteration = 0
|
|
|
|
while iteration < max_iterations:
|
|
iteration += 1
|
|
|
|
# Invoke agent
|
|
result = screener(state)
|
|
state["messages"].extend(result["messages"])
|
|
last_msg = result["messages"][-1]
|
|
|
|
# Check if tool call
|
|
if not last_msg.tool_calls:
|
|
# No more tools, this is the final response
|
|
print(f"Screening Agent Recommendation: {last_msg.content}")
|
|
|
|
# Extract ticker (simple heuristic)
|
|
import re
|
|
tickers = re.findall(r'\b[A-Z]{2,5}\b', last_msg.content)
|
|
|
|
if tickers:
|
|
target_ticker = tickers[0] # Pick the first one
|
|
print(f"Selected Ticker: {target_ticker}")
|
|
|
|
# 3. Run TradingAgentsGraph on the selected ticker
|
|
print(f"--- Starting Analysis for {target_ticker} ---")
|
|
ta = TradingAgentsGraph(debug=True, config=config)
|
|
|
|
from datetime import datetime
|
|
today = datetime.now().strftime("%Y-%m-%d")
|
|
|
|
_, decision = ta.propagate(target_ticker, today)
|
|
print("\nFinal Decision:")
|
|
print(decision)
|
|
else:
|
|
print("No tickers found in recommendation.")
|
|
|
|
break
|
|
|
|
else:
|
|
print(f"Tool Call (Iter {iteration}): {last_msg.tool_calls}")
|
|
|
|
# Execute tools
|
|
from tradingagents.agents.utils.core_stock_tools import get_market_movers, get_earnings_calendar
|
|
from tradingagents.agents.utils.news_data_tools import get_insider_transactions
|
|
from tradingagents.agents.utils.technical_indicators_tools import get_indicators
|
|
from tradingagents.dataflows.social_sentiment import get_trending_social
|
|
|
|
tool_outputs = []
|
|
for tool_call in last_msg.tool_calls:
|
|
output = "Error: Tool not found"
|
|
try:
|
|
if tool_call["name"] == "get_market_movers":
|
|
output = get_market_movers.invoke(tool_call["args"])
|
|
elif tool_call["name"] == "get_earnings_calendar":
|
|
output = get_earnings_calendar.invoke(tool_call["args"])
|
|
elif tool_call["name"] == "get_insider_transactions":
|
|
args = tool_call["args"]
|
|
if "curr_date" not in args:
|
|
from datetime import datetime
|
|
args["curr_date"] = datetime.now().strftime("%Y-%m-%d")
|
|
output = get_insider_transactions.invoke(args)
|
|
elif tool_call["name"] == "get_indicators":
|
|
args = tool_call["args"]
|
|
if "curr_date" not in args:
|
|
from datetime import datetime
|
|
args["curr_date"] = datetime.now().strftime("%Y-%m-%d")
|
|
output = get_indicators.invoke(args)
|
|
elif tool_call["name"] == "get_trending_social":
|
|
output = get_trending_social.invoke(tool_call["args"])
|
|
except Exception as e:
|
|
output = f"Tool execution failed: {str(e)}"
|
|
|
|
tool_outputs.append(
|
|
{"tool_call_id": tool_call["id"], "content": str(output)}
|
|
)
|
|
|
|
# Add tool outputs to messages
|
|
from langchain_core.messages import ToolMessage
|
|
for output in tool_outputs:
|
|
state["messages"].append(ToolMessage(content=output["content"], tool_call_id=output["tool_call_id"]))
|
|
|
|
if __name__ == "__main__":
|
|
main()
|