Spaces:

ColettoGS
/

zico-agent

Sleeping

github-actions[bot]
Deploy from GitHub Actions: 9f340581b793cd99a67574e021bd63e1e5fbdc84
12cf769 about 2 months ago
57.8 kB
import asyncio
import hashlib
import base64
import json
import os
import time
from typing import Any, Dict, List, Literal, Optional

from fastapi import Depends, FastAPI, File, Form, HTTPException, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from langchain_core.messages import HumanMessage
from pydantic import BaseModel

from src.infrastructure.logging import setup_logging, get_logger
from src.infrastructure.rate_limiter import setup_rate_limiter, limiter
from src.agents.config import Config
from src.graphs.factory import build_graph
from src.graphs.nodes import initialize_agents
from src.graphs.utils import get_text_content
from src.models.chatMessage import ChatMessage
from src.routes.chat_manager_routes import router as chat_manager_router
from src.diagnostics.routes import router as diagnostics_router
from src.service.chat_manager import StoragePersistenceError, chat_manager_instance
from src.security.chat_auth import PanoramaPrincipal, require_chat_principal
from src.agents.crypto_data.tools import get_coingecko_id, get_tradingview_symbol
from src.agents.metadata import metadata
from src.diagnostics.response_boundary import start_trace, update_trace

# Setup structured logging
log_level = os.getenv("LOG_LEVEL", "INFO")
log_format = os.getenv("LOG_FORMAT", "color")
setup_logging(level=log_level, format_type=log_format)
logger = get_logger(__name__)

logger.info("Starting Zico Agent API (StateGraph architecture)")

# Initialize FastAPI app
app = FastAPI(
    title="Zico Agent API",
    version="3.0",
    description="Multi-agent AI assistant with deterministic StateGraph routing",
)

# Setup rate limiting
setup_rate_limiter(app)

# Enable CORS for local/frontend dev
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize agents and compile the StateGraph
initialize_agents()
graph = build_graph()


class ChatRequest(BaseModel):
    message: ChatMessage
    chain_id: str = "default"
    wallet_address: str = "default"
    conversation_id: str = "default"
    user_id: str = "anonymous"
    metadata: Dict[str, Any] | None = None
    response_mode: Literal["fast", "reasoning"] = "fast"


# Lightweight in-memory agent config for frontend integrations
AVAILABLE_AGENTS = [
    {"name": "default", "human_readable_name": "Default General Purpose", "description": "General chat and meta-queries about agents."},
    {"name": "crypto data", "human_readable_name": "Crypto Data Fetcher", "description": "Real-time cryptocurrency prices, market cap, FDV, TVL."},
    {"name": "token swap", "human_readable_name": "Token Swap Agent", "description": "Swap tokens using supported DEX APIs."},
    {"name": "realtime search", "human_readable_name": "Real-Time Search", "description": "Search the web for recent information."},
    {"name": "dexscreener", "human_readable_name": "DexScreener Analyst", "description": "Fetches and analyzes DEX trading data."},
    {"name": "rugcheck", "human_readable_name": "Token Safety Analyzer", "description": "Analyzes token safety and trends (Solana)."},
    {"name": "imagen", "human_readable_name": "Image Generator", "description": "Generate images from text prompts."},
    {"name": "rag", "human_readable_name": "Document Assistant", "description": "Answer questions about uploaded documents."},
    {"name": "tweet sizzler", "human_readable_name": "Tweet / X-Post Generator", "description": "Generate engaging tweets."},
    {"name": "dca", "human_readable_name": "DCA Strategy Manager", "description": "Plan and manage DCA strategies."},
    {"name": "base", "human_readable_name": "Base Transaction Manager", "description": "Handle transactions on Base network."},
    {"name": "mor rewards", "human_readable_name": "MOR Rewards Tracker", "description": "Track MOR rewards and balances."},
    {"name": "mor claims", "human_readable_name": "MOR Claims Agent", "description": "Claim MOR tokens."},
    {"name": "lending", "human_readable_name": "Lending Agent", "description": "Supply, borrow, repay, or withdraw assets."},
    {"name": "yield", "human_readable_name": "Yield / Liquidity Agent", "description": "Enter, exit, or claim rewards from LP pools on Base."},
]

# Default to a small, reasonable subset
SELECTED_AGENTS = [agent["name"] for agent in AVAILABLE_AGENTS[:6]]

# Commands exposed to the ChatInput autocomplete
AGENT_COMMANDS = [
    {"command": "morpheus", "name": "Default General Purpose", "description": "General assistant for simple queries and meta-questions."},
    {"command": "crypto", "name": "Crypto Data Fetcher", "description": "Get prices, market cap, FDV, TVL and more."},
    {"command": "document", "name": "Document Assistant", "description": "Ask questions about uploaded documents."},
    {"command": "tweet", "name": "Tweet / X-Post Generator", "description": "Create engaging tweets about crypto and web3."},
    {"command": "search", "name": "Real-Time Search", "description": "Search the web for recent events or updates."},
    {"command": "dexscreener", "name": "DexScreener Analyst", "description": "Analyze DEX trading data on supported chains."},
    {"command": "rugcheck", "name": "Token Safety Analyzer", "description": "Check token safety and view trending tokens."},
    {"command": "dca", "name": "DCA Strategy Manager", "description": "Plan a dollar-cost averaging strategy."},
    {"command": "base", "name": "Base Transaction Manager", "description": "Send tokens and swap on Base."},
    {"command": "rewards", "name": "MOR Rewards Tracker", "description": "Check rewards balance and accrual."},
    {"command": "lending", "name": "Lending Agent", "description": "Supply, borrow, repay, or withdraw assets."},
    {"command": "yield", "name": "Yield / Liquidity Agent", "description": "Manage LP positions and rewards on Base."},
]


# Agents endpoints expected by the frontend
@app.get("/agents/available")
def get_available_agents():
    return {
        "selected_agents": SELECTED_AGENTS,
        "available_agents": AVAILABLE_AGENTS,
    }


