| from fastapi import FastAPI, Header |
| from fastapi.middleware.cors import CORSMiddleware |
| from pydantic import BaseModel |
| from typing import Optional, List |
| from contextlib import asynccontextmanager |
| import os |
|
|
| from auth import verify_api_key |
| from credits import check_and_deduct_credits |
| from acra import run_acra_pipeline |
| from classifier_inference import warm_up |
|
|
| @asynccontextmanager |
| async def lifespan(app: FastAPI): |
| print("Loading ACRA classifier...") |
| warm_up() |
| print("Ready \u2713") |
| yield |
|
|
| app = FastAPI(title="ACRA API", |
| description="Adaptive Contextual Retrieval Architecture \u2014 NurricAI", |
| version="1.0.0", lifespan=lifespan) |
|
|
| app.add_middleware(CORSMiddleware, allow_origins=["*"], |
| allow_methods=["*"], allow_headers=["*"]) |
|
|
| class IngestRequest(BaseModel): |
| texts: List[str] |
| metadata: Optional[List[dict]] = None |
| namespace: Optional[str] = "default" |
|
|
| class QueryRequest(BaseModel): |
| query: str |
| namespace: Optional[str] = "default" |
| top_k: Optional[int] = 5 |
| rerank: Optional[bool] = True |
| use_web: Optional[bool] = False |
| llm_endpoint: Optional[str] = None |
| llm_api_key: Optional[str] = None |
| llm_model: Optional[str] = None |
|
|
| class QueryResponse(BaseModel): |
| answer: str |
| sources: List[dict] |
| credits_used: int |
| credits_remaining: int |
| complexity: Optional[dict] = None |
| retrieval_source: Optional[str] = None |
| cost: Optional[dict] = None |
| cost: Optional[dict] = None |
|
|
| @app.get("/") |
| def root(): return {"status": "ACRA API is live \U0001F680", "docs": "/docs"} |
|
|
| @app.get("/health") |
| def health(): return {"status": "ok"} |
|
|
| @app.post("/v1/ingest") |
| async def ingest(body: IngestRequest, |
| x_api_key: str = Header(..., alias="X-API-Key")): |
| user = await verify_api_key(x_api_key) |
| cost = max(1, len(body.texts) // 10) |
| remaining = await check_and_deduct_credits(user["id"], cost) |
| inserted = await run_acra_pipeline( |
| mode="ingest", texts=body.texts, |
| metadata=body.metadata or [{} for _ in body.texts], |
| namespace=body.namespace, user_id=user["id"]) |
| return {"status": "success", "chunks_indexed": inserted, |
| "credits_used": cost, "credits_remaining": remaining} |
|
|
| @app.post("/v1/query", response_model=QueryResponse) |
| async def query(body: QueryRequest, |
| x_api_key: str = Header(..., alias="X-API-Key")): |
| user = await verify_api_key(x_api_key) |
| remaining = await check_and_deduct_credits(user["id"], 1) |
| result = await run_acra_pipeline( |
| mode="query", query=body.query, namespace=body.namespace, |
| top_k=body.top_k, rerank=body.rerank, user_id=user["id"], |
| use_web=body.use_web) |
| return QueryResponse( |
| answer = result["answer"], |
| sources = result["sources"], |
| credits_used = 1, |
| credits_remaining = remaining, |
| complexity = result.get("complexity"), |
| retrieval_source = result.get("retrieval_source"), |
| cost = result.get("cost"), |
| ) |
|
|
| @app.get("/v1/usage") |
| async def usage(x_api_key: str = Header(..., alias="X-API-Key")): |
| user = await verify_api_key(x_api_key) |
| return {"plan": user["plan"], |
| "credits_remaining": user["credits_remaining"], |
| "credits_reset": user["credits_reset_at"]} |
|
|