tools.py

"""Agent tools. Each returns JSON-serializable data; every number carries its
source accession so the verifier and the UI can prove it.
"""

import logging
from functools import lru_cache

from db import query
from market import get_quote as market_get_quote
from retrieval import search_passages

logger = logging.getLogger(__name__)

# Popular abbreviations that differ from the exchange symbol we store. Indian
# companies are the main source of this: the agent tends to reach for the
# colloquial name (RIL, SBI, Airtel) rather than the NSE symbol.
TICKER_ALIASES = {
    "RIL": "RELIANCE", "INFOSYS": "INFY", "HDFC": "HDFCBANK", "ICICI": "ICICIBANK",
    "SBI": "SBIN", "STATEBANK": "SBIN", "LNT": "LT", "L&T": "LT", "LARSEN": "LT",
    "MAHINDRA": "M&M", "BHARTI": "BHARTIARTL", "AIRTEL": "BHARTIARTL",
    "KOTAK": "KOTAKBANK", "AXIS": "AXISBANK", "TATAMOTORS": "TMPV",
    "MARUTISUZUKI": "MARUTI", "ASIANPAINTS": "ASIANPAINT", "BAJAJFINANCE": "BAJFINANCE",
    "SUNPHARMACEUTICAL": "SUNPHARMA", "ULTRATECH": "ULTRACEMCO", "ZOMATO": "ETERNAL",
}


@lru_cache(maxsize=1)
def _directory() -> tuple[frozenset[str], tuple[tuple[str, str], ...]]:
    rows = query("select ticker, name from companies")
    tickers = frozenset(t for t, _ in rows)
    names = tuple((t, n.lower()) for t, n in rows)
    return tickers, names


def resolve_ticker(raw: str) -> str:
    """Map a user/agent-supplied name or alias to a symbol we actually hold.

    Tries: exact symbol, curated alias, then a company-name match (so 'Reliance'
    or 'Tata Consultancy' resolve even without the exact ticker). Falls back to
    the input unchanged so unknown tickers still surface a clean 'no data'."""
    token = raw.upper().strip()
    tickers, names = _directory()
    if token in tickers:
        return token
    stripped = token.replace(" ", "").replace(".", "")
    if stripped in TICKER_ALIASES:
        return TICKER_ALIASES[stripped]
    needle = raw.lower().strip()
    if len(needle) >= 3:
        for ticker, name in names:
            if needle in name or name.startswith(needle):
                return ticker
    return token

# Canonical metric -> ordered fallback chain of concepts. First concept with
# data wins (companies tag revenue differently). Chains mix us-gaap (SEC) and
# Indian results-taxonomy concepts (ind-as / banking / insurance); a ticker
# only ever matches its own region's concepts, so the union is safe.
METRICS: dict[str, list[str]] = {
    "revenue": [
        "RevenueFromContractWithCustomerExcludingAssessedTax",
        "Revenues",
        "SalesRevenueNet",
        "RevenueFromOperations",     # ind-as
        "InterestEarned",            # banks: interest income is the top line
        "GrossPremiumIncome",        # life insurers
    ],
    "cost_of_revenue": ["CostOfRevenue", "CostOfGoodsAndServicesSold"],
    "gross_profit": ["GrossProfit"],
    "operating_income": ["OperatingIncomeLoss", "OperatingProfitBeforeProvisionAndContingencies"],
    "net_income": [
        "NetIncomeLoss",
        "ProfitLossForPeriod",                     # ind-as
        "ProfitLossForThePeriod",                  # banking
        "ProfitLossAfterTaxAndExtraordinaryItems", # insurance
    ],
    "profit_before_tax": [
        "ProfitBeforeTax",                          # ind-as
        "ProfitLossFromOrdinaryActivitiesBeforeTax",  # banking
        "ProfitLossBeforeTax",                      # insurance
    ],
    "rnd_expense": ["ResearchAndDevelopmentExpense"],
    "sga_expense": ["SellingGeneralAndAdministrativeExpense"],
    "eps_basic": [
        "EarningsPerShareBasic",
        "BasicEarningsLossPerShareFromContinuingAndDiscontinuedOperations",
        "BasicEarningsPerShareAfterExtraordinaryItems",
        "BasicAndDilutedEPSAfterExtraordinaryItemsNetOfTaxExpenseForThePeriodNotToBeAnnualized",
    ],
    "eps_diluted": [
        "EarningsPerShareDiluted",
        "DilutedEarningsLossPerShareFromContinuingAndDiscontinuedOperations",
        "DilutedEarningsPerShareAfterExtraordinaryItems",
        "BasicAndDilutedEPSAfterExtraordinaryItemsNetOfTaxExpenseForThePeriodNotToBeAnnualized",
    ],
    "shares_diluted": ["WeightedAverageNumberOfDilutedSharesOutstanding"],
    "total_assets": ["Assets"],  # same concept name in us-gaap and ind-as
    "total_liabilities": ["Liabilities"],
    "stockholders_equity": ["StockholdersEquity"],
    "cash": ["CashAndCashEquivalentsAtCarryingValue", "CashAndCashEquivalentsCashFlowStatement"],
    "long_term_debt": ["LongTermDebtNoncurrent", "LongTermDebt", "BorrowingsNoncurrent"],
    "inventory": ["InventoryNet"],
    "operating_cash_flow": [
        "NetCashProvidedByUsedInOperatingActivities",
        "CashFlowsFromUsedInOperatingActivities",   # ind-as (half-yearly)
    ],
    "capex": ["PaymentsToAcquirePropertyPlantAndEquipment"],
    "buybacks": ["PaymentsForRepurchaseOfCommonStock"],
    "dividends_paid": ["PaymentsOfDividendsCommonStock"],
    # India-specific: banks
    "interest_expended": ["InterestExpended"],
    "provisions": ["ProvisionsOtherThanTaxAndContingencies"],
    "gross_npa_pct": ["PercentageOfGrossNpa"],
    "net_npa_pct": ["PercentageOfNpa"],
    "cet1_ratio": ["CET1Ratio"],
    "return_on_assets": ["ReturnOnAssets"],
    # India-specific: insurers
    "net_premium_income": ["NetPremiumIncome"],
}

