1"""Optional AI data assistant (phase 6): natural language → *validated* read-only SQL over the catalog views.23Guarantees: the model never executes anything — it only proposes SQL; FinData validates it with the same sandbox as4the playground (single SELECT, allow-listed directories, timeout, row cap), shows the SQL + EXPLAIN plan, cites the5dataset versions used, and refuses questions about data that is not in the catalog. Server-side BYOK (Anthropic)6only; status is `pending_configuration` until ANTHROPIC_API_KEY is set and the capability test has been recorded.7"""8from __future__ import annotations910import json11import time12from typing import Any1314from fastapi import APIRouter, Depends, Request15from findata_core.config import get_settings16from findata_core.db import connection17from findata_ingest.reader import dataset_views, run_sql18from pydantic import BaseModel, Field1920from findata_api import auth as A21from findata_api.errors import ApiError2223router = APIRouter(prefix="/assistant", tags=["assistant"])24MODEL = "claude-opus-5-5"252627class AskIn(BaseModel):28 question: str = Field(..., min_length=3, max_length=2000)29 max_rows: int = Field(200, ge=1, le=5000)30 execute: bool = True313233def _catalog_context() -> tuple[str, dict[str, Any]]:34 with connection() as conn:35 ds = conn.execute("SELECT id, name, interval, adjustment, status, current_version, timestamp_convention, timezone FROM datasets WHERE kind IN ('market_bars','derived') ORDER BY id").fetchall()36 cov = conn.execute("SELECT dataset_id, min(session_date) d0, max(session_date) d1, count(DISTINCT symbol) syms FROM coverage GROUP BY 1").fetchall()37 covm = {c["dataset_id"]: c for c in cov}38 lines = ["Available DuckDB views (one per dataset; columns: instrument_id, symbol_as_reported, asset_class, interval, timestamp_utc (bar OPEN, UTC), session_date, open, high, low, close, volume, currency, trade_count, vwap, adjustment_type, source_id, ingested_at, quality_flags, revision_id):"]39 versions = {}40 for d in ds:41 c = covm.get(d["id"])42 v = f"view {d['id'].replace('.', '_')}: {d['name']} — interval {d['interval']}, adjustment {d['adjustment']}, status {d['status']}, version {d['current_version']}"43 if c:44 v += f", sessions {c['d0']}→{c['d1']}, {c['syms']} symbols"45 else:46 v += ", NO DATA LOADED"47 lines.append(v)48 versions[d["id"]] = d["current_version"]49 lines.append("Fundamentals are NOT available through SQL views (use the fundamentals endpoint); say so if asked.")50 return "\n".join(lines), versions515253SYSTEM = """You translate research questions about a private market-data lake into ONE read-only DuckDB SQL statement.54Rules: SELECT/WITH only; use only the listed views and columns; never invent datasets, symbols, dates or columns that are not listed; if the question needs data that is not available (e.g. a dataset with NO DATA LOADED, fundamentals, options, ticks, a date outside coverage), set "sql" to null and explain in "refusal". Prices are raw (unadjusted) unless the view says otherwise. timestamp_utc marks the bar open; convert with timezone 'America/New_York' when the user talks about Eastern time. Always LIMIT results reasonably.55Respond with JSON only: {"sql": string|null, "explanation": string, "assumptions": [string], "refusal": string|null}."""565758def _ask_model(question: str, context: str) -> dict[str, Any]:59 import anthropic60 s = get_settings()61 client = anthropic.Anthropic(api_key=s.anthropic_api_key.get_secret_value(), max_retries=1, timeout=90.0)62 try:63 resp = client.messages.create(64 model=MODEL, max_tokens=4000, output_config={"effort": "medium"},65 system=[{"type": "text", "text": SYSTEM, "cache_control": {"type": "ephemeral"}}, {"type": "text", "text": context}],66 messages=[{"role": "user", "content": question}],67 )68 except anthropic.AuthenticationError:69 raise ApiError(502, "provider_auth", "Anthropic rejected the configured key") from None70 except anthropic.RateLimitError:71 raise ApiError(503, "provider_rate_limited", "Anthropic rate limit; retry later") from None72 except anthropic.APIStatusError as e:73 raise ApiError(502, "provider_error", f"Anthropic error {e.status_code}") from None74 except anthropic.APIConnectionError:75 raise ApiError(502, "provider_unreachable", "Could not reach Anthropic") from None76 if resp.stop_reason == "refusal":77 return {"sql": None, "explanation": "", "assumptions": [], "refusal": "The model declined this request."}78 text = "".join(b.text for b in resp.content if b.type == "text").strip()79 if text.startswith("```"):80 text = text.strip("`").split("\n", 1)[1] if "\n" in text else text.strip("`")81 text = text.rsplit("```", 1)[0]82 try:83 out = json.loads(text)84 except json.JSONDecodeError:85 return {"sql": None, "explanation": text[:500], "assumptions": [], "refusal": "Model output was not valid JSON"}86 out["_usage"] = {"input_tokens": resp.usage.input_tokens, "output_tokens": resp.usage.output_tokens,87 "cache_read_input_tokens": getattr(resp.usage, "cache_read_input_tokens", None), "model": resp.model}88 return out899091@router.get("/status")92def status(p: A.Principal = Depends(A.require("query"))):93 s = get_settings()94 configured = bool(s.anthropic_api_key.get_secret_value())95 return {"status": "active" if configured else "pending_configuration", "provider": "anthropic" if configured else None, "model": MODEL if configured else None,96 "note": None if configured else "Set ANTHROPIC_API_KEY in the node .env (server-side BYOK) and run the capability/cost test described in docs/STATUS.md.",97 "guarantees": ["model proposes SQL only; FinData validates and executes in the read-only sandbox", "generated SQL and EXPLAIN plan are always shown",98 "dataset versions used are cited", "questions about data not in the catalog are refused"]}99100101@router.post("")102def ask(body: AskIn, request: Request, p: A.Principal = Depends(A.require("query"))):103 s = get_settings()104 if not s.anthropic_api_key.get_secret_value():105 raise ApiError(409, "pending_configuration", "AI assistant is not configured (ANTHROPIC_API_KEY missing on the server)")106 context, versions = _catalog_context()107 t0 = time.monotonic()108 proposal = _ask_model(body.question, context)109 out: dict[str, Any] = {"question": body.question, "sql": proposal.get("sql"), "explanation": proposal.get("explanation"),110 "assumptions": proposal.get("assumptions") or [], "refusal": proposal.get("refusal"),111 "dataset_versions": versions, "model_usage": proposal.get("_usage"), "executed": False}112 if not proposal.get("sql"):113 A.audit(request, "assistant.ask", None, refused=True, question=body.question[:300])114 return out115 views = dataset_views()116 try:117 plan, _ = run_sql(f"EXPLAIN {proposal['sql']}", 50, 10, views)118 out["plan"] = "\n".join(str(r.get("explain_value") or r) for r in plan.to_pylist())119 except Exception as e:120 out["refusal"] = f"Generated SQL rejected by the sandbox: {str(e)[:300]}"121 A.audit(request, "assistant.ask", None, rejected=True, question=body.question[:300], sql=proposal["sql"][:500])122 return out123 if body.execute:124 try:125 tbl, meta = run_sql(proposal["sql"], min(body.max_rows, s.query_max_rows), s.query_timeout_seconds, views)126 except Exception as e:127 out["refusal"] = f"Execution failed: {str(e)[:300]}"128 return out129 rows = tbl.to_pylist()130 for r in rows:131 for k, v in r.items():132 if hasattr(v, "isoformat"):133 r[k] = v.isoformat()134 out.update({"executed": True, "columns": [{"name": f.name, "type": str(f.type)} for f in tbl.schema], "rows": rows,135 "truncated": meta.get("truncated"), "elapsed_ms": int((time.monotonic() - t0) * 1000)})136 A.audit(request, "assistant.ask", None, question=body.question[:300], sql=proposal["sql"][:500], rows=len(out.get("rows", [])))137 return out138