@app.post("/agents/selected")
async def set_selected_agents(request: Request):
    global SELECTED_AGENTS
    data = await request.json()
    agents = data.get("agents", [])
    available_names = {a["name"] for a in AVAILABLE_AGENTS}
    valid_agents = [a for a in agents if a in available_names]
    if not valid_agents:
        return {"status": "no_change", "agents": SELECTED_AGENTS}
    SELECTED_AGENTS = valid_agents[:6]
    return {"status": "success", "agents": SELECTED_AGENTS}


@app.get("/agents/commands")
def get_agent_commands():
    return {"commands": AGENT_COMMANDS}


# Map agent runtime names to high-level types for storage/analytics
def _map_agent_type(agent_name: str) -> str:
    mapping = {
        "crypto_agent": "crypto data",
        "default_agent": "default",
        "database_agent": "analysis",
        "search_agent": "realtime search",
        "swap_agent": "token swap",
        "lending_agent": "lending",
        "staking_agent": "staking",
        "liquidity_agent": "yield",
        "strategy_agent": "yield strategy",
        "portfolio_advisor": "portfolio analysis",
        "supervisor": "supervisor",
    }
    return mapping.get(agent_name, "supervisor")


def _sanitize_user_message_content(content: str | None) -> str | None:
    """Strip wrapper prompts (e.g., 'User Message: ...') the frontend might send."""
    if not content:
        return content
    text = content.strip()
    marker = "user message:"
    lowered = text.lower()
    idx = lowered.rfind(marker)
    if idx != -1:
        candidate = text[idx + len(marker) :].strip()
        if candidate:
            return candidate
    return text


def _validate_non_empty_user_content(content: str | None, *, user_id: str, conversation_id: str) -> str:
    cleaned = (content or "").strip()
    if cleaned:
        return cleaned
    logger.warning(
        "Rejecting empty user message before LLM invocation user=%s conversation=%s",
        user_id,
        conversation_id,
    )
    raise HTTPException(status_code=400, detail="message.content must be a non-empty string")


def _storage_http_exception(exc: StoragePersistenceError) -> HTTPException:
    logger.warning("Chat storage failure: %s", exc, exc_info=True)
    return HTTPException(
        status_code=503,
        detail="Chat persistence is currently unavailable. Your message was not stored.",
    )


def _resolve_identity(
    request: ChatRequest,
    principal: PanoramaPrincipal,
) -> tuple[str, str]:
    """Resolve authenticated ownership and the requested conversation ID."""
    user_id = principal.user_id
    conversation_id = (request.conversation_id or "").strip() or "default"
    return user_id, conversation_id


def _invoke_graph(
    conversation_messages,
    user_id,
    conversation_id,
    *,
    wallet_address: str | None = None,
    pre_classified: Dict[str, Any] | None = None,
    response_mode: str = "fast",
):
    """Invoke the StateGraph and return the result state.
    If *pre_classified* is provided (e.g. from audio transcription) its
    fields are merged into the initial state so the semantic router can
    skip the embedding call.
    """
    initial_state: Dict[str, Any] = {
        "messages": conversation_messages,
        "user_id": user_id,
        "conversation_id": conversation_id,
        "wallet_address": wallet_address,
        "response_mode": response_mode,
    }
    if pre_classified:
        initial_state.update(pre_classified)
    return graph.invoke(initial_state)


def _build_response_payload(result, user_id, conversation_id, extra_fields=None):
    """Build the HTTP response from graph result state."""
    final_response = result.get("final_response", "No response available")
    response_agent = result.get("response_agent", "supervisor")
    response_metadata = result.get("response_metadata", {})
    nodes_executed = result.get("nodes_executed", [])

    agent_name = _map_agent_type(response_agent)

    # Build response metadata and enrich
    full_metadata = {"supervisor_result": result}
    swap_meta_snapshot = None

    if response_metadata:
        full_metadata.update(response_metadata)
    elif agent_name == "token swap":
        swap_meta = metadata.get_swap_agent(user_id=user_id, conversation_id=conversation_id)
        if swap_meta:
            full_metadata.update(swap_meta)
            swap_meta_snapshot = swap_meta
    elif agent_name == "lending":
        lending_meta = metadata.get_lending_agent(user_id=user_id, conversation_id=conversation_id)
        if lending_meta:
            full_metadata.update(lending_meta)
    elif agent_name == "staking":
        staking_meta = metadata.get_staking_agent(user_id=user_id, conversation_id=conversation_id)
        if staking_meta:
            full_metadata.update(staking_meta)
    elif agent_name == "yield":
        liquidity_meta = metadata.get_liquidity_agent(user_id=user_id, conversation_id=conversation_id)
        if liquidity_meta:
            full_metadata.update(liquidity_meta)
    elif agent_name == "yield strategy":
        strategy_meta = metadata.get_strategy_agent(user_id=user_id, conversation_id=conversation_id)
        if strategy_meta:
            full_metadata.update(strategy_meta)

    # Create and store the response message
    response_message = ChatMessage(
        role="assistant",
        content=final_response,
        agent_name=agent_name,
        agent_type=_map_agent_type(agent_name),
        metadata=response_metadata,
        conversation_id=conversation_id,
        user_id=user_id,
        requires_action=True if agent_name in ["token swap", "lending", "staking", "yield", "yield strategy"] else False,
        action_type="swap" if agent_name == "token swap" else "lending" if agent_name == "lending" else "staking" if agent_name == "staking" else "liquidity" if agent_name == "yield" else "strategy" if agent_name == "yield strategy" else None,
    )

    chat_manager_instance.add_message(
        message=response_message.dict(),
        conversation_id=conversation_id,
        user_id=user_id,
    )

    # Build payload
    response_payload = {
        "response": final_response,
        "agentName": agent_name,
        "nodesExecuted": nodes_executed,
    }

    # Add extra fields (e.g. transcription for audio)
    if extra_fields:
        response_payload.update(extra_fields)

    # Resolve metadata for payload
    response_meta = response_metadata or {}
    if agent_name == "token swap" and not response_meta:
        if swap_meta_snapshot:
            response_meta = swap_meta_snapshot
        else:
            swap_meta = metadata.get_swap_agent(user_id=user_id, conversation_id=conversation_id)
            if swap_meta:
                response_meta = swap_meta
    elif agent_name == "yield" and not response_meta:
        liquidity_meta = metadata.get_liquidity_agent(user_id=user_id, conversation_id=conversation_id)
        if liquidity_meta:
            response_meta = liquidity_meta

    if response_meta:
        response_payload["metadata"] = response_meta

    # Clear metadata after ready events
    _clear_ready_metadata(agent_name, response_meta, user_id, conversation_id)

    return response_payload


