agent/chat_agent.py

"""Bounded agentic loop for free-form Q&A over stored documents.

Reuses the same LangChain tools as the brief agent (search_filing, search_transcript,
get_financial_metrics, get_analyst_expectations) but WITHOUT search_news — answers are
grounded exclusively in ingested SEC filings, transcripts, and structured metrics.

Anti-hallucination layers:
  - temperature=0
  - System prompt bans training-data facts
  - Data-availability preamble prevents invented fiscal periods
  - Inline source citations (chunk_context headers from tool outputs)
  - sources list returned for UI transparency expander
"""
from __future__ import annotations

import datetime
from typing import Any

from langchain_core.messages import (
    AIMessage,
    HumanMessage,
    ToolMessage,
)

from agent.llm import RunConfig, default_config, make_chat_model, build_system_message
from agent.prompts import CHAT_SYSTEM_PROMPT
from agent.tools import (
    get_analyst_expectations,
    get_financial_metrics,
    search_filing,
    search_transcript,
)
from storage import metrics_db

# Chat tools — intentionally excludes search_news (Tavily live web search).
CHAT_TOOLS = [
    get_financial_metrics,
    get_analyst_expectations,
    search_filing,
    search_transcript,
]
_CHAT_TOOLS_BY_NAME: dict[str, Any] = {t.name: t for t in CHAT_TOOLS}

MAX_CHAT_ROUNDS = 5


# ── Helpers ───────────────────────────────────────────────────────────────────

def _build_data_preamble(ticker: str) -> str:
    """Return a plain-text summary of available ingested periods for the ticker.

    Injected at the top of every user message so the model cannot invent fiscal
    periods or misidentify "the most recent" filing.
    """
    rows = metrics_db.get_all_metrics(ticker)
    today = datetime.date.today().isoformat()
    if not rows:
        return (
            f"## Available data for {ticker}\n\n"
            "No data ingested yet. Run: python ingest.py "
            f"{ticker}\n\nToday's date: {today}"
        )

    lines = [
        f"## Available data for {ticker}",
        f"Today's date: {today}",
        "",
        "Ingested periods (most recent first):",
    ]
    for i, row in enumerate(rows):
        marker = "  ← MOST RECENT" if i == 0 else ""
        lines.append(
            f"  - {row['period']} | {row['form_type']} | Filed: {row['filing_date']}{marker}"
        )
    lines += [
        "",
        f"Most recent period: {rows[0]['period']} (filed {rows[0]['filing_date']})",
    ]
    return "\n".join(lines)


def _history_to_messages(history: list[dict]) -> list:
    """Convert simplified history dicts to LangChain message objects.

    Each entry must have ``role`` ("user" | "assistant") and ``content`` (str).
    The optional ``sources`` key on assistant entries is ignored here.
    """
    messages = []
    for entry in history:
        role = entry.get("role")
        content = entry.get("content", "")
        if role == "user":
            messages.append(HumanMessage(content=content))
        elif role == "assistant":
            messages.append(AIMessage(content=content))
    return messages


def _extract_text(content: Any) -> str:
    """Extract plain text from an Anthropic content value (str or list of blocks)."""
    if isinstance(content, str):
        return content.strip()
    if isinstance(content, list):
        parts = []
        for block in content:
            if isinstance(block, dict) and block.get("type") == "text":
                parts.append(block["text"])
            elif hasattr(block, "type") and block.type == "text":
                parts.append(block.text)
        return "\n".join(parts).strip()
    return str(content).strip()


# ── Public API ────────────────────────────────────────────────────────────────

def answer_question(
    ticker: str, question: str, history: list[dict], config: RunConfig | None = None
) -> dict:
    """Run a bounded agentic loop to answer a question about a ticker's documents.

    Args:
        ticker:   Upper-cased company ticker, already validated by the UI.
        question: The user's current question.
        history:  Previous Q&A turns as ``[{"role": ..., "content": ...}, ...]``.
                  Should NOT include the current question (it is passed separately).
        config:   Provider/model/key snapshot; defaults to Anthropic Haiku via env key.

    Returns:
        ``{"answer": str, "sources": list[dict]}``
        Each source dict has ``tool_name``, ``args``, and ``output`` (truncated to
        800 chars for readability in the UI expander).
    """
    ticker = ticker.upper()
    cfg = config or default_config()

    # System message with prompt caching (mirrors graph.py agent_node).
    system_msg = build_system_message(cfg, CHAT_SYSTEM_PROMPT.format(ticker=ticker))

    # Inject data-availability preamble so the model never invents periods.
    preamble = _build_data_preamble(ticker)
    full_question = f"{preamble}\n\n---\n\n{question}"

    llm_with_tools = make_chat_model(cfg).bind_tools(CHAT_TOOLS)

    messages: list = (
        [system_msg]
        + _history_to_messages(history)
        + [HumanMessage(content=full_question)]
    )

    sources: list[dict] = []

    for _round in range(MAX_CHAT_ROUNDS):
        response = llm_with_tools.invoke(messages)
        messages.append(response)

        tool_calls = getattr(response, "tool_calls", None) or []
        if not tool_calls:
            return {"answer": _extract_text(response.content), "sources": sources}

        # Execute tool calls sequentially (chat pace; no concurrency needed).
        for tc in tool_calls:
            name = tc["name"]
            args = dict(tc.get("args") or {})
            tool_call_id = tc.get("id", "")

            # Defensively inject ticker — all 4 tools require it.
            if "ticker" not in args:
                args["ticker"] = ticker

            try:
                fn = _CHAT_TOOLS_BY_NAME.get(name)
                if fn is None:
                    raise ValueError(f"Unknown tool: {name!r}")
                raw = fn.invoke(args)
                output_str = raw if isinstance(raw, str) else str(raw)
            except Exception as exc:
                output_str = f"Error: {exc}"

            # Capture for the UI sources expander (truncate for readability).
            sources.append({
                "tool_name": name,
                "args": args,
                "output": output_str[:800] + ("…" if len(output_str) > 800 else ""),
            })

            messages.append(ToolMessage(
                tool_call_id=tool_call_id,
                name=name,
                content=output_str,
            ))

    # Round cap reached — elicit a final answer without further tool calls.
    messages.append(HumanMessage(
        content=(
            "Round cap reached. Summarise your answer based solely on what "
            "you have retrieved so far. Do not issue any more tool calls."
        )
    ))
    final = make_chat_model(cfg).invoke(messages)
    return {"answer": _extract_text(final.content), "sources": sources}