agents/pipeline.py

import logging
from agents.trader import Trader
from agents.analysts import TechnicalAnalyst, NewsAnalyst, SentimentAnalyst
from agents.researcher import Researcher
from agents.risk_manager import RiskManager
from llm.openrouter import OpenRouterClient

logger = logging.getLogger(__name__)


class SingleAgentPipeline:
    """Benchmark A: single Trader agent with price + indicators."""

    def __init__(self, trader: Trader):
        self.trader = trader

    def decide(self, market_data: dict) -> dict:
        decision = self.trader.run(market_data)
        return {
            "decision": decision,
            "agent_outputs": {"trader": decision},
        }


class SimplePipeline:
    """Benchmark B: TechnicalAnalyst + NewsAnalyst -> Trader."""

    def __init__(self, analysts: list, trader: Trader):
        self.technical = next(a for a in analysts if isinstance(a, TechnicalAnalyst))
        self.news = next(a for a in analysts if isinstance(a, NewsAnalyst))
        self.trader = trader

    def decide(self, market_data: dict) -> dict:
        tech_analysis = self.technical.run(market_data)
        news_analysis = self.news.run(market_data)

        trader_context = {
            **market_data,
            "tech_analysis": tech_analysis,
            "news_analysis": news_analysis,
        }
        decision = self.trader.run(trader_context)

        return {
            "decision": decision,
            "agent_outputs": {
                "technical_analyst": tech_analysis,
                "news_analyst": news_analysis,
                "trader": decision,
            },
        }


class FullPipeline:
    """Benchmark C: Technical + Sentiment + News -> Researcher -> RiskManager -> Trader."""

    def __init__(self, analysts: list, researcher: Researcher, risk_manager: RiskManager, trader: Trader):
        self.technical = next(a for a in analysts if isinstance(a, TechnicalAnalyst))
        self.sentiment = next(a for a in analysts if isinstance(a, SentimentAnalyst))
        self.news = next(a for a in analysts if isinstance(a, NewsAnalyst))
        self.researcher = researcher
        self.risk_manager = risk_manager
        self.trader = trader

    def decide(self, market_data: dict) -> dict:
        # Phase 1: analysts
        tech_analysis = self.technical.run(market_data)
        news_analysis = self.news.run(market_data)
        sentiment_analysis = self.sentiment.run(market_data)

        # Phase 2: researcher bull/bear debate
        research_context = {
            **market_data,
            "tech_analysis": tech_analysis,
            "news_analysis": news_analysis,
            "sentiment_analysis": sentiment_analysis,
        }
        research = self.researcher.run(research_context)

        # Phase 3: risk manager
        portfolio = market_data.get("portfolio", {})
        risk_context = {
            "recommendation": research,
            "portfolio": portfolio,
        }
        risk_decision = self.risk_manager.run(risk_context)

        # Phase 4: final trader decision
        trader_context = {
            **market_data,
            "research": research,
            "risk_decision": risk_decision,
        }
        decision = self.trader.run(trader_context)

        return {
            "decision": decision,
            "agent_outputs": {
                "technical_analyst": tech_analysis,
                "news_analyst": news_analysis,
                "sentiment_analyst": sentiment_analysis,
                "researcher": research,
                "risk_manager": risk_decision,
                "trader": decision,
            },
        }


def build_pipeline(benchmark: str, model: str):
    """Factory: build the correct pipeline for benchmark A/B/C."""
    llm = OpenRouterClient(model=model)

    if benchmark == "A":
        return SingleAgentPipeline(trader=Trader(llm, benchmark="A"))

    if benchmark == "B":
        return SimplePipeline(
            analysts=[TechnicalAnalyst(llm), NewsAnalyst(llm)],
            trader=Trader(llm, benchmark="B"),
        )

    if benchmark == "C":
        return FullPipeline(
            analysts=[TechnicalAnalyst(llm), SentimentAnalyst(llm), NewsAnalyst(llm)],
            researcher=Researcher(llm),
            risk_manager=RiskManager(llm),
            trader=Trader(llm, benchmark="C"),
        )

    raise ValueError(f"Unknown benchmark: {benchmark}")