def _clear_ready_metadata(agent_name, response_meta, user_id, conversation_id):
    """Clear DeFi metadata when intent is ready for execution."""
    if not response_meta or not isinstance(response_meta, dict):
        return

    status = response_meta.get("status")
    event = response_meta.get("event")

    if agent_name == "token swap" and (status == "ready" or event == "swap_intent_ready"):
        metadata.set_swap_agent({}, user_id=user_id, conversation_id=conversation_id)
    elif agent_name == "lending" and (status == "ready" or event == "lending_intent_ready"):
        metadata.set_lending_agent({}, user_id=user_id, conversation_id=conversation_id)
    elif agent_name == "staking" and (status == "ready" or event == "staking_intent_ready"):
        metadata.set_staking_agent({}, user_id=user_id, conversation_id=conversation_id)
    elif agent_name == "yield" and (status == "ready" or event == "liquidity_intent_ready"):
        metadata.set_liquidity_agent({}, user_id=user_id, conversation_id=conversation_id)
    elif agent_name == "yield strategy" and (status == "ready" or event == "strategy_intent_ready"):
        metadata.set_strategy_agent({}, user_id=user_id, conversation_id=conversation_id)


@app.get("/health")
def health_check():
    return {"status": "ok"}


@app.get("/costs")
def get_costs():
    """Get current LLM cost summary."""
    cost_tracker = Config.get_cost_tracker()
    return cost_tracker.get_summary()


@app.get("/costs/detailed")
def get_detailed_costs():
    """Get detailed LLM cost report."""
    cost_tracker = Config.get_cost_tracker()
    return cost_tracker.get_detailed_report()


@app.get("/costs/conversation")
def get_conversation_costs(request: Request):
    """Get accumulated LLM costs for a specific conversation."""
    params = request.query_params
    conversation_id = params.get("conversation_id")
    user_id = params.get("user_id")

    if not conversation_id or not user_id:
        raise HTTPException(
            status_code=400,
            detail="Both 'conversation_id' and 'user_id' query parameters are required.",
        )

    costs = chat_manager_instance.get_conversation_costs(
        conversation_id=conversation_id,
        user_id=user_id,
    )
    return {
        "conversation_id": conversation_id,
        "user_id": user_id,
        "costs": costs,
    }


@app.get("/models")
def get_available_models():
    """List available LLM models."""
    return {
        "models": Config.list_available_models(),
        "providers": Config.list_available_providers(),
        "default": Config.DEFAULT_MODEL,
    }


@app.post("/chat")
def chat(
    request: ChatRequest,
    principal: PanoramaPrincipal = Depends(require_chat_principal),
):
    user_id: str | None = None
    conversation_id: str | None = None
    try:
        logger.debug("Received chat payload: %s", request.model_dump())
        user_id, conversation_id = _resolve_identity(
            request,
            principal,
        )
        logger.debug(
            "Resolved chat identity user=%s conversation=%s wallet=%s",
            user_id,
            conversation_id,
            (request.wallet_address or "").strip() if request.wallet_address else None,
        )
        logger.info(
            "chat.identity_resolved user_id=%s conversation_id=%s source=chat_endpoint",
            user_id,
            conversation_id,
        )

        wallet = request.wallet_address.strip() if request.wallet_address else None
        if wallet and wallet.lower() == "default":
            wallet = None
        display_name = None
        if isinstance(request.message.metadata, dict):
            display_name = request.message.metadata.get("display_name")

        try:
            chat_manager_instance.ensure_session(
                user_id,
                conversation_id,
                wallet_address=wallet,
                display_name=display_name,
            )
        except StoragePersistenceError as exc:
            raise _storage_http_exception(exc) from exc

        if request.message.role == "user":
            clean_content = _sanitize_user_message_content(request.message.content)
            request.message.content = _validate_non_empty_user_content(
                clean_content,
                user_id=user_id,
                conversation_id=conversation_id,
            )

        # Add the user message to the conversation
        try:
            chat_manager_instance.add_message(
                message=request.message.dict(),
                conversation_id=conversation_id,
                user_id=user_id,
            )
        except StoragePersistenceError as exc:
            raise _storage_http_exception(exc) from exc

        # Get all messages from the conversation
        conversation_messages = chat_manager_instance.get_messages(
            conversation_id=conversation_id,
            user_id=user_id,
        )

        # Take cost snapshot before invoking
        cost_tracker = Config.get_cost_tracker()
        cost_snapshot = cost_tracker.get_snapshot()

        # Invoke the StateGraph
        result = _invoke_graph(
            conversation_messages, user_id, conversation_id,
            wallet_address=wallet, response_mode=request.response_mode,
        )

        # Calculate and save cost delta
        cost_delta = cost_tracker.calculate_delta(cost_snapshot)
        if cost_delta.get("cost", 0) > 0 or cost_delta.get("calls", 0) > 0:
            chat_manager_instance.update_conversation_costs(
                cost_delta,
                conversation_id=conversation_id,
                user_id=user_id,
            )

        logger.debug(
            "Graph returned result for user=%s conversation=%s nodes=%s",
            user_id,
            conversation_id,
            result.get("nodes_executed", []),
        )

        if result:
            return _build_response_payload(result, user_id, conversation_id)

        return {"response": "No response available", "agentName": "supervisor"}
    except StoragePersistenceError as exc:
        raise _storage_http_exception(exc) from exc
    except HTTPException:
        raise
    except Exception as e:
        logger.exception(
            "Chat handler failed for user=%s conversation=%s",
            user_id,
            conversation_id,
        )
        raise HTTPException(status_code=500, detail=str(e))


# ---------------------------------------------------------------------------
# SSE Streaming endpoint
# ---------------------------------------------------------------------------

