agent/tools.py

import os
from typing import Any
from langchain_core.tools import tool
from agent.evidence import evidence_envelope, make_evidence_record, parse_evidence_envelope
from storage import metrics_db, news_cache

TOP_K_FILINGS = 5
TOP_K_TRANSCRIPTS = 4

_tavily_client: Any = None


class _LazyVectorStore:
    """Keep heavy vector dependencies lazy while preserving a patchable seam."""

    @staticmethod
    def search(*args, **kwargs):
        from storage import vector_store as implementation

        return implementation.search(*args, **kwargs)


vector_store = _LazyVectorStore()


def TavilyClient(*args, **kwargs):
    """Lazily construct Tavily's client; kept as a named dependency seam."""
    from tavily import TavilyClient as implementation

    return implementation(*args, **kwargs)


def _empty(tool_name: str, query: dict, message: str) -> str:
    return evidence_envelope(tool=tool_name, query=query, status="EMPTY", message=message)


def _error(tool_name: str, query: dict, exc: Exception) -> str:
    return evidence_envelope(
        tool=tool_name,
        query=query,
        status="ERROR",
        message=f"{type(exc).__name__}: {exc}",
        error_code="TOOL_EXECUTION_ERROR",
    )


def _vector_record(result: dict, fallback_source: str):
    metadata = dict(result.get("metadata") or {})
    source = metadata.get("source") or fallback_source
    metadata.setdefault(
        "chunk_context",
        f"Source: {source} | Section: {metadata.get('section', '')} | Date: {metadata.get('filing_date') or metadata.get('date') or ''}",
    )
    return make_evidence_record(
        source=source,
        content=result.get("text", ""),
        document_id=str(metadata.get("document_id") or ":".join(filter(None, [
            source,
            str(metadata.get("ticker", "")),
            str(metadata.get("period", "")),
            str(metadata.get("filing_date") or metadata.get("date") or ""),
            str(metadata.get("section", "")),
        ]))),
        chunk_id=str(metadata.get("chunk_id")) if metadata.get("chunk_id") else None,
        source_url=metadata.get("source_url"),
        as_of=metadata.get("filing_date") or metadata.get("date"),
        metadata=metadata,
    )


def _get_tavily_client() -> Any:
    global _tavily_client
    if _tavily_client is None:
        _tavily_client = TavilyClient(api_key=os.environ["TAVILY_API_KEY"])
    return _tavily_client


def _fmt_billions(v) -> str:
    return f"${v / 1e9:.1f}B" if v is not None else "N/A"


def _fmt_pct(v) -> str:
    return f"{v:.1%}" if v is not None else "N/A"


def _fmt_millions(v) -> str:
    return f"${v / 1e6:.0f}M" if v is not None else "N/A"


def _fmt_row(data: dict) -> str:
    yoy = f" (YoY: {data['revenue_yoy_pct']:+.1f}%)" if data.get("revenue_yoy_pct") is not None else ""
    lines = [
        f"  Filing: {data.get('form_type', '')} {data.get('period', '')} | Date: {data['filing_date']}",
        f"  Revenue: {_fmt_billions(data['revenue'])}{yoy}",
        f"  EPS (diluted): {data['eps']:.2f}" if data.get("eps") else "  EPS: N/A",
        f"  Gross Margin: {_fmt_pct(data['gross_margin'])}",
        f"  Operating Margin: {_fmt_pct(data['operating_margin'])}",
        f"  Free Cash Flow: {_fmt_billions(data['free_cash_flow'])}",
    ]
    # Valuation parameter inputs
    shares = data.get("shares_diluted")
    lines.append(f"  Shares Outstanding (diluted): {shares / 1e6:.0f}M" if shares else "  Shares Outstanding: N/A")
    lines.append(f"  Effective Tax Rate: {_fmt_pct(data.get('effective_tax_rate'))}")
    lines.append(f"  Interest Expense: {_fmt_millions(data.get('interest_expense'))}")
    lines.append(f"  Total Debt (LT): {_fmt_billions(data.get('total_debt'))}")
    lines.append(f"  Capex: {_fmt_millions(data.get('capex'))}")
    lines.append(f"  Dividends Paid: {_fmt_millions(data.get('dividends_paid'))}")
    lines.append(f"  Share Buybacks: {_fmt_millions(data.get('buybacks'))}")
    equity = data.get("stockholders_equity")
    debt = data.get("total_debt")
    de_ratio = f" (D/E: {debt / equity:.2f}x)" if (debt and equity and equity > 0) else ""
    lines.append(f"  Stockholders' Equity: {_fmt_billions(equity)}{de_ratio}")
    if data.get("guidance_text"):
        lines.append(f"  Guidance: {data['guidance_text']}")
    warnings = data.get("quality_warnings") or []
    contexts = data.get("metric_contexts") or {}
    lineage = []
    if isinstance(contexts, dict):
        for metric, context in sorted(contexts.items()):
            if not isinstance(context, dict):
                continue
            source = context.get("source") or "derived"
            selection = context.get("selection") or "unspecified"
            concept = context.get("concept")
            detail = f"{metric}={source}/{selection}"
            if concept:
                detail += f"/{concept}"
            lineage.append(detail)
    lines.extend([
        f"  Data Quality: {data.get('data_quality_status') or 'LEGACY_UNVERIFIED'}",
        f"  Accession: {data.get('accession') or 'N/A'} | Report Date: {data.get('report_date') or 'N/A'}",
        f"  Quality Warnings: {', '.join(str(item) for item in warnings) if warnings else 'none'}",
        f"  Metric Lineage: {'; '.join(lineage) if lineage else 'not recorded'}",
    ])
    return "\n".join(lines)