# Some filers zero-fill ratio fields in the consolidated XBRL and disclose the
# real numbers only in the standalone filing (HDFC Bank, SBI). A 0.00% GNPA or
# CET1 is regulatorily impossible, so for these metrics zero means "blank".
ZERO_MEANS_UNREPORTED = {"gross_npa_pct", "net_npa_pct", "cet1_ratio", "return_on_assets"}


def query_facts(ticker: str, metrics: list[str], annual: bool = True, years: int = 5) -> dict:
    """Reported fundamentals from XBRL (SEC 10-K/10-Q or NSE results filings)."""
    ticker = resolve_ticker(ticker)
    years = max(1, min(int(years), 15))
    unknown = [m for m in metrics if m not in METRICS]
    if unknown:
        return {"error": f"Unknown metrics {unknown}. Valid: {sorted(METRICS)}"}

    results: dict[str, list[dict]] = {}
    for metric in metrics:
        # Query the whole fallback chain at once: companies switch concepts
        # between years (e.g. Revenues vs RevenueFromContract...), so picking
        # "first concept with data" can silently return only old years.
        chain = METRICS[metric]
        zero_filter = "and f.value <> 0" if metric in ZERO_MEANS_UNREPORTED else ""
        if annual:
            # Full-year duration facts from 10-Ks (or instant facts for balance sheet)
            rows = query(
                f"""
                select distinct on (f.fiscal_year)
                       f.fiscal_year, f.period_end::text, f.value, f.unit, f.accession, f.form
                from xbrl_facts f join companies c on c.cik = f.cik
                where c.ticker = %s and f.concept = any(%s)
                  and f.form in ('10-K', 'RESULTS')
                  and (f.period_start is null or f.period_end - f.period_start > 300)
                  and f.fiscal_period = 'FY' {zero_filter}
                order by f.fiscal_year desc, f.period_end desc,
                         array_position(%s, f.concept)
                limit %s
                """,
                (ticker, chain, chain, years),
            )
        else:
            rows = query(
                f"""
                select distinct on (f.period_end)
                       f.fiscal_year, f.period_end::text, f.value, f.unit, f.accession, f.form
                from xbrl_facts f join companies c on c.cik = f.cik
                where c.ticker = %s and f.concept = any(%s)
                  and (f.period_start is null or f.period_end - f.period_start between 60 and 120)
                  {zero_filter}
                order by f.period_end desc, array_position(%s, f.concept)
                limit %s
                """,
                (ticker, chain, chain, years * 4),
            )

        results[metric] = [
            {
                "fiscal_year": fy,
                "period_end": period_end,
                "value": float(value),
                "unit": unit,
                "source_accession": accession,
                "source_form": form,
            }
            for fy, period_end, value, unit, accession, form in rows
        ]
    return {
        "ticker": ticker,
        "annual": annual,
        "metrics": results,
        "note": (
            "Show these values explicitly in the answer, labeled by fiscal year, in a "
            "table — never just a qualitative summary. Compute the ratios the question "
            "needs (margins, growth %, CAGR) from them and show the working. Values are "
            "absolute in the stated unit: INR figures should be presented in ₹ crore "
            "(divide by 10,000,000), USD figures in $B/$M as appropriate."
        ),
    }