# Human-readable labels for each graph node
_NODE_LABELS: Dict[str, str] = {
    "entry_node": "Preparing context...",
    "semantic_router_node": "Routing your request...",
    "llm_router_node": "Analyzing intent...",
    "swap_agent_node": "Consulting swap protocols...",
    "lending_agent_node": "Checking lending markets...",
    "staking_agent_node": "Reviewing staking options...",
    "liquidity_agent_node": "Preparing liquidity flow...",
    "dca_agent_node": "Planning DCA strategy...",
    "strategy_agent_node": "Designing AVAX strategy...",
    "crypto_agent_node": "Fetching market data...",
    "search_agent_node": "Searching the web...",
    "default_agent_node": "Thinking...",
    "database_agent_node": "Querying portfolio...",
    "portfolio_advisor_node": "Analyzing your portfolio...",
    "formatter_node": "Formatting response...",
    "error_node": "Validating parameters...",
}


def _sse(event_type: str, data: dict) -> str:
    """Format a Server-Sent Event string."""
    return f"event: {event_type}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"


async def _persist_response_bg(
    full_response: str,
    response_agent: str,
    response_metadata: dict,
    user_id: str,
    conversation_id: str,
    cost_delta: dict,
) -> None:
    """Durably persist the streamed assistant response before acknowledgement."""
    agent_name = _map_agent_type(response_agent)
    response_message = ChatMessage(
        role="assistant",
        content=full_response,
        agent_name=agent_name,
        agent_type=_map_agent_type(agent_name),
        metadata=response_metadata,
        conversation_id=conversation_id,
        user_id=user_id,
        requires_action=(
            True if agent_name in ("token swap", "lending", "staking", "yield", "yield strategy") else False
        ),
        action_type=(
            "swap"
            if agent_name == "token swap"
            else "lending"
            if agent_name == "lending"
            else "staking"
            if agent_name == "staking"
            else "liquidity"
            if agent_name == "yield"
            else "strategy"
            if agent_name == "yield strategy"
            else None
        ),
    )

    # Assistant-message durability is part of successful chat completion.
    # Let this failure propagate to the stream so it cannot emit a false
    # successful completion acknowledgement.
    await asyncio.to_thread(
        chat_manager_instance.add_message,
        response_message.dict(),
        conversation_id,
        user_id,
    )

    # Cost accounting is secondary to durable chat history.
    if cost_delta.get("cost", 0) > 0 or cost_delta.get("calls", 0) > 0:
        try:
            await asyncio.to_thread(
                chat_manager_instance.update_conversation_costs,
                cost_delta,
                conversation_id,
                user_id,
            )
        except Exception:
            logger.exception(
                "Failed to update costs for streamed response "
                "user=%s conversation=%s",
                user_id,
                conversation_id,
            )

    # Metadata cleanup is also secondary once the response is durable.
    try:
        _clear_ready_metadata(
            agent_name,
            response_metadata,
            user_id,
            conversation_id,
        )
    except Exception:
        logger.exception(
            "Failed to clear metadata after streamed response "
            "user=%s conversation=%s",
            user_id,
            conversation_id,
        )


