lorebase-guide 0.2.3__py3-none-any.whl

lorebase_guide/app.py

@@ -0,0 +1,451 @@
+"""The Guide HTTP service.
+
+- `POST /embed` `{text, mode: query|document}` -> `{model, dims, vector, mode}`;
+- `POST /explore` `{domain, branch?, question, max_steps?, budget?}` -> server-sent events:
+  `session`, then `thought`, `tool_call` and `tool_result` per step, then `answer` (or `error`
+  followed by `answer`), and `done`;
+- `GET /healthz` (liveness) and `GET /readyz` (Keeper reachable).
+
+Every request is authenticated with the caller's own credential and checked by Keeper before
+any model work (see `auth.py`); limits per credential protect the model budget (`limits.py`).
+Errors are problem documents (`application/problem+json`) with a `code`.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+import uuid
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
+from contextlib import AsyncExitStack
+from dataclasses import dataclass
+from typing import Annotated, Any, Literal
+
+from fastapi import FastAPI, Request
+from fastapi.exceptions import RequestValidationError
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse, Response, StreamingResponse
+from opentelemetry.trace import SpanKind
+from pydantic import BaseModel, ConfigDict, Field
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+from lorebase_client.errors import LorebaseError
+from lorebase_guide import __version__
+from lorebase_guide.agent import Budget, Exploration, Explorer
+from lorebase_guide.auth import AuthRefusedError, Caller, read_caller, verify
+from lorebase_guide.embed import Embedder
+from lorebase_guide.keeper import KeeperClients
+from lorebase_guide.limits import Gate, Refusal, SlidingWindow, Slots
+from lorebase_guide.settings import GuideSettings
+from lorebase_guide.telemetry import (
+    EMBEDDINGS,
+    EXPLORATIONS,
+    REFUSED,
+    REQUEST_DURATION,
+    REQUESTS,
+    tracer,
+)
+from lorebase_guide.trace import TraceSession, new_session_id
+from lorebase_inference import ChatModel, InferenceError, InferenceSettings, create_chat_model
+from lorebase_platform.logging import request_context
+
+logger = logging.getLogger("lorebase.guide")
+access_logger = logging.getLogger("lorebase.guide.access")
+
+NAME = r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$"
+BRANCH = r"^[A-Za-z0-9][A-Za-z0-9._/-]{0,127}$"
+PROBLEM = "application/problem+json"
+
+
+class EmbedRequest(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    text: str = Field(min_length=1, description="Text to embed.")
+    mode: Literal["query", "document"] = Field(
+        default="query",
+        description="query: inference.embeddings.query_prefix is put before the text; "
+        "document: document_prefix, as Scribe does for Record texts.",
+    )
+
+
+class EmbedResponse(BaseModel):
+    model: str = Field(description="Model name vectors are stored under (Keeper's model).")
+    dims: int
+    mode: Literal["query", "document"]
+    vector: list[float]
+
+
+class ExploreRequest(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    domain: Annotated[str, Field(pattern=NAME, description="Domain id.")]
+    branch: Annotated[str | None, Field(default=None, pattern=BRANCH)] = None
+    question: str = Field(min_length=1)
+    max_steps: int | None = Field(
+        default=None, ge=1, le=50, description="Lower than guide.agent.max_steps only."
+    )
+    budget: int | None = Field(
+        default=None,
+        ge=100,
+        description="Token budget; lower than guide.agent.max_total_tokens only.",
+    )
+
+
+def problem(
+    status: int, code: str, detail: str, *, headers: dict[str, str] | None = None
+) -> JSONResponse:
+    return JSONResponse(
+        {"type": "about:blank", "title": code, "status": status, "code": code, "detail": detail},
+        status_code=status,
+        media_type=PROBLEM,
+        headers=headers,
+    )
+
+
+def _refused(refusal: Refusal) -> JSONResponse:
+    REFUSED.add(1, {"reason": refusal.reason})
+    detail = {
+        "rate": "too many requests from this caller; retry later",
+        "concurrent": "this caller already has an exploration running",
+        "busy": "Guide is busy; retry later",
+    }.get(refusal.reason, "refused")
+    return problem(429, refusal.reason, detail, headers={"Retry-After": str(refusal.retry_after_s)})
+
+
+def _auth_refused(exc: AuthRefusedError) -> JSONResponse:
+    REFUSED.add(1, {"reason": exc.code})
+    headers = {"WWW-Authenticate": "Bearer"} if exc.status == 401 else None
+    return problem(exc.status, exc.code, exc.message, headers=headers)
+
+
+def sse(name: str, data: dict[str, Any]) -> bytes:
+    return f"event: {name}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n".encode()
+
+
+class _EventStream(StreamingResponse):
+    """A stream whose resources (concurrency slot, Keeper client) are released however the
+    response ends, also when the client goes away before the first event."""
+
+    def __init__(
+        self,
+        content: AsyncIterator[bytes],
+        cleanup: Callable[[], Awaitable[None]],
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(content, **kwargs)
+        self._cleanup = cleanup
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        try:
+            await super().__call__(scope, receive, send)
+        finally:
+            await self._cleanup()
+
+
+@dataclass
+class Runtime:
+    settings: GuideSettings
+    keepers: KeeperClients
+    embedder: Embedder
+    chat: Callable[[], ChatModel]
+    explore_window: SlidingWindow
+    embed_window: SlidingWindow
+    slots: Slots
+    verifying: Gate
+    _explorer: Explorer | None = None
+
+    def explorer(self) -> Explorer:
+        if self._explorer is None:
+            self._explorer = Explorer(self.chat(), self.settings.agent, self.embedder)
+        return self._explorer
+
+
+class AccessMiddleware:
+    """Request id, server span, body size limit, timing and an access log line without
+    values (method, route, status, duration)."""
+
+    def __init__(self, app: ASGIApp, *, max_body_bytes: int) -> None:
+        self._app = app
+        self._max_body = max_body_bytes
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self._app(scope, receive, send)
+            return
+        request_id = uuid.uuid4().hex
+        path = scope.get("path", "")
+        started = time.monotonic()
+        status = 500
+
+        async def tracked_send(message: Message) -> None:
+            nonlocal status
+            if message["type"] == "http.response.start":
+                status = message["status"]
+                headers = list(message.get("headers", []))
+                headers.append((b"x-request-id", request_id.encode()))
+                headers.append((b"x-content-type-options", b"nosniff"))
+                message = {**message, "headers": headers}
+            await send(message)
+
+        declared = _content_length(scope)
+        refusal: JSONResponse | None = None
+        if scope.get("method") == "POST":
+            if declared is None:
+                refusal = problem(411, "length_required", "send the body with a Content-Length")
+            elif declared > self._max_body:
+                refusal = problem(413, "payload_too_large", "request body too large")
+        with (
+            request_context(request_id),
+            tracer.start_as_current_span(
+                f"{scope.get('method', '')} {path}", kind=SpanKind.SERVER
+            ) as span,
+        ):
+            try:
+                if refusal is not None:
+                    await refusal(scope, receive, tracked_send)
+                else:
+                    await self._app(scope, _capped(receive, self._max_body), tracked_send)
+            finally:
+                elapsed = time.monotonic() - started
+                span.set_attribute("http.response.status_code", status)
+                route = path if path in _ROUTES else "other"
+                REQUESTS.add(1, {"route": route, "status": status})
+                REQUEST_DURATION.record(elapsed, {"route": route})
+                access_logger.info(
+                    "request",
+                    extra={
+                        "method": scope.get("method"),
+                        "route": route,
+                        "status": status,
+                        "duration_ms": round(elapsed * 1000),
+                    },
+                )
+
+
+def _capped(receive: Receive, limit: int) -> Receive:
+    """`receive` that ends the body once `limit` bytes were read (a Content-Length that lies
+    cannot make the service read more)."""
+    received = 0
+
+    async def capped() -> Message:
+        nonlocal received
+        message = await receive()
+        if message["type"] == "http.request":
+            body = message.get("body", b"")
+            if received + len(body) > limit:
+                body = body[: max(0, limit - received)]
+                message = {"type": "http.request", "body": body, "more_body": False}
+            received += len(body)
+        return message
+
+    return capped
+
+
+_ROUTES = frozenset({"/embed", "/explore", "/healthz", "/readyz"})
+
+
+def _content_length(scope: Scope) -> int | None:
+    for name, value in scope.get("headers", []):
+        if name == b"content-length":
+            try:
+                return int(value)
+            except ValueError:
+                return None
+    return None
+
+
+def _headers(request: Request) -> dict[str, str]:
+    return {k.lower(): v for k, v in request.headers.items()}
+
+
+def create_app(
+    settings: GuideSettings,
+    inference: InferenceSettings,
+    *,
+    chat: Callable[[], ChatModel] | None = None,
+    embedder: Embedder | None = None,
+    keeper_transport: Any = None,
+    clock: Callable[[], float] = time.monotonic,
+) -> FastAPI:
+    """The Guide application. `chat`, `embedder` and `keeper_transport` replace the configured
+    model, embedding model and network (tests)."""
+    runtime = Runtime(
+        settings=settings,
+        keepers=KeeperClients(settings, transport=keeper_transport),
+        embedder=embedder or Embedder(inference),
+        chat=chat or (lambda: create_chat_model(inference, task=settings.agent.task)),
+        explore_window=SlidingWindow(
+            settings.limits.explorations_per_minute,
+            clock,
+            max_keys=settings.limits.max_tracked_callers,
+        ),
+        embed_window=SlidingWindow(
+            settings.limits.embeds_per_minute,
+            clock,
+            max_keys=settings.limits.max_tracked_callers,
+        ),
+        slots=Slots(settings.limits.max_concurrent, settings.limits.max_concurrent_per_caller),
+        verifying=Gate(settings.limits.max_concurrent_verifications),
+    )
+    app = FastAPI(
+        title="Lorebase Guide",
+        version=__version__,
+        docs_url=None,
+        redoc_url=None,
+        openapi_url=None,
+    )
+    app.state.runtime = runtime
+    http = settings.http
+    if http.origins_fell_back:
+        logger.warning(
+            "guide.http.allowed_origins is an empty list; using guide.http.default_origins",
+            extra={"origins": list(http.origins)},
+        )
+
+    @app.exception_handler(RequestValidationError)
+    async def _invalid(request: Request, exc: RequestValidationError) -> JSONResponse:
+        problems = "; ".join(
+            f"{'.'.join(str(p) for p in e['loc'][1:]) or 'body'}: {e['msg']}"
+            for e in exc.errors()[:5]
+        )
+        return problem(422, "invalid_request", problems[:500])
+
+    @app.get("/healthz")
+    async def healthz() -> dict[str, str]:
+        return {"status": "ok"}
+
+    @app.get("/readyz")
+    async def readyz() -> JSONResponse:
+        ok = await runtime.keepers.probe(settings.readiness_timeout_s)
+        body = {"status": "ready" if ok else "not_ready", "checks": {"keeper": ok}}
+        return JSONResponse(body, status_code=200 if ok else 503)
+
+    async def _verified(caller: Caller) -> Caller | JSONResponse:
+        if not runtime.verifying.enter():
+            return _refused(Refusal("busy", 1))
+        try:
+            async with runtime.keepers.open(caller) as client:
+                return await verify(caller, client)
+        except AuthRefusedError as exc:
+            return _auth_refused(exc)
+        except LorebaseError:
+            return problem(502, "keeper_unavailable", "Keeper could not be reached")
+        finally:
+            runtime.verifying.leave()
+
+    @app.post("/embed", response_model=EmbedResponse)
+    async def embed(body: EmbedRequest, request: Request) -> Response:
+        try:
+            caller = read_caller(_headers(request), http)
+        except AuthRefusedError as exc:
+            return _auth_refused(exc)
+        if len(body.text) > settings.limits.max_embed_chars:
+            return problem(422, "text_too_long", "text is longer than allowed")
+        verified = await _verified(caller)
+        if isinstance(verified, JSONResponse):
+            return verified
+        refusal = runtime.embed_window.take(verified.identity)
+        if refusal is not None:
+            return _refused(refusal)
+        try:
+            result = await runtime.embedder.embed(body.text, body.mode)
+        except InferenceError as exc:
+            logger.warning("embedding failed", extra={"error": type(exc).__name__})
+            return problem(502, "embedding_failed", "the embedding model failed")
+        EMBEDDINGS.add(1, {"mode": body.mode})
+        payload = EmbedResponse(
+            model=result.model, dims=result.dims, mode=body.mode, vector=result.vector
+        )
+        return JSONResponse(payload.model_dump(), headers={"Cache-Control": "no-store"})
+
+    @app.post("/explore")
+    async def explore(body: ExploreRequest, request: Request) -> Response:
+        try:
+            caller = read_caller(_headers(request), http)
+        except AuthRefusedError as exc:
+            return _auth_refused(exc)
+        if len(body.question) > settings.agent.max_question_chars:
+            return problem(422, "question_too_long", "the question is longer than allowed")
+        if not runtime.verifying.enter():
+            return _refused(Refusal("busy", 1))
+        stack = AsyncExitStack()
+        try:
+            trace = TraceSession(settings.trace, new_session_id(), body.domain)
+            client = await stack.enter_async_context(runtime.keepers.open(caller, trace.headers()))
+            caller = await verify(caller, client)
+            explorer = runtime.explorer()
+        except AuthRefusedError as exc:
+            await stack.aclose()
+            return _auth_refused(exc)
+        except LorebaseError:
+            await stack.aclose()
+            return problem(502, "keeper_unavailable", "Keeper could not be reached")
+        except BaseException:
+            await stack.aclose()
+            raise
+        finally:
+            runtime.verifying.leave()
+        # Limits count verified callers only, keyed on the principal Keeper reported.
+        key = caller.identity
+        refusal = runtime.explore_window.take(key) or runtime.slots.acquire(key)
+        if refusal is not None:
+            await stack.aclose()
+            return _refused(refusal)
+        stack.callback(runtime.slots.release, key)
+        budget = Budget.of(settings.agent, max_steps=body.max_steps, budget=body.budget)
+        state = Exploration()
+
+        async def stream() -> AsyncGenerator[bytes]:
+            started = time.monotonic()
+            try:
+                async for event in explorer.run(
+                    client=client,
+                    trace=trace,
+                    domain=body.domain,
+                    branch=body.branch,
+                    question=body.question,
+                    budget=budget,
+                    state=state,
+                ):
+                    yield sse(event.name, event.data)
+            finally:
+                if state.outcome == "running":
+                    state.outcome = "aborted"
+                EXPLORATIONS.add(1, {"outcome": state.outcome})
+                logger.info(
+                    "exploration finished",
+                    extra={
+                        "session": trace.session_id,
+                        "outcome": state.outcome,
+                        "steps": state.steps,
+                        "input_tokens": state.usage.input_tokens,
+                        "output_tokens": state.usage.output_tokens,
+                        "duration_ms": round((time.monotonic() - started) * 1000),
+                    },
+                )
+                await stack.aclose()
+
+        events = stream()
+
+        async def cleanup() -> None:
+            await events.aclose()
+            await stack.aclose()
+
+        return _EventStream(
+            events,
+            cleanup,
+            media_type="text/event-stream",
+            headers={"Cache-Control": "no-store", "X-Accel-Buffering": "no"},
+        )
+
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=list(http.origins),
+        allow_methods=["GET", "POST"],
+        allow_headers=[http.credential_header, http.csrf_header, "Content-Type"],
+        allow_credentials=bool(http.session_cookie),
+        max_age=600,
+    )
+    app.add_middleware(AccessMiddleware, max_body_bytes=http.max_body_bytes)
+    return app

lorebase_guide/auth.py

@@ -0,0 +1,183 @@
+"""The caller's credential: read from the request, checked by Keeper, forwarded unchanged.
+
+Guide has no Keeper credential of its own and never substitutes one. A caller authenticates
+in one of three ways:
+
+- `bearer`: the credential header (`Authorization: Bearer <OIDC token or API key>`, a bare API
+  key is accepted too). Sent to Keeper as is.
+- `session`: Keeper's session cookie (`guide.http.session_cookie`) with its CSRF token in
+  `guide.http.csrf_header`, when Guide is served under the web app's origin. Only that one
+  cookie and the token are forwarded; Guide asks Keeper for the session's token
+  (`GET /auth/session`) and refuses the request unless the two match, so a cross-site page
+  cannot spend the user's model budget with the cookie alone.
+- `anonymous`: no credential, only with `guide.http.allow_anonymous` (Keeper in auth mode none).
+
+Before any model call, `verify` asks Keeper who the caller is: an unknown or expired
+credential is refused before it costs tokens. A request carrying an `Origin` that is not in
+`guide.http.allowed_origins` (or `default_origins` while that is empty), or marked cross-site by
+the browser, is refused.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import hmac
+import re
+from collections.abc import Mapping
+from dataclasses import dataclass
+from http.cookies import CookieError, SimpleCookie
+from typing import Any, Literal
+
+from lorebase_client import AsyncLorebaseClient
+from lorebase_client.core import Call, parse_json
+from lorebase_client.errors import LorebaseError
+from lorebase_guide.settings import HttpSettings
+
+CredentialKind = Literal["bearer", "session", "anonymous"]
+_MAX_CREDENTIAL_CHARS = 8_192
+_CROSS_SITE = "cross-site"
+_SESSION_ID = re.compile(r"[A-Za-z0-9_-]{32,128}")
+"""Shape of Keeper's session ids; any other cookie value is refused, never forwarded."""
+_CSRF_TOKEN = re.compile(r"[A-Za-z0-9_.~+/=-]{1,256}")
+
+
+class AuthRefusedError(Exception):
+    """The request is refused before any work: `status` is 401 or 403."""
+
+    def __init__(self, status: int, code: str, message: str) -> None:
+        super().__init__(message)
+        self.status = status
+        self.code = code
+        self.message = message
+
+
+@dataclass(frozen=True)
+class Caller:
+    """The credential of one request. Its repr never shows the secret parts."""
+
+    kind: CredentialKind
+    authorization: str | None = None
+    session_id: str | None = None
+    csrf_token: str | None = None
+    principal: str | None = None
+    """Principal id reported by Keeper once verified."""
+
+    def __repr__(self) -> str:
+        return f"Caller(kind={self.kind!r}, principal={self.principal!r})"
+
+    @property
+    def identity(self) -> str:
+        """Key of the per-caller limits: the verified principal, else the credential's
+        fingerprint (anonymous callers share one key)."""
+        return f"principal:{self.principal}" if self.principal else self.fingerprint
+
+    @property
+    def fingerprint(self) -> str:
+        """Stable, non-reversible id of the credential, used for per-caller limits."""
+        secret = self.authorization or self.session_id or ""
+        if not secret:
+            return "anonymous"
+        return hashlib.sha256(f"{self.kind}\x00{secret}".encode()).hexdigest()[:32]
+
+    def keeper_headers(self, http: HttpSettings) -> dict[str, str]:
+        """Headers that carry this credential to Keeper, and nothing else of the request."""
+        if self.kind == "bearer" and self.authorization:
+            return {"Authorization": self.authorization}
+        if self.kind == "session" and self.session_id and self.csrf_token:
+            return {
+                "Cookie": f"{http.session_cookie}={self.session_id}",
+                http.csrf_header: self.csrf_token,
+            }
+        return {}
+
+
+def _single_line(value: str) -> bool:
+    return "\r" not in value and "\n" not in value and len(value) <= _MAX_CREDENTIAL_CHARS
+
+
+def _cookie(headers: Mapping[str, str], name: str) -> str | None:
+    raw = headers.get("cookie")
+    if not raw or not name:
+        return None
+    jar: SimpleCookie = SimpleCookie()
+    try:
+        jar.load(raw)
+    except CookieError:
+        return None
+    morsel = jar.get(name)
+    return morsel.value if morsel is not None and morsel.value else None
+
+
+def check_origin(headers: Mapping[str, str], http: HttpSettings) -> None:
+    """Refuse browser requests from origins that are not configured, and cross-site ones."""
+    origin = headers.get("origin")
+    if origin is not None and origin.rstrip("/").lower() not in http.origins:
+        raise AuthRefusedError(
+            403, "origin_not_allowed", "requests from this origin are not allowed"
+        )
+    site = headers.get("sec-fetch-site")
+    if site is not None and site.lower() == _CROSS_SITE and origin is None:
+        raise AuthRefusedError(403, "cross_site", "cross-site requests are not allowed")
+
+
+def read_caller(headers: Mapping[str, str], http: HttpSettings) -> Caller:
+    """The credential of a request (lower-case header names), not yet verified."""
+    check_origin(headers, http)
+    authorization = headers.get(http.credential_header.lower())
+    if authorization is not None and authorization.strip():
+        value = authorization.strip()
+        if not _single_line(value):
+            raise AuthRefusedError(401, "unauthorized", "malformed credential")
+        scheme, _, rest = value.partition(" ")
+        if scheme.lower() == "bearer":
+            if not rest.strip():
+                raise AuthRefusedError(401, "unauthorized", "malformed credential")
+            value = f"Bearer {rest.strip()}"
+        elif rest:
+            raise AuthRefusedError(401, "unauthorized", "only Bearer credentials are accepted")
+        else:
+            value = f"Bearer {value}"
+        return Caller(kind="bearer", authorization=value)
+    session_id = _cookie(headers, http.session_cookie)
+    if session_id is not None:
+        token = (headers.get(http.csrf_header.lower()) or "").strip()
+        if not _SESSION_ID.fullmatch(session_id):
+            raise AuthRefusedError(401, "unauthorized", "malformed session cookie")
+        if not _CSRF_TOKEN.fullmatch(token):
+            raise AuthRefusedError(
+                403, "csrf", f"a session cookie needs its CSRF token in {http.csrf_header}"
+            )
+        return Caller(kind="session", session_id=session_id, csrf_token=token)
+    if http.allow_anonymous:
+        return Caller(kind="anonymous")
+    raise AuthRefusedError(
+        401,
+        "unauthorized",
+        f"send `{http.credential_header}: Bearer <token or API key>` or the session cookie",
+    )
+
+
+async def verify(caller: Caller, client: AsyncLorebaseClient) -> Caller:
+    """Ask Keeper who the caller is, with the caller's own credential; for a session, also
+    check the CSRF token against the session's. Returns the caller with its principal."""
+    try:
+        who = await client.whoami()
+        if caller.kind == "session":
+            session: Any = await client.send(Call("GET", "/auth/session", parse_json, safe=True))
+            expected = str(session.get("csrf_token") or "") if isinstance(session, dict) else ""
+            given = caller.csrf_token or ""
+            if not expected or not hmac.compare_digest(expected.encode(), given.encode()):
+                raise AuthRefusedError(403, "csrf", "the CSRF token does not match the session")
+    except LorebaseError as exc:
+        if exc.status in (401, 403):
+            raise AuthRefusedError(
+                exc.status, "unauthorized", "Keeper refused the credential"
+            ) from None
+        raise
+    return Caller(
+        kind=caller.kind,
+        authorization=caller.authorization,
+        session_id=caller.session_id,
+        csrf_token=caller.csrf_token,
+        principal=who.id or None,
+    )