Spaces:
Running
Running
File size: 3,559 Bytes
817f25b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """FinChat HTTP API (FastAPI).
A thin, headless wrapper around the existing RAG engine in ``src/rag.py``.
This is what the Vercel front-end calls. It exposes:
GET / -> service banner
GET /health -> liveness probe
GET /companies -> [{ticker, name}, ...] loaded in the corpus
POST /chat -> {answer, routed_to, sources[]}
GET /docs -> interactive Swagger UI (FastAPI built-in)
Run locally from the repo root:
uvicorn api.main:app --reload --port 7860
"""
from __future__ import annotations
import os
from contextlib import asynccontextmanager
from fastapi import FastAPI, Header, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from src.rag import answer, available_companies, get_vectorstore
# --- request / response models ---------------------------------------------
class ChatRequest(BaseModel):
question: str
class Source(BaseModel):
index: int
source: str
ticker: str | None = None
company: str | None = None
type: str | None = None
excerpt: str
class ChatResponse(BaseModel):
answer: str
routed_to: str | None = None
sources: list[Source]
# --- lifespan: warm the model once at boot ---------------------------------
@asynccontextmanager
async def lifespan(app: FastAPI):
# Load the vector store + embedding model a single time at startup, so the
# first real user request doesn't pay the cold-load cost.
get_vectorstore()
yield
app = FastAPI(title="FinChat API", version="1.0.0", lifespan=lifespan)
# Allow the Vercel front-end (and localhost during development) to call us.
# Set FRONTEND_ORIGIN in the Space secrets to your Vercel URL(s), comma-separated.
_origins = [o.strip() for o in os.getenv("FRONTEND_ORIGIN", "*").split(",") if o.strip()]
app.add_middleware(
CORSMiddleware,
allow_origins=_origins or ["*"],
allow_credentials=False,
allow_methods=["GET", "POST", "OPTIONS"],
allow_headers=["*"],
)
# Optional shared secret. If API_SECRET is set (as a Space secret), callers must
# send it as the "X-API-Key" header. The Vercel proxy adds it automatically, so
# random traffic can't hit this Space directly and burn the Groq daily quota.
_API_SECRET = os.getenv("API_SECRET")
def _check_auth(x_api_key: str | None) -> None:
if _API_SECRET and x_api_key != _API_SECRET:
raise HTTPException(status_code=401, detail="Invalid or missing API key.")
@app.get("/")
def root():
return {"service": "FinChat API", "status": "ok", "docs": "/docs"}
@app.get("/health")
def health():
return {"status": "healthy"}
@app.get("/companies")
def companies():
return [{"ticker": t, "name": n} for t, n in available_companies()]
@app.post("/chat", response_model=ChatResponse)
def chat(req: ChatRequest, x_api_key: str | None = Header(default=None)):
_check_auth(x_api_key)
question = (req.question or "").strip()
if not question:
raise HTTPException(status_code=400, detail="question must not be empty.")
result = answer(question)
sources = [
Source(
index=i,
source=d.metadata.get("source", ""),
ticker=d.metadata.get("ticker"),
company=d.metadata.get("company"),
type=d.metadata.get("type"),
excerpt=d.page_content[:500],
)
for i, d in enumerate(result["sources"], 1)
]
return ChatResponse(
answer=result["answer"],
routed_to=result["routed_to"],
sources=sources,
)
|