def _build_event_generator(
    initial_state: Dict[str, Any],
    user_id: str,
    conversation_id: str,
    response_mode: str,
):
    """Factory that returns an async generator yielding SSE events from the graph."""

    async def event_generator():
        start_trace(user_id, conversation_id)
        cost_tracker = Config.get_cost_tracker()
        cost_snapshot = cost_tracker.get_snapshot()

        final_response_chunks: List[str] = []
        authoritative_graph_response = ""
        model_runs: Dict[str, Dict[str, Any]] = {}
        model_run_order: List[str] = []
        response_agent = "supervisor"
        response_metadata: Dict[str, Any] = {}
        nodes_executed: List[str] = []
        current_agent_node: str | None = None
        streaming_tokens = False

        try:
            async for event in graph.astream_events(
                initial_state, version="v2"
            ):
                kind = event["event"]
                name = event.get("name", "")

                if kind == "on_chain_start" and name in _NODE_LABELS:
                    nodes_executed.append(name)
                    if name.endswith("_agent_node"):
                        current_agent_node = name
                    yield _sse("status", {
                        "step": name,
                        "label": _NODE_LABELS[name],
                        "ts": time.time(),
                    })

                elif kind == "on_chain_end" and name == "semantic_router_node":
                    output = event.get("data", {}).get("output", {})
                    if isinstance(output, dict):
                        yield _sse("status", {
                            "step": "routed",
                            "agent": output.get("route_agent", "unknown"),
                            "confidence": output.get("route_confidence", 0),
                            "ts": time.time(),
                        })

                elif kind == "on_tool_start":
                    yield _sse("status", {
                        "step": "tool",
                        "tool": name,
                        "label": f"Using {name}...",
                        "ts": time.time(),
                    })

                elif kind == "on_tool_end":
                    tool_output = event.get("data", {}).get("output", "")
                    preview = str(tool_output)[:200]
                    yield _sse("tool_io", {
                        "tool": name,
                        "output": preview,
                        "ts": time.time(),
                    })

                elif kind == "on_chat_model_stream":
                    chunk = event.get("data", {}).get("chunk")
                    if chunk and hasattr(chunk, "content") and chunk.content:
                        text = chunk.content if isinstance(chunk.content, str) else ""
                        if text:
                            parent_tags = event.get("tags", [])
                            is_formatter = "formatter" in name.lower() or any(
                                "formatter" in t for t in parent_tags
                            )
                            if not is_formatter and current_agent_node:
                                run_id = str(event.get("run_id") or "")
                                if run_id:
                                    if run_id not in model_runs:
                                        model_runs[run_id] = {
                                            "sequence": len(model_run_order) + 1,
                                            "name": name,
                                            "parent_depth": len(event.get("parent_ids", []) or []),
                                            "chunks": [],
                                            "end_length": None,
                                            "end_sha256_16": None,
                                            "tool_call_count": None,
                                        }
                                        model_run_order.append(run_id)
                                    model_runs[run_id]["chunks"].append(text)

                                if not streaming_tokens:
                                    streaming_tokens = True
                                    yield _sse("status", {
                                        "step": "generating",
                                        "label": "Generating response...",
                                        "ts": time.time(),
                                    })
                                final_response_chunks.append(text)
                                yield _sse("token", {"t": text})

                elif kind == "on_chat_model_end":
                    run_id = str(event.get("run_id") or "")
                    if run_id and run_id in model_runs:
                        output = event.get("data", {}).get("output")
                        end_text = get_text_content(output) or ""
                        tool_calls = getattr(output, "tool_calls", None)
                        model_runs[run_id]["end_length"] = len(end_text)
                        model_runs[run_id]["end_sha256_16"] = (
                            hashlib.sha256(end_text.encode("utf-8")).hexdigest()[:16]
                            if end_text
                            else None
                        )
                        model_runs[run_id]["tool_call_count"] = (
                            len(tool_calls) if isinstance(tool_calls, list) else 0
                        )

                elif kind == "on_chain_end" and name == "LangGraph":
                    output = event.get("data", {}).get("output", {})
                    if isinstance(output, dict):
                        response_agent = output.get("response_agent", response_agent)
                        response_metadata = output.get("response_metadata", response_metadata)
                        graph_response = output.get("final_response", "")
                        if isinstance(graph_response, str) and graph_response:
                            authoritative_graph_response = graph_response

        except Exception as exc:
            logger.exception("Stream error for user=%s conversation=%s", user_id, conversation_id)
            yield _sse("error", {"message": str(exc)})
            return

        # Token callbacks are presentation-only. They are not guaranteed to
        # represent a complete user-facing response across nested LangGraph /
        # agent model calls. The completed graph state is authoritative for
        # durability and the final completion acknowledgement.
        streamed_response = "".join(final_response_chunks)
        full_response = authoritative_graph_response or streamed_response

        def _digest(value: str) -> str:
            return hashlib.sha256(value.encode("utf-8")).hexdigest()[:16]

        logger.info(
            "response.boundary.stream_complete streamed_length=%d streamed_sha256_16=%s "
            "graph_length=%d graph_sha256_16=%s selected_length=%d selected_sha256_16=%s",
            len(streamed_response),
            _digest(streamed_response),
            len(authoritative_graph_response),
            _digest(authoritative_graph_response),
            len(full_response),
            _digest(full_response),
        )
        model_run_evidence = []
        for run_id in model_run_order:
            run = model_runs[run_id]
            run_streamed = "".join(run["chunks"])
            model_run_evidence.append({
                "sequence": run["sequence"],
                "name": run["name"],
                "parent_depth": run["parent_depth"],
                "streamed_length": len(run_streamed),
                "streamed_sha256_16": _digest(run_streamed) if run_streamed else None,
                "end_length": run["end_length"],
                "end_sha256_16": run["end_sha256_16"],
                "tool_call_count": run["tool_call_count"],
            })

        update_trace(
            user_id,
            conversation_id,
            section="stream",
            value={
                "streamed_length": len(streamed_response),
                "streamed_sha256_16": _digest(streamed_response),
                "graph_length": len(authoritative_graph_response),
                "graph_sha256_16": _digest(authoritative_graph_response),
                "selected_length": len(full_response),
                "selected_sha256_16": _digest(full_response),
                "model_run_count": len(model_run_evidence),
                "model_runs": model_run_evidence,
            },
        )

        cost_delta = cost_tracker.calculate_delta(cost_snapshot)
        agent_name = _map_agent_type(response_agent)

        if not response_metadata:
            if agent_name == "token swap":
                swap_meta = metadata.get_swap_agent(user_id=user_id, conversation_id=conversation_id)
                if swap_meta:
                    response_metadata = swap_meta
            elif agent_name == "lending":
                lending_meta = metadata.get_lending_agent(user_id=user_id, conversation_id=conversation_id)
                if lending_meta:
                    response_metadata = lending_meta
            elif agent_name == "staking":
                staking_meta = metadata.get_staking_agent(user_id=user_id, conversation_id=conversation_id)
                if staking_meta:
                    response_metadata = staking_meta
            elif agent_name == "yield":
                liquidity_meta = metadata.get_liquidity_agent(user_id=user_id, conversation_id=conversation_id)
                if liquidity_meta:
                    response_metadata = liquidity_meta
            elif agent_name == "yield strategy":
                strategy_meta = metadata.get_strategy_agent(user_id=user_id, conversation_id=conversation_id)
                if strategy_meta:
                    response_metadata = strategy_meta

        try:
            await _persist_response_bg(
                full_response,
                response_agent,
                response_metadata,
                user_id,
                conversation_id,
                cost_delta,
            )
        except Exception:
            logger.exception(
                "Failed to persist streamed response "
                "user=%s conversation=%s",
                user_id,
                conversation_id,
            )
            yield _sse(
                "error",
                {
                    "message": (
                        "The assistant response could not be durably "
                        "saved. Please retry."
                    )
                },
            )
            return

        yield _sse("done", {
            "agent": agent_name,
            "nodes": nodes_executed,
            "metadata": response_metadata,
            "response": full_response,
            "response_mode": response_mode,
            "costs": {"total_usd": cost_delta.get("cost", 0)},
        })

    return event_generator


