backend/agents/supervisor_agent.py

"""
Supervisor Agent — LangGraph orchestrator for HITL chat mode.
Routes user queries to the appropriate tools/specialists.
"""
import json
import logging
import operator
from typing import Annotated, Sequence, TypedDict

from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import ToolNode
from langchain_core.tools import tool

from backend.agents.scan_pipeline import run_full_scan
from backend.agents.regime_agent import detect_regime
from backend.agents.news_agent import fetch_market_news

logger = logging.getLogger(__name__)

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]

@tool
def get_market_regime(market: str = "india") -> str:
    """Get the current market regime, trend, and volatility state for a given market ('india' or 'us')."""
    res = detect_regime(market)
    return json.dumps(res, indent=2)

@tool
def run_swing_scan(market: str = "both", top_n: int = 5) -> str:
    """Run a full pipeline scan to find the top swing trading opportunities. Returns scan stats."""
    res = run_full_scan(market=market, top_n=top_n)
    # Return stats summary to avoid massive text in chat
    return json.dumps(res.get("stats", {"error": "Scan failed"}), indent=2)

@tool
def get_market_news() -> str:
    """Fetch the latest market news from various RSS feeds."""
    news = fetch_market_news()
    # Summarize to avoid massive context
    summary = [{"title": n["title"], "source": n["source"]} for n in news[:10]]
    return json.dumps(summary, indent=2)

tools = [get_market_regime, run_swing_scan, get_market_news]
tool_node = ToolNode(tools)

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
llm_with_tools = llm.bind_tools(tools)

def supervisor_node(state: AgentState):
    messages = state["messages"]
    
    # Prepend system message if not present
    if not any(isinstance(m, SystemMessage) for m in messages):
        sys_msg = SystemMessage(
            content="You are a senior quantitative trading supervisor agent. "
                    "You have access to specialists via tools. "
                    "Use tools to answer queries about market regime or swing trading scans. "
                    "If a tool fails, report the error to the user gracefully."
        )
        messages = [sys_msg] + messages

    response = llm_with_tools.invoke(messages)
    return {"messages": [response]}

def should_continue(state: AgentState) -> str:
    messages = state["messages"]
    last_message = messages[-1]
    if last_message.tool_calls:
        return "tools"
    return END

workflow = StateGraph(AgentState)
workflow.add_node("supervisor", supervisor_node)
workflow.add_node("tools", tool_node)

workflow.add_edge(START, "supervisor")
workflow.add_conditional_edges("supervisor", should_continue, ["tools", END])
workflow.add_edge("tools", "supervisor")

app = workflow.compile()

def invoke_supervisor(query: str) -> str:
    """Entry point for the HITL chat."""
    initial_state = {"messages": [HumanMessage(content=query)]}
    result = app.invoke(initial_state)
    return result["messages"][-1].content