@tool
def get_financial_metrics(ticker: str) -> str:
    """Retrieve structured financial metrics for a ticker across all ingested periods.
    Always call this first to anchor quantitative claims. Returns an evidence.v1
    JSON envelope; cite a record by copying its `ref` into `evidence_ref`."""
    query = {"ticker": ticker.upper()}
    try:
        rows = metrics_db.get_all_metrics(ticker)
    except Exception as exc:
        return _error("get_financial_metrics", query, exc)
    if not rows:
        return _empty(
            "get_financial_metrics",
            query,
            f"Ticker {ticker.upper()} not ingested. Run: python ingest.py {ticker.upper()}",
        )

    usable_rows = [
        row for row in rows
        if row.get("data_quality_status") in {"VERIFIED", "CHECK_REQUIRED"}
    ]
    if not usable_rows:
        statuses = sorted({
            str(row.get("data_quality_status") or "LEGACY_UNVERIFIED")
            for row in rows
        })
        return _empty(
            "get_financial_metrics",
            query,
            (
                f"Metrics exist for {ticker.upper()}, but none carries SEC period lineage "
                f"(statuses: {', '.join(statuses)}). Run: python ingest.py "
                f"{ticker.upper()} --full to restore citable metrics."
            ),
        )

    try:
        records = []
        for row in usable_rows:
            content = f"Company: {row['company_name']} ({row['ticker']})\n{_fmt_row(row)}"
            period = str(row.get("period") or "unknown")
            filing_date = str(row.get("filing_date") or "")
            accession = str(row.get("accession") or "")
            records.append(make_evidence_record(
                source="metrics",
                content=content,
                document_id=(
                    f"metrics:{ticker.upper()}:{accession}"
                    if accession else f"metrics:{ticker.upper()}:{period}:{filing_date}"
                ),
                source_url=row.get("source_url") or None,
                as_of=filing_date or None,
                metadata={
                    "ticker": ticker.upper(),
                    "company_name": row.get("company_name"),
                    "period": period,
                    "filing_date": filing_date,
                    "form_type": row.get("form_type"),
                    "period_basis": row.get("period_basis", "unknown"),
                    "report_date": row.get("report_date"),
                    "accession": row.get("accession"),
                    "source_url": row.get("source_url"),
                    "metric_contexts": row.get("metric_contexts") or {},
                    "quality_warnings": row.get("quality_warnings") or [],
                    "data_quality_status": row.get("data_quality_status") or "LEGACY_UNVERIFIED",
                },
            ))
    except Exception as exc:
        return _error("get_financial_metrics", query, exc)
    return evidence_envelope(tool="get_financial_metrics", query=query, records=records)


@tool
def search_filing(query: str, ticker: str, since: str = "", period: str = "") -> str:
    """Search SEC filings (10-Q and 10-K) for relevant text using semantic search.
    Use for MD&A, risk factors, and outlook sections.
    Use `since` (YYYY-MM-DD) to restrict results to filings after that date.
    Use `period` (e.g. 'Q12024', 'FY2023') to fetch chunks from a specific past period for cross-year/cross-quarter comparison.
    Returns an evidence.v1 JSON envelope; copy the supporting record's `ref` exactly."""
    request = {"query": query, "ticker": ticker.upper(), "since": since, "period": period}
    try:
        results = vector_store.search(
            "filings", query, ticker, n_results=TOP_K_FILINGS,
            min_filing_date=since if since else None,
            period=period if period else None,
        )
    except Exception as exc:
        return _error("search_filing", request, exc)
    if not results:
        return _empty(
            "search_filing",
            request,
            f"No filing data for {ticker.upper()} matching query (period={period or 'any'}). Run: python ingest.py {ticker.upper()}",
        )
    try:
        records = [_vector_record(result, "10-Q") for result in results]
    except Exception as exc:
        return _error("search_filing", request, exc)
    return evidence_envelope(tool="search_filing", query=request, records=records)