def _streaming_response(
    initial_state: Dict[str, Any],
    user_id: str,
    conversation_id: str,
    response_mode: str,
) -> StreamingResponse:
    """Build a StreamingResponse from the shared event generator factory."""
    gen = _build_event_generator(initial_state, user_id, conversation_id, response_mode)
    return StreamingResponse(
        gen(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


@app.post("/chat/stream")
async def chat_stream(
    request: ChatRequest,
    principal: PanoramaPrincipal = Depends(require_chat_principal),
):
    """SSE streaming endpoint — streams thought process + tokens in real-time.
    Event types:
      - ``status``  : node lifecycle (step label, routing info)
      - ``token``   : incremental text chunks from the final LLM response
      - ``tool_io`` : tool invocation results (truncated for the wire)
      - ``done``    : final metadata envelope (agent, nodes, costs)
      - ``error``   : unrecoverable error
    """
    uid: str | None = None
    cid: str | None = None
    try:
        uid, cid = _resolve_identity(
            request,
            principal,
        )
    except HTTPException as exc:
        # Return error as a streaming event so the client can parse it
        async def _err():
            yield _sse("error", {"message": exc.detail})

        return StreamingResponse(_err(), media_type="text/event-stream")

    user_id, conversation_id = uid, cid
    logger.info(
        "chat.identity_resolved user_id=%s conversation_id=%s source=chat_stream",
        user_id,
        conversation_id,
    )

    wallet = request.wallet_address.strip() if request.wallet_address else None
    if wallet and wallet.lower() == "default":
        wallet = None
    display_name = None
    if isinstance(request.message.metadata, dict):
        display_name = request.message.metadata.get("display_name")

    # Session setup, message persistence, and history fetch (non-blocking)
    try:
        await asyncio.to_thread(
            chat_manager_instance.ensure_session,
            user_id,
            conversation_id,
            wallet_address=wallet,
            display_name=display_name,
        )
    except StoragePersistenceError as exc:
        error_detail = _storage_http_exception(exc).detail
        async def _err():
            yield _sse("error", {"message": error_detail})

        return StreamingResponse(_err(), media_type="text/event-stream")

    if request.message.role == "user":
        clean_content = _sanitize_user_message_content(request.message.content)
        try:
            request.message.content = _validate_non_empty_user_content(
                clean_content,
                user_id=user_id,
                conversation_id=conversation_id,
            )
        except HTTPException as exc:
            async def _err():
                yield _sse("error", {"message": exc.detail})

            return StreamingResponse(_err(), media_type="text/event-stream")

    try:
        await asyncio.to_thread(
            chat_manager_instance.add_message,
            request.message.dict(),
            conversation_id,
            user_id,
        )
    except StoragePersistenceError as exc:
        error_detail = _storage_http_exception(exc).detail
        async def _err():
            yield _sse("error", {"message": error_detail})

        return StreamingResponse(_err(), media_type="text/event-stream")

    conversation_messages = await asyncio.to_thread(
        chat_manager_instance.get_messages,
        conversation_id,
        user_id,
    )

    initial_state: Dict[str, Any] = {
        "messages": conversation_messages,
        "user_id": user_id,
        "conversation_id": conversation_id,
        "wallet_address": wallet,
        "response_mode": request.response_mode,
    }

    return _streaming_response(initial_state, user_id, conversation_id, request.response_mode)


# ---------------------------------------------------------------------------
# Streaming chat with file attachments (images + documents)
# ---------------------------------------------------------------------------

@app.post("/chat/stream-with-files")
async def chat_stream_with_files(
    message: str = Form(...),
    user_id: str = Form(...),
    conversation_id: str = Form("default"),
    wallet_address: str = Form("default"),
    response_mode: Literal["fast", "reasoning"] = Form("fast"),
    files: List[UploadFile] = File(default=[]),
    principal: PanoramaPrincipal = Depends(require_chat_principal),
):
    """SSE streaming with file attachments (images & documents).
    Accepts ``multipart/form-data`` with one or more files.
    Images are encoded as base64 for Gemini multimodal content.
    Documents (PDF, txt, md, csv) are text-extracted and injected as context.
    """
    from src.service.file_processing import (
        get_file_mime_type, is_image, is_document,
        encode_image_for_gemini, extract_text_from_document,
        MAX_IMAGE_SIZE, MAX_DOCUMENT_SIZE,
    )

    # Ownership comes exclusively from the authenticated Panorama principal.
    uid = principal.user_id
    cid = (conversation_id or "").strip() or "default"
    wallet = wallet_address.strip() if wallet_address else None
    if wallet and wallet.lower() == "default":
        wallet = None
    logger.info(
        "chat.identity_resolved user_id=%s conversation_id=%s source=chat_stream_with_files",
        uid,
        cid,
    )

    clean_message = _sanitize_user_message_content(message)
    try:
        message = _validate_non_empty_user_content(
            clean_message,
            user_id=uid,
            conversation_id=cid,
        )
    except HTTPException as exc:
        async def _err():
            yield _sse("error", {"message": exc.detail})

        return StreamingResponse(_err(), media_type="text/event-stream")

    # --- Process file attachments ---
    file_attachments: List[Dict[str, Any]] = []

    for f in files:
        content = await f.read()
        mime = get_file_mime_type(f.filename or "", f.content_type)

        if is_image(mime):
            if len(content) > MAX_IMAGE_SIZE:
                async def _err():
                    yield _sse("error", {"message": f"Image '{f.filename}' exceeds {MAX_IMAGE_SIZE // (1024*1024)}MB limit."})
                return StreamingResponse(_err(), media_type="text/event-stream")
            encoded = encode_image_for_gemini(content, mime)
            file_attachments.append({
                "type": "image",
                "filename": f.filename,
                "mime_type": mime,
                "data": encoded["data"],
            })
        elif is_document(mime):
            if len(content) > MAX_DOCUMENT_SIZE:
                async def _err():
                    yield _sse("error", {"message": f"Document '{f.filename}' exceeds {MAX_DOCUMENT_SIZE // (1024*1024)}MB limit."})
                return StreamingResponse(_err(), media_type="text/event-stream")
            try:
                text = await asyncio.to_thread(extract_text_from_document, content, mime)
                logger.info("Extracted %d chars from document %s (%s)", len(text), f.filename, mime)
            except Exception as exc:
                logger.exception("Failed to extract text from %s (%s): %s", f.filename, mime, exc)
                text = ""
            file_attachments.append({
                "type": "document",
                "filename": f.filename,
                "mime_type": mime,
                "text": text,
            })
        else:
            async def _err():
                yield _sse("error", {"message": f"Unsupported file type: {f.filename} ({mime or 'unknown'})"})
            return StreamingResponse(_err(), media_type="text/event-stream")

    # --- Session setup + persist user message ---
    try:
        await asyncio.to_thread(
            chat_manager_instance.ensure_session,
            uid, cid, wallet_address=wallet,
        )
    except StoragePersistenceError as exc:
        error_detail = _storage_http_exception(exc).detail
        async def _err():
            yield _sse("error", {"message": error_detail})

        return StreamingResponse(_err(), media_type="text/event-stream")

    user_msg = ChatMessage(
        role="user",
        content=message,
        metadata={"source": "miniapp-chat", "has_attachments": True, "attachment_count": len(file_attachments)},
    )
    try:
        await asyncio.to_thread(
            chat_manager_instance.add_message, user_msg.dict(), cid, uid,
        )
    except StoragePersistenceError as exc:
        error_detail = _storage_http_exception(exc).detail
        async def _err():
            yield _sse("error", {"message": error_detail})

        return StreamingResponse(_err(), media_type="text/event-stream")

    conversation_messages = await asyncio.to_thread(
        chat_manager_instance.get_messages, cid, uid,
    )

    initial_state: Dict[str, Any] = {
        "messages": conversation_messages,
        "user_id": uid,
        "conversation_id": cid,
        "wallet_address": wallet,
        "response_mode": response_mode,
        "file_attachments": file_attachments,
    }

    return _streaming_response(initial_state, uid, cid, response_mode)


# Supported audio MIME types
AUDIO_MIME_TYPES = {
    ".mp3": "audio/mpeg",
    ".wav": "audio/wav",
    ".flac": "audio/flac",
    ".ogg": "audio/ogg",
    ".webm": "audio/webm",
    ".m4a": "audio/mp4",
    ".aac": "audio/aac",
}

MAX_AUDIO_SIZE = 20 * 1024 * 1024


def _get_audio_mime_type(filename: str, content_type: str | None) -> str:
    """Determine the MIME type for an audio file."""
    if filename:
        ext = os.path.splitext(filename.lower())[1]
        if ext in AUDIO_MIME_TYPES:
            return AUDIO_MIME_TYPES[ext]
    if content_type and content_type.startswith("audio/"):
        return content_type
    return "audio/mpeg"


# ---------------------------------------------------------------------------
# Audio: combined transcription + intent classification prompt
# ---------------------------------------------------------------------------

_AUDIO_TRANSCRIBE_AND_CLASSIFY_PROMPT = """\
You will receive an audio clip. Perform TWO tasks:
1. **Transcribe** exactly what is being said.
2. **Classify** the user's intent into one of these categories:
   swap, lending, staking, liquidity, dca, strategy, market_data, search, education, general
Return ONLY a JSON object (no markdown fences) with these fields:
{"transcription": "<exact transcription>", "intent": "<category>", "confidence": <0.0-1.0>}
"""

_AUDIO_TRANSCRIBE_ONLY_PROMPT = """\
You will receive an audio clip. Transcribe exactly what is being said.
Return ONLY the transcription text, nothing else. No JSON, no markdown, no labels.
"""

# Maps audio classification intents to agent runtime names
_AUDIO_INTENT_AGENT_MAP: Dict[str, str] = {
    "swap": "swap_agent",
    "lending": "lending_agent",
    "staking": "staking_agent",
    "liquidity": "liquidity_agent",
    "dca": "dca_agent",
    "strategy": "strategy_agent",
    "market_data": "crypto_agent",
    "portfolio": "portfolio_advisor",
    "search": "search_agent",
    "education": "default_agent",
    "general": "default_agent",
}


def _parse_audio_classification(raw_content: str) -> tuple[str, str | None, float]:
    """Parse the combined transcription + classification JSON response.
    Returns ``(transcription, intent, confidence)``.  Falls back gracefully
    if the model doesn't return valid JSON — treats the entire response as
    plain transcription.
    """
    text = raw_content.strip()

    # Strip markdown code fences if present
    if text.startswith("```"):
        text = text.split("\n", 1)[-1]
    if text.endswith("```"):
        text = text.rsplit("```", 1)[0]
    text = text.strip()

    try:
        data = json.loads(text)
        transcription = (data.get("transcription") or "").strip()
        intent = (data.get("intent") or "").strip().lower()
        confidence = float(data.get("confidence", 0.0))

        if not transcription:
            # JSON parsed but no transcription field — use raw
            return raw_content.strip(), None, 0.0

        valid_intents = set(_AUDIO_INTENT_AGENT_MAP.keys())
        if intent not in valid_intents:
            intent = None
            confidence = 0.0

        return transcription, intent, confidence
    except (json.JSONDecodeError, ValueError, TypeError):
        # Not JSON — treat entire response as transcription
        return raw_content.strip(), None, 0.0


@app.post("/chat/audio")
async def chat_audio(
    audio: UploadFile = File(..., description="Audio file (mp3, wav, flac, ogg, webm, m4a)"),
    user_id: str = Form(..., description="User ID"),
    conversation_id: str = Form(..., description="Conversation ID"),
    wallet_address: str = Form("default", description="Wallet address"),
    principal: PanoramaPrincipal = Depends(require_chat_principal),
):
    """Process audio input through the agent pipeline.
    Optimisations over the naive sequential approach:
    1. Combined transcription + intent classification in a single LLM call
    2. Session setup + history fetch run in parallel with transcription
    3. All blocking calls run in a thread pool (asyncio.to_thread)
    4. Pre-classified intent is injected into graph state so semantic_router
       can skip the embedding call (~200 ms saved)
    """
    request_user_id: str = principal.user_id
    request_conversation_id: str = (
        (conversation_id or "").strip() or "default"
    )

    try:
        logger.debug(
            "Received audio chat request user=%s conversation=%s filename=%s",
            request_user_id,
            request_conversation_id,
            audio.filename,
        )

        # Validate file size
        audio_content = await audio.read()
        if len(audio_content) > MAX_AUDIO_SIZE:
            raise HTTPException(
                status_code=413,
                detail=f"Audio file too large. Maximum size is {MAX_AUDIO_SIZE // (1024 * 1024)}MB.",
            )

        if len(audio_content) == 0:
            raise HTTPException(status_code=400, detail="Audio file is empty.")

        mime_type = _get_audio_mime_type(audio.filename or "", audio.content_type)
        logger.debug("Audio MIME type: %s, size: %d bytes", mime_type, len(audio_content))

        encoded_audio = base64.b64encode(audio_content).decode("utf-8")

        wallet = wallet_address.strip() if wallet_address else None
        if wallet and wallet.lower() == "default":
            wallet = None

        # Take cost snapshot
        cost_tracker = Config.get_cost_tracker()
        cost_snapshot = cost_tracker.get_snapshot()

        # ── Parallel phase: transcription + session/history ──────────────
        # These are independent — run concurrently.

        from src.llm.tiers import ModelTier
        llm = Config.get_llm(model=ModelTier.TRANSCRIPTION, with_cost_tracking=True)

        transcription_message = HumanMessage(
            content=[
                {"type": "text", "text": _AUDIO_TRANSCRIBE_AND_CLASSIFY_PROMPT},
                {"type": "media", "data": encoded_audio, "mime_type": mime_type},
            ]
        )

        async def _transcribe():
            return await asyncio.to_thread(llm.invoke, [transcription_message])

        async def _setup_session_and_history():
            try:
                await asyncio.to_thread(
                    chat_manager_instance.ensure_session,
                    request_user_id,
                    request_conversation_id,
                    wallet_address=wallet,
                )
            except StoragePersistenceError as exc:
                raise _storage_http_exception(exc) from exc
            return await asyncio.to_thread(
                chat_manager_instance.get_messages,
                request_conversation_id,
                request_user_id,
            )

        transcription_response, conversation_messages = await asyncio.gather(
            _transcribe(),
            _setup_session_and_history(),
        )

        # ── Parse combined transcription + classification ────────────────

        raw_content = transcription_response.content
        if isinstance(raw_content, list):
            text_parts = []
            for part in raw_content:
                if isinstance(part, dict) and part.get("text"):
                    text_parts.append(part["text"])
                elif isinstance(part, str):
                    text_parts.append(part)
            raw_content = " ".join(text_parts).strip()

        transcribed_text, audio_intent, audio_confidence = _parse_audio_classification(
            raw_content or "",
        )

        if not transcribed_text:
            raise HTTPException(
                status_code=400,
                detail="Could not transcribe the audio. Please try again with a clearer recording.",
            )

        logger.info(
            "Audio transcribed: '%s' | intent=%s confidence=%.2f",
            transcribed_text[:200],
            audio_intent,
            audio_confidence,
        )

        # ── Store user message ───────────────────────────────────────────

        user_message = ChatMessage(
            role="user",
            content=transcribed_text,
            metadata={
                "source": "audio",
                "audio_filename": audio.filename,
                "audio_size": len(audio_content),
                "audio_mime_type": mime_type,
            },
        )
        try:
            await asyncio.to_thread(
                chat_manager_instance.add_message,
                user_message.dict(),
                request_conversation_id,
                request_user_id,
            )
        except StoragePersistenceError as exc:
            raise _storage_http_exception(exc) from exc

        # Re-fetch messages with the newly added user message
        conversation_messages = await asyncio.to_thread(
            chat_manager_instance.get_messages,
            request_conversation_id,
            request_user_id,
        )

        # ── Invoke graph with pre-classified intent ──────────────────────

        pre_classified: Dict[str, Any] | None = None
        if audio_intent and audio_confidence > 0.0:
            pre_classified = {
                "route_intent": audio_intent,
                "route_confidence": audio_confidence,
                "route_agent": _AUDIO_INTENT_AGENT_MAP.get(audio_intent, "default_agent"),
            }

        result = await asyncio.to_thread(
            _invoke_graph,
            conversation_messages,
            request_user_id,
            request_conversation_id,
            wallet_address=wallet,
            pre_classified=pre_classified,
            response_mode="fast",  # Audio path defaults to fast
        )

        # ── Cost tracking ────────────────────────────────────────────────

        cost_delta = cost_tracker.calculate_delta(cost_snapshot)
        if cost_delta.get("cost", 0) > 0 or cost_delta.get("calls", 0) > 0:
            chat_manager_instance.update_conversation_costs(
                cost_delta,
                conversation_id=request_conversation_id,
                user_id=request_user_id,
            )

        logger.debug(
            "Graph returned result for audio user=%s conversation=%s nodes=%s",
            request_user_id,
            request_conversation_id,
            result.get("nodes_executed", []),
        )

        if result:
            return _build_response_payload(
                result,
                request_user_id,
                request_conversation_id,
                extra_fields={"transcription": transcribed_text},
            )

        return {
            "response": "No response available",
            "agentName": "supervisor",
            "transcription": transcribed_text,
        }

    except StoragePersistenceError as exc:
        raise _storage_http_exception(exc) from exc
    except HTTPException:
        raise
    except Exception as e:
        logger.exception(
            "Audio chat handler failed for user=%s conversation=%s",
            request_user_id,
            request_conversation_id,
        )
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/transcribe")
async def transcribe_audio(
    audio: UploadFile = File(..., description="Audio file (mp3, wav, flac, ogg, webm, m4a)"),
):
    """Transcribe audio to text without invoking the agent pipeline.
    Stateless endpoint — no user_id, conversation_id, or persistence.
    Returns ``{"text": "<transcription>"}``.
    """
    try:
        audio_content = await audio.read()
        if len(audio_content) > MAX_AUDIO_SIZE:
            raise HTTPException(
                status_code=413,
                detail=f"Audio file too large. Maximum size is {MAX_AUDIO_SIZE // (1024 * 1024)}MB.",
            )
        if len(audio_content) == 0:
            raise HTTPException(status_code=400, detail="Audio file is empty.")

        mime_type = _get_audio_mime_type(audio.filename or "", audio.content_type)
        encoded_audio = base64.b64encode(audio_content).decode("utf-8")

        from src.llm.tiers import ModelTier
        llm = Config.get_llm(model=ModelTier.TRANSCRIPTION, with_cost_tracking=True)

        message = HumanMessage(
            content=[
                {"type": "text", "text": _AUDIO_TRANSCRIBE_ONLY_PROMPT},
                {"type": "media", "data": encoded_audio, "mime_type": mime_type},
            ]
        )

        response = await asyncio.to_thread(llm.invoke, [message])

        raw_content = response.content
        if isinstance(raw_content, list):
            text_parts = []
            for part in raw_content:
                if isinstance(part, dict) and part.get("text"):
                    text_parts.append(part["text"])
                elif isinstance(part, str):
                    text_parts.append(part)
            raw_content = " ".join(text_parts).strip()

        text = (raw_content or "").strip()
        if not text:
            raise HTTPException(
                status_code=400,
                detail="Could not transcribe the audio. Please try again with a clearer recording.",
            )

        return {"text": text}

    except HTTPException:
        raise
    except Exception as e:
        logger.exception("Transcribe endpoint failed")
        raise HTTPException(status_code=500, detail=str(e))


# Include chat manager router
app.include_router(chat_manager_router)
app.include_router(diagnostics_router)