server.py

import os
import uuid
from contextlib import asynccontextmanager
from typing import Annotated

import dotenv
from fastapi import FastAPI, Body, Response, Cookie
from jixia.structs import LeanName
from psycopg.rows import scalar_row, class_row
from psycopg.types.json import Jsonb
from psycopg_pool import ConnectionPool
from pydantic import BaseModel, field_validator
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.errors import RateLimitExceeded
from slowapi.middleware import SlowAPIMiddleware
from slowapi.util import get_remote_address
from starlette.requests import Request

from fastapi.middleware.cors import CORSMiddleware

from query_expansion import QueryExpander
from engine import QueryResult, PhyslibSearchEngine, Record


@asynccontextmanager
async def lifespan(app: FastAPI):
    dotenv.load_dotenv()
    with ConnectionPool(
        os.environ["CONNECTION_STRING"],
        kwargs={"autocommit": True},
        check=ConnectionPool.check_connection,
    ) as pool:
        app.expander = QueryExpander(os.environ["GEMINI_FAST_MODEL"])
        app.retriever = PhyslibSearchEngine(os.environ["CHROMA_PATH"], None)
        app.pool = pool
        yield


def client_key(request: Request) -> str:
    # Requests reach this app through two proxies (the platform router, then the
    # Next.js rewrite), so the socket peer is always 127.0.0.1 -- keying on it
    # would put every visitor in one shared bucket and let a handful of users
    # rate-limit everyone else. Use the originating client from X-Forwarded-For
    # (leftmost entry) and fall back to the peer address only if it's absent.
    forwarded = request.headers.get("x-forwarded-for")
    if forwarded:
        return forwarded.split(",")[0].strip()
    return get_remote_address(request)


limiter = Limiter(key_func=client_key, default_limits=["1/second"])
app = FastAPI(lifespan=lifespan)
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
app.add_middleware(
    CORSMiddleware,
    allow_origins=os.environ.get("ALLOWED_ORIGINS", "http://localhost:3000").split(","),
    allow_methods=["*"],
    allow_headers=["*"],
    allow_credentials=True,
)


@app.middleware("http")
async def set_connection(request: Request, call_next):
    with app.pool.connection() as conn:
        app.retriever.conn = conn
        return await call_next(request)


app.add_middleware(SlowAPIMiddleware)


@app.post("/search")
# /search inherited the global 1/second default because it had no decorator of
# its own, while cheaper endpoints (/fetch at 10/second) were given explicit,
# larger budgets. That was an omission rather than a decision, and it was the
# binding constraint on the theorem-proving client, which issues one search per
# subgoal attempt.
#
# 20/second is still a ceiling, deliberately: every /search makes a paid Gemini
# embedding call before it reaches Chroma, so an unbounded endpoint is an
# unbounded bill as well as a way to exhaust the threadpool — the handler is
# sync and GeminiEmbedding retries for up to 50s before giving up.
@limiter.limit("20/second")
def search(
    request: Request,
    response: Response,
    query: list[str],
    num_results: Annotated[int, Body(gt=0, le=150)] = 10,
) -> list[list[QueryResult]]:
    if len(query) == 1:
        with app.retriever.conn.cursor(row_factory=scalar_row) as cursor:
            cursor.execute(
                """
                           INSERT INTO physlibsearch.query(id, query, time)
                           VALUES (GEN_RANDOM_UUID(), %s, NOW())
                           RETURNING id
                           """,
                (query[0],),
            )
            session_id = cursor.fetchone()
            response.set_cookie("session", str(session_id))
    else:
        with app.retriever.conn.cursor() as cursor:
            cursor.executemany(
                """
                               INSERT INTO physlibsearch.query(id, query, time)
                               VALUES (GEN_RANDOM_UUID(), %s, NOW())
                               """,
                [(q,) for q in query],
            )

    return app.retriever.find_declarations(query, num_results)


@app.post("/fetch")
@limiter.limit("10/second")
def fetch(request: Request, query: list[LeanName]) -> list[Record]:
    return app.retriever.fetch_declarations(query)