@tool
def search_transcript(query: str, ticker: str, since: str = "", period: str = "") -> str:
    """Search earnings call transcripts for management commentary.
    Use for CEO/CFO tone, forward guidance, and analyst Q&A themes.
    Use `since` (YYYY-MM-DD) to restrict results to transcripts after that date.
    Use `period` (e.g. 'Q12024', 'FY2023') to fetch chunks from a specific past period for cross-year/cross-quarter comparison.
    Returns an evidence.v1 JSON envelope; copy the supporting record's `ref` exactly."""
    request = {"query": query, "ticker": ticker.upper(), "since": since, "period": period}
    try:
        results = vector_store.search(
            "transcripts", query, ticker, n_results=TOP_K_TRANSCRIPTS,
            min_filing_date=since if since else None,
            period=period if period else None,
        )
    except Exception as exc:
        return _error("search_transcript", request, exc)
    if not results:
        return _empty(
            "search_transcript",
            request,
            f"No transcript for {ticker.upper()} matching query (period={period or 'any'}). It may not have been ingested yet.",
        )
    try:
        records = [_vector_record(result, "transcript") for result in results]
    except Exception as exc:
        return _error("search_transcript", request, exc)
    return evidence_envelope(tool="search_transcript", query=request, records=records)


def _fmt_revision_signal(pct: float) -> str:
    if pct >= 1.0:
        return "positive momentum"
    if pct <= -1.0:
        return "negative momentum"
    return "stable"


@tool
def get_analyst_expectations(ticker: str) -> str:
    """Fetch analyst consensus EPS / revenue estimates, 30-day estimate revision %,
    and post-earnings d1/d5/since-release stock price reaction. Use to compare reported
    actuals against market expectations and gauge market reception of the latest filing.
    Returns an evidence.v1 JSON envelope with explicit period/event alignment flags."""
    from ingestion.analyst import fetch_analyst_estimates, fetch_price_reaction

    request = {"ticker": ticker.upper()}
    try:
        rows = metrics_db.get_all_metrics(ticker)
        latest_filing_date = rows[0]["filing_date"] if rows else None
        target_period = rows[0].get("period") if rows else None

        est, est_err = fetch_analyst_estimates(ticker, target_period=target_period)
        price, price_err = (fetch_price_reaction(ticker, latest_filing_date)
                             if latest_filing_date else (None, "no filing date in metrics_db"))
    except Exception as exc:
        return _error("get_analyst_expectations", request, exc)

    if est is None and price is None:
        return _empty(
            "get_analyst_expectations",
            request,
            f"No analyst data available for {ticker.upper()}. Errors: {est_err}; {price_err}",
        )

    # Use exact schema field names so the LLM can copy values verbatim without renaming.
    lines = [f"Analyst expectations for {ticker.upper()} (field names match MarketExpectations schema):"]
    if est:
        estimates_aligned = bool(est.get("period_aligned") and est.get("comparison_allowed"))
        eps = est.get("consensus_eps_est") if estimates_aligned else None
        rev = est.get("consensus_rev_est") if estimates_aligned else None
        rev30 = est.get("estimate_revision_30d_pct") if estimates_aligned else None
        lines.append(f"  consensus_eps_est: {eps:.4f}" if eps is not None else "  consensus_eps_est: null")
        lines.append(
            f"  consensus_rev_est_bn: {rev / 1e9:.4f}"  # already converted to billions
            if rev is not None else "  consensus_rev_est_bn: null"
        )
        lines.append(
            f"  revision_30d_pct: {rev30:.4f}  # signal: {_fmt_revision_signal(rev30)}"
            if rev30 is not None else "  revision_30d_pct: null"
        )
        lines.append(f"  target_period: {est.get('target_period') or 'null'}")
        lines.append(f"  provider_period_codes: {est.get('provider_period_codes') or {}}")
        lines.append(f"  period_aligned: {str(bool(est.get('period_aligned'))).lower()}")
        lines.append(f"  comparison_allowed: {str(bool(est.get('comparison_allowed'))).lower()}")
        lines.append(f"  alignment_status: {est.get('alignment_status') or 'UNVERIFIED'}")
    else:
        lines.append(f"  consensus_eps_est: null  # unavailable: {est_err}")
        lines.append("  consensus_rev_est_bn: null")
        lines.append("  revision_30d_pct: null")

    if price:
        event_allowed = bool(price.get("event_aligned") and price.get("comparison_allowed"))
        d1 = price.get("d1_pct") if event_allowed else None
        d5 = price.get("d5_pct") if event_allowed else None
        since = price.get("since_release_pct") if event_allowed else None
        lines.append(f"  d1_price_reaction_pct: {d1:.4f}" if d1 is not None else "  d1_price_reaction_pct: null")
        lines.append(f"  d5_price_reaction_pct: {d5:.4f}" if d5 is not None else "  d5_price_reaction_pct: null")
        lines.append(f"  since_release_price_reaction_pct: {since:.4f}" if since is not None else "  since_release_price_reaction_pct: null")
        lines.append(f"  event_date: {price.get('event_date') or 'null'}")
        lines.append(f"  event_kind: {price.get('event_kind') or 'unknown'}")
        lines.append(f"  event_timing: {price.get('event_timing') or 'unknown'}")
        lines.append(f"  event_aligned: {str(bool(price.get('event_aligned'))).lower()}")
        lines.append(f"  event_comparison_allowed: {str(bool(price.get('comparison_allowed'))).lower()}")
        lines.append(f"  price_alignment_status: {price.get('alignment_status') or 'UNVERIFIED'}")
    else:
        lines.append(f"  d1_price_reaction_pct: null  # unavailable: {price_err}")
        lines.append("  d5_price_reaction_pct: null")
        lines.append("  since_release_price_reaction_pct: null")

    try:
        content = "\n".join(lines)
        record = make_evidence_record(
            source="analyst",
            content=content,
            document_id=f"analyst:{ticker.upper()}:{latest_filing_date or 'current'}",
        as_of=(est or {}).get("as_of") or latest_filing_date,
            metadata={
                "ticker": ticker.upper(),
                "latest_filing_date": latest_filing_date,
                "target_period": target_period,
                "estimate_error": est_err,
                "price_error": price_err,
            },
        )
    except Exception as exc:
        return _error("get_analyst_expectations", request, exc)
    return evidence_envelope(tool="get_analyst_expectations", query=request, records=[record])


