193 lines
8.0 KiB
Python
193 lines
8.0 KiB
Python
from collections import deque
|
|
import datetime
|
|
|
|
|
|
class MessageBuffer:
|
|
# Fixed teams that always run (not user-selectable)
|
|
FIXED_AGENTS = {
|
|
"Research Team": ["Bull Researcher", "Bear Researcher", "Research Manager"],
|
|
"Trading Team": ["Trader"],
|
|
"Risk Management": ["Aggressive Analyst", "Neutral Analyst", "Conservative Analyst"],
|
|
"Portfolio Management": ["Portfolio Manager"],
|
|
}
|
|
|
|
# Analyst name mapping
|
|
ANALYST_MAPPING = {
|
|
"market": "Market Analyst",
|
|
"social": "Social Analyst",
|
|
"news": "News Analyst",
|
|
"fundamentals": "Fundamentals Analyst",
|
|
}
|
|
|
|
# Report section mapping: section -> (analyst_key for filtering, finalizing_agent)
|
|
# analyst_key: which analyst selection controls this section (None = always included)
|
|
# finalizing_agent: which agent must be "completed" for this report to count as done
|
|
REPORT_SECTIONS = {
|
|
"market_report": ("market", "Market Analyst"),
|
|
"sentiment_report": ("social", "Social Analyst"),
|
|
"news_report": ("news", "News Analyst"),
|
|
"fundamentals_report": ("fundamentals", "Fundamentals Analyst"),
|
|
"investment_plan": (None, "Research Manager"),
|
|
"trader_investment_plan": (None, "Trader"),
|
|
"final_trade_decision": (None, "Portfolio Manager"),
|
|
}
|
|
|
|
def __init__(self, max_length=100):
|
|
self.messages = deque(maxlen=max_length)
|
|
self.tool_calls = deque(maxlen=max_length)
|
|
self.current_report = None
|
|
self.final_report = None # Store the complete final report
|
|
self.agent_status = {}
|
|
self.current_agent = None
|
|
self.report_sections = {}
|
|
self.selected_analysts = []
|
|
self._last_message_id = None
|
|
self._last_updated_section = None
|
|
|
|
def init_for_analysis(self, selected_analysts):
|
|
"""Initialize agent status and report sections based on selected analysts.
|
|
|
|
Args:
|
|
selected_analysts: List of analyst type strings (e.g., ["market", "news"])
|
|
"""
|
|
self.selected_analysts = [a.lower() for a in selected_analysts]
|
|
|
|
# Build agent_status dynamically
|
|
self.agent_status = {}
|
|
|
|
# Add selected analysts
|
|
for analyst_key in self.selected_analysts:
|
|
if analyst_key in self.ANALYST_MAPPING:
|
|
self.agent_status[self.ANALYST_MAPPING[analyst_key]] = "pending"
|
|
|
|
# Add fixed teams
|
|
for team_agents in self.FIXED_AGENTS.values():
|
|
for agent in team_agents:
|
|
self.agent_status[agent] = "pending"
|
|
|
|
# Build report_sections dynamically
|
|
self.report_sections = {}
|
|
for section, (analyst_key, _) in self.REPORT_SECTIONS.items():
|
|
if analyst_key is None or analyst_key in self.selected_analysts:
|
|
self.report_sections[section] = None
|
|
|
|
# Reset other state
|
|
self.current_report = None
|
|
self.final_report = None
|
|
self.current_agent = None
|
|
self.messages.clear()
|
|
self.tool_calls.clear()
|
|
self._last_message_id = None
|
|
self._last_updated_section = None
|
|
|
|
def get_completed_reports_count(self):
|
|
"""Count reports that are finalized (their finalizing agent is completed).
|
|
|
|
A report is considered complete when:
|
|
1. The report section has content (not None), AND
|
|
2. The agent responsible for finalizing that report has status "completed"
|
|
|
|
This prevents interim updates (like debate rounds) from counting as completed.
|
|
"""
|
|
count = 0
|
|
for section in self.report_sections:
|
|
if section not in self.REPORT_SECTIONS:
|
|
continue
|
|
_, finalizing_agent = self.REPORT_SECTIONS[section]
|
|
# Report is complete if it has content AND its finalizing agent is done
|
|
has_content = self.report_sections.get(section) is not None
|
|
agent_done = self.agent_status.get(finalizing_agent) == "completed"
|
|
if has_content and agent_done:
|
|
count += 1
|
|
return count
|
|
|
|
def add_message(self, message_type, content):
|
|
timestamp = datetime.datetime.now().strftime("%H:%M:%S")
|
|
self.messages.append((timestamp, message_type, content))
|
|
|
|
def add_tool_call(self, tool_name, args):
|
|
timestamp = datetime.datetime.now().strftime("%H:%M:%S")
|
|
self.tool_calls.append((timestamp, tool_name, args))
|
|
|
|
def update_agent_status(self, agent, status):
|
|
if agent in self.agent_status:
|
|
self.agent_status[agent] = status
|
|
self.current_agent = agent
|
|
|
|
def update_report_section(self, section_name, content):
|
|
if section_name in self.report_sections:
|
|
self.report_sections[section_name] = content
|
|
self._last_updated_section = section_name
|
|
self._update_current_report()
|
|
|
|
def _update_current_report(self):
|
|
# For the panel display, only show the most recently updated section
|
|
latest_section = self._last_updated_section
|
|
latest_content = self.report_sections.get(latest_section) if latest_section else None
|
|
|
|
# Fallback if section tracking is unavailable
|
|
if latest_content is None:
|
|
for section, content in self.report_sections.items():
|
|
if content is not None:
|
|
latest_section = section
|
|
latest_content = content
|
|
|
|
if latest_section and latest_content:
|
|
# Format the current section for display
|
|
section_titles = {
|
|
"market_report": "Market Analysis",
|
|
"sentiment_report": "Social Sentiment",
|
|
"news_report": "News Analysis",
|
|
"fundamentals_report": "Fundamentals Analysis",
|
|
"investment_plan": "Research Team Decision",
|
|
"trader_investment_plan": "Trading Team Plan",
|
|
"final_trade_decision": "Portfolio Management Decision",
|
|
}
|
|
self.current_report = (
|
|
f"### {section_titles[latest_section]}\n{latest_content}"
|
|
)
|
|
|
|
# Update the final complete report
|
|
self._update_final_report()
|
|
|
|
def _update_final_report(self):
|
|
report_parts = []
|
|
|
|
# Analyst Team Reports - use .get() to handle missing sections
|
|
analyst_sections = ["market_report", "sentiment_report", "news_report", "fundamentals_report"]
|
|
if any(self.report_sections.get(section) for section in analyst_sections):
|
|
report_parts.append("## Analyst Team Reports")
|
|
if self.report_sections.get("market_report"):
|
|
report_parts.append(
|
|
f"### Market Analysis\n{self.report_sections['market_report']}"
|
|
)
|
|
if self.report_sections.get("sentiment_report"):
|
|
report_parts.append(
|
|
f"### Social Sentiment\n{self.report_sections['sentiment_report']}"
|
|
)
|
|
if self.report_sections.get("news_report"):
|
|
report_parts.append(
|
|
f"### News Analysis\n{self.report_sections['news_report']}"
|
|
)
|
|
if self.report_sections.get("fundamentals_report"):
|
|
report_parts.append(
|
|
f"### Fundamentals Analysis\n{self.report_sections['fundamentals_report']}"
|
|
)
|
|
|
|
# Research Team Reports
|
|
if self.report_sections.get("investment_plan"):
|
|
report_parts.append("## Research Team Decision")
|
|
report_parts.append(f"{self.report_sections['investment_plan']}")
|
|
|
|
# Trading Team Reports
|
|
if self.report_sections.get("trader_investment_plan"):
|
|
report_parts.append("## Trading Team Plan")
|
|
report_parts.append(f"{self.report_sections['trader_investment_plan']}")
|
|
|
|
# Portfolio Management Decision
|
|
if self.report_sections.get("final_trade_decision"):
|
|
report_parts.append("## Portfolio Management Decision")
|
|
report_parts.append(f"{self.report_sections['final_trade_decision']}")
|
|
|
|
self.final_report = "\n\n".join(report_parts) if report_parts else None
|