def retrieve_passages(query_text: str, ticker: str = "", form: str = "", k: int = 6) -> dict:
    """Semantic search over indexed 10-K/10-Q / annual report text."""
    resolved = resolve_ticker(ticker) if ticker else None
    hits = search_passages(query_text, ticker=resolved, form=form or None, k=k)
    return {"passages": hits}


def get_quote(ticker: str) -> dict:
    """Live market quote."""
    try:
        return market_get_quote(ticker)
    except Exception as exc:
        return {"error": f"No market data for '{ticker}': {exc}"}


def list_companies() -> dict:
    """Companies available in the local dataset."""
    rows = query("select ticker, name from companies order by ticker")
    return {"companies": [{"ticker": t, "name": n} for t, n in rows]}


# ---- Gemini function declarations ----

TOOL_DECLARATIONS = [
    {
        "name": "query_facts",
        "description": (
            "Audited fundamentals (XBRL) for one company — US SEC filings ($) or "
            "Indian NSE results (₹). The source for every quantitative claim: revenue, "
            "income, EPS, balance sheet, cash flow, buybacks, and the inputs for margins "
            "and growth; for Indian banks also NPA ratios, provisions and capital, for "
            "insurers premium income. Batch EVERY metric the question needs into ONE "
            "call — it is far cheaper than several calls. Each value carries its source "
            "filing accession for citation."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "Company symbol or name — NSE symbol for Indian names "
                                   "(RELIANCE, INFY, HDFCBANK) or SEC ticker for US (AAPL). "
                                   "Common names/aliases are resolved automatically.",
                },
                "metrics": {
                    "type": "array",
                    "items": {"type": "string", "enum": sorted(METRICS)},
                    "description": "All metrics needed to answer, in one call. For a margin "
                                   "question include revenue plus the profit lines; for growth "
                                   "include the metric across enough periods.",
                },
                "annual": {
                    "type": "boolean",
                    "description": "true = full fiscal years (default, for yearly trends); "
                                   "false = individual quarters.",
                },
                "years": {
                    "type": "integer",
                    "description": "Periods of history to return. Default 5; use 4-5+ for a "
                                   "trend so the answer isn't thin.",
                },
            },
            "required": ["ticker", "metrics"],
        },
    },
    {
        "name": "retrieve_passages",
        "description": (
            "Semantic search over indexed filing text — US 10-K/10-Q filings and "
            "Indian annual reports. Use for qualitative questions: risks, strategy, "
            "management discussion, segments, guidance language. Returns passages "
            "with source filing and location (SEC Item section, or annual-report page)."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "query_text": {"type": "string"},
                "ticker": {"type": "string", "description": "Optional ticker filter"},
                "form": {
                    "type": "string",
                    "enum": ["10-K", "10-Q", "ANNUAL"],
                    "description": "Optional filter; ANNUAL = Indian annual report. Omit to search all forms",
                },
                "k": {"type": "integer", "description": "Number of passages (default 6)"},
            },
            "required": ["query_text"],
        },
    },
    {
        "name": "get_quote",
        "description": "Live market quote: price, day change, market cap, 52-week range.",
        "parameters": {
            "type": "object",
            "properties": {"ticker": {"type": "string"}},
            "required": ["ticker"],
        },
    },
    {
        "name": "list_companies",
        "description": "List the companies available in the local SEC dataset.",
        "parameters": {"type": "object", "properties": {}},
    },
]

TOOL_FUNCTIONS = {
    "query_facts": query_facts,
    "retrieve_passages": retrieve_passages,
    "get_quote": get_quote,
    "list_companies": list_companies,
}