@tool
def search_news(query: str, ticker: str, days: int = 30) -> str:
    """Search for recent news about a company using Tavily.
    Use to find events after the most recent filing date.

    Args:
        query: The search query (e.g., "earnings", "acquisition")
        ticker: Stock ticker (e.g., "AAPL") — used to ground the search
        days: Number of days back to search (default: 30)

    Returns:
        An evidence.v1 JSON envelope (OK, EMPTY, or ERROR). Copy a supporting
        record's `ref` exactly when citing news."""
    request = {"query": query, "ticker": ticker.upper(), "days": days}
    # Check cache first (keyed by original args, 1-hour TTL)
    cached = news_cache.get(ticker, query, days)
    if cached is not None:
        if parse_evidence_envelope(cached):
            return cached
        # Legacy cache payloads have no immutable locator or source-level
        # provenance.  Never repackage them as newly verified evidence.
        return _empty(
            "search_news",
            request,
            "Legacy cached news omitted because its provenance cannot be verified.",
        )

    try:
        # Look up company name from metrics
        metrics_rows = metrics_db.get_all_metrics(ticker)
        company_name = metrics_rows[0]["company_name"] if metrics_rows else ticker

        # Get the latest filing date for filtering
        latest_filing_date = metrics_rows[0]["filing_date"] if metrics_rows else None

        # Compose ticker-grounded query
        full_query = f"{ticker.upper()} {company_name} {query}"

        results = _get_tavily_client().search(
            query=full_query,
            max_results=5,
            search_depth="basic",
            days=days,
            topic="news"
        ).get("results", [])
    except Exception as exc:
        return _error("search_news", request, exc)

    # Filter results by filing date
    if results and latest_filing_date:
        results = [r for r in results if r.get("published_date", "") >= latest_filing_date]

    if not results:
        result_str = _empty("search_news", request, "No recent news found.")
    else:
        try:
            records = []
            for r in results:
                published = str(r.get("published_date") or "unknown")
                url = str(r.get("url") or "")
                title = str(r.get("title") or "Untitled")
                content = f"Source: news | {published}\nTitle: {title}\nURL: {url}\n{r.get('content', '')}"
                records.append(make_evidence_record(
                    source="news",
                    content=content,
                    document_id=url or f"news:{ticker.upper()}:{published}:{title}",
                    source_url=url or None,
                    as_of=published if published != "unknown" else None,
                    metadata={
                        "ticker": ticker.upper(),
                        "title": title,
                        "published_date": published,
                    },
                ))
        except Exception as exc:
            return _error("search_news", request, exc)
        result_str = evidence_envelope(tool="search_news", query=request, records=records)

    news_cache.put(ticker, query, days, result_str)
    return result_str