@app.post("/expand")
@limiter.limit("15/minute")
async def expand(request: Request, query: Annotated[str, Body()]) -> str:
    expanded = await app.expander.expand(query)
    return expanded


class Feedback(BaseModel):
    declaration: LeanName
    action: str
    cancel: bool | None = None


class UserFeedback(BaseModel):
    tab_name: str
    rating: int
    feedback_type: str = "general"
    message: str


class ModuleInfo(BaseModel):
    name: LeanName
    count: int
    docstring: str | None = None

    @field_validator("docstring", mode="before")
    @classmethod
    def join_sections(cls, v: object) -> str | None:
        if v is None:
            return None
        if isinstance(v, list):
            return "\n\n".join(str(s) for s in v if s)
        if isinstance(v, str) and v.startswith("{") and v.endswith("}"):
            # Parse PostgreSQL text[] literal: {"section1","section2",...}
            inner = v[1:-1]
            parts: list[str] = []
            i = 0
            while i < len(inner):
                if inner[i] == '"':
                    j = i + 1
                    buf: list[str] = []
                    while j < len(inner):
                        if inner[j] == "\\" and j + 1 < len(inner):
                            buf.append(inner[j + 1])
                            j += 2
                        elif inner[j] == '"':
                            j += 1
                            break
                        else:
                            buf.append(inner[j])
                            j += 1
                    parts.append("".join(buf))
                    i = j
                    if i < len(inner) and inner[i] == ",":
                        i += 1
                else:
                    end = inner.find(",", i)
                    if end == -1:
                        end = len(inner)
                    token = inner[i:end].strip()
                    if token.upper() not in ("NULL", ""):
                        parts.append(token)
                    i = end + 1
            return "\n\n".join(p for p in parts if p.strip())
        return v


@app.get("/modules")
@limiter.limit("30/minute")
def list_modules(request: Request) -> list[ModuleInfo]:
    with app.pool.connection() as conn:
        with conn.cursor(row_factory=class_row(ModuleInfo)) as cursor:
            cursor.execute("""
                SELECT d.module_name AS name, COUNT(*) AS count, mo.docstring
                FROM declaration d
                INNER JOIN module mo ON mo.name = d.module_name
                WHERE d.visible = TRUE
                  AND EXISTS(SELECT 1 FROM informal i WHERE i.symbol_name = d.name)
                GROUP BY d.module_name, mo.docstring
                ORDER BY d.module_name
            """)
            return cursor.fetchall()


@app.post("/modules/declarations")
@limiter.limit("30/minute")
def module_declarations(request: Request, module_name: LeanName) -> list[Record]:
    with app.pool.connection() as conn:
        with conn.cursor(row_factory=class_row(Record)) as cursor:
            cursor.execute(
                """
                SELECT r.*
                FROM record r
                INNER JOIN declaration d ON r.name = d.name
                WHERE r.module_name = %s AND d.visible = TRUE
                ORDER BY d.index
            """,
                (Jsonb(module_name),),
            )
            return cursor.fetchall()


@app.post("/user-feedback")
@limiter.limit("10/minute")
async def user_feedback(request: Request, body: UserFeedback):
    if not (1 <= body.rating <= 5):
        from fastapi import HTTPException

        raise HTTPException(status_code=422, detail="Rating must be between 1 and 5")
    with app.pool.connection() as conn:
        with conn.cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO physlibsearch.user_feedback
                    (tab_name, rating, feedback_type, message, created_at)
                VALUES (%s, %s, %s, %s, NOW())
                """,
                (body.tab_name, body.rating, body.feedback_type, body.message),
            )


@app.post("/feedback")
async def feedback(session: Annotated[str, Cookie()], body: Feedback):
    query_id = uuid.UUID(session)
    if body.cancel:
        with app.retriever.conn.cursor() as cursor:
            cursor.execute("DELETE FROM physlibsearch.feedback WHERE query_id = %s AND declaration_name = %s", (query_id, Jsonb(body.declaration)))
    else:
        with app.retriever.conn.cursor() as cursor:
            cursor.execute("INSERT INTO physlibsearch.feedback(query_id, declaration_name, action) VALUES (%s, %s, %s)", (query_id, Jsonb(body.declaration), body.action))