Add power-user AI Chat page; centralize demo-key model pinning
AI Chat is a plain scratchpad for talking to a model directly — no story context, scripts or world state — for poking at models, prompts and endpoints without starting an adventure. Power users only: the router 404s (rather than 403s) for everyone else and the nav link is hidden. The conversation lives in localStorage, so there's no new table or migration. is_power_user() now also returns True in local mode: it's the operator's own machine and their own key, the same reasoning that makes the provider debug log local-only. Alongside that, the rule keeping the shared demo key off paid models now lives in exactly one place. It had been duplicated into the chat router, which is how one copy eventually drifts: - resolve_provider_config() takes an optional model_override and is the only place the whitelist is applied, so turns, AI Chat and the connection test all inherit it. An override is a per-request preference, never a grant. - ProviderConfig.__post_init__ refuses to exist when api_key is the demo key and the model isn't whitelisted. It keys on the key itself rather than the using_demo flag, so a mislabelled config can't slip past, and it raises so a future path that bypasses the resolver fails loudly instead of billing. - The demo branch still pins endpoint_url too — a user-controlled endpoint would leak the key itself, which is worse than spending it. Provider gained chat(messages, ...) beside generate(), both delegating to a shared _stream(url, body); completion-mode endpoints get the messages flattened into a labelled transcript. Settings' /models fetch moved to list_endpoint_models() and is shared with /api/chat/config. Tests: 10 new in tests/test_chat.py (70 total). These deliberately do not stub resolve_provider_config — the point is to exercise the real BYOK-vs-demo decision and assert on what the provider actually received: off-whitelist override pinned, off-whitelist Settings.model pinned, redirected endpoint pinned, BYOK passed through untouched. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014FGY1yvzSeKgTtRfeVtDmx
This commit is contained in:
co-authored by
Claude Opus 5
parent
57c24a07d8
commit
1dd31086c1
+41
-5
@@ -81,19 +81,50 @@ def demo_enabled() -> bool:
|
||||
@dataclass
|
||||
class ProviderConfig:
|
||||
"""What the turn engine should actually connect with, after the
|
||||
BYOK-vs-demo decision."""
|
||||
BYOK-vs-demo decision. Build these with resolve_provider_config()."""
|
||||
|
||||
endpoint_url: str
|
||||
api_key: str
|
||||
model: str
|
||||
using_demo: bool
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Belt and braces around the shared demo key. resolve_provider_config()
|
||||
# already pins the model, but this makes it a property of the config
|
||||
# object itself: however it was built, and by whichever caller, the
|
||||
# server-funded key can never be paired with an off-whitelist (i.e.
|
||||
# possibly paid) model. Unreachable by design — a 500 here means a new
|
||||
# code path tried to bypass the pinning, which is worth failing loudly
|
||||
# rather than silently billing.
|
||||
if DEMO_API_KEY and self.api_key == DEMO_API_KEY and self.model not in DEMO_MODELS:
|
||||
raise ValueError(
|
||||
f"Refusing to use the shared demo key with non-whitelisted model {self.model!r}"
|
||||
)
|
||||
|
||||
def resolve_provider_config(settings: models.Settings) -> ProviderConfig:
|
||||
|
||||
def resolve_provider_config(
|
||||
settings: models.Settings, *, model_override: str | None = None
|
||||
) -> ProviderConfig:
|
||||
"""BYOK when the user has their own key, the shared demo key otherwise.
|
||||
|
||||
THE security-relevant branch is the demo one, and it is the only place the
|
||||
whitelist rule lives — every caller must come through here rather than
|
||||
building a ProviderConfig itself. On the demo key:
|
||||
|
||||
- the model is pinned to DEMO_MODELS, so a caller-supplied override (the AI
|
||||
Chat page) or a hand-edited Settings row cannot aim a server-funded key at
|
||||
a paid model; anything unrecognised falls back to DEMO_MODELS[0];
|
||||
- the endpoint is pinned to DEMO_ENDPOINT_URL, so the key itself can't be
|
||||
redirected to a URL the user controls and harvested.
|
||||
|
||||
`model_override` is a per-request preference (never a grant): it's honoured
|
||||
verbatim under BYOK, and only if whitelisted on the demo key.
|
||||
"""
|
||||
key = settings.api_key_plain
|
||||
requested = (model_override or "").strip() or settings.model
|
||||
if key or not demo_enabled():
|
||||
return ProviderConfig(settings.endpoint_url, key, settings.model, False)
|
||||
model = settings.model if settings.model in DEMO_MODELS else DEMO_MODELS[0]
|
||||
return ProviderConfig(settings.endpoint_url, key, requested, False)
|
||||
model = requested if requested in DEMO_MODELS else DEMO_MODELS[0]
|
||||
return ProviderConfig(DEMO_ENDPOINT_URL, DEMO_API_KEY, model, True)
|
||||
|
||||
|
||||
@@ -102,7 +133,12 @@ def _today() -> str:
|
||||
|
||||
|
||||
def is_power_user(user: models.User) -> bool:
|
||||
"""Trusted testers (email allowlist) bypass the demo turn cap."""
|
||||
"""Trusted testers: unmetered demo turns, plus tooling that isn't part of
|
||||
the game (the AI Chat scratchpad). Local installs are always trusted — it's
|
||||
the operator's own machine and their own API key, same reasoning as the
|
||||
provider debug log being local-only."""
|
||||
if not MULTI_USER:
|
||||
return True
|
||||
return bool(user.email) and user.email.lower() in POWER_USERS
|
||||
|
||||
|
||||
|
||||
@@ -25,6 +25,7 @@ from . import auth, models
|
||||
# scope -> (max requests, window seconds)
|
||||
RATE_LIMITS: dict[str, tuple[int, int]] = {
|
||||
"turn": (10, 60), # AI turn generation (demo key also has a daily cap)
|
||||
"chat": (30, 60), # AI Chat scratchpad (power users only)
|
||||
"script-test": (30, 60), # sandboxed, but each run costs up to 2s CPU
|
||||
"connection-test": (10, 60), # outbound HTTP to a user-supplied URL
|
||||
"import": (30, 60), # large writes
|
||||
|
||||
+2
-1
@@ -10,7 +10,7 @@ from .auth import MULTI_USER
|
||||
from .database import engine
|
||||
from .limits import BodySizeLimitMiddleware
|
||||
from .migrations import bootstrap
|
||||
from .routers import adventures, auth, debug, scenarios, scripts, settings, story_cards
|
||||
from .routers import adventures, auth, chat, debug, scenarios, scripts, settings, story_cards
|
||||
from .seed import seed_public_scenarios
|
||||
|
||||
bootstrap(engine)
|
||||
@@ -88,6 +88,7 @@ app.include_router(adventures.router)
|
||||
app.include_router(story_cards.router)
|
||||
app.include_router(scripts.router)
|
||||
app.include_router(settings.router)
|
||||
app.include_router(chat.router)
|
||||
app.include_router(debug.router)
|
||||
|
||||
|
||||
|
||||
@@ -54,6 +54,12 @@ _running: set[int] = set()
|
||||
_tasks: set[asyncio.Task] = set()
|
||||
|
||||
|
||||
# BYOK-only by construction: both factories below take the user's own
|
||||
# endpoint/key straight from Settings and never auth.DEMO_*, so summarization
|
||||
# and embedding can't spend the shared demo key (their call sites are also
|
||||
# skipped when using_demo). Don't "fix" this by passing a ProviderConfig in —
|
||||
# summary_model/embedding_model are free-form user input and are not on the
|
||||
# demo whitelist.
|
||||
def summary_provider(settings: models.Settings) -> OpenAICompatibleProvider:
|
||||
return OpenAICompatibleProvider(
|
||||
settings.endpoint_url,
|
||||
|
||||
@@ -10,6 +10,17 @@ from .base import PromptParts, Provider, ProviderError
|
||||
# continuing prose instead of replying conversationally.
|
||||
CHAT_CONTINUE_HINT = "\n\n[Continue the story directly. Output only story text.]"
|
||||
|
||||
# Completion endpoints have no roles, so a plain chat has to be flattened into
|
||||
# one labelled transcript that trails off on "Assistant:" for the model to continue.
|
||||
_ROLE_LABELS = {"system": "System", "user": "User", "assistant": "Assistant"}
|
||||
|
||||
|
||||
def flatten_messages(messages: list[dict]) -> str:
|
||||
turns = "\n\n".join(
|
||||
f"{_ROLE_LABELS.get(m['role'], m['role'])}: {m['content']}" for m in messages
|
||||
)
|
||||
return f"{turns}\n\nAssistant:"
|
||||
|
||||
|
||||
class OpenAICompatibleProvider(Provider):
|
||||
"""Adapter for any /v1-style endpoint: Ollama, LM Studio, OpenAI, OpenRouter, vLLM, Groq…"""
|
||||
@@ -110,7 +121,42 @@ class OpenAICompatibleProvider(Provider):
|
||||
if not self.model:
|
||||
raise ProviderError("No model configured — set one in Settings.")
|
||||
url, body = self._request(parts, temperature, max_tokens)
|
||||
async for event in self._stream(url, body):
|
||||
yield event
|
||||
|
||||
async def chat(
|
||||
self, messages: list[dict], *, temperature: float, max_tokens: int
|
||||
) -> AsyncIterator[tuple[str, str]]:
|
||||
"""Plain multi-turn chat — no story framing, no context assembly. Takes
|
||||
[{"role", "content"}, ...] straight to the endpoint. Used by the AI Chat
|
||||
scratchpad; the turn engine uses generate()."""
|
||||
if not self.model:
|
||||
raise ProviderError("No model configured — set one in Settings.")
|
||||
if self.api_mode == "completion":
|
||||
url = f"{self.base_url}/completions"
|
||||
body = {
|
||||
"model": self.model,
|
||||
"prompt": flatten_messages(messages),
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
"stream": True,
|
||||
}
|
||||
else:
|
||||
url = f"{self.base_url}/chat/completions"
|
||||
body = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
"stream": True,
|
||||
}
|
||||
self._apply_reasoning_budget(body)
|
||||
async for event in self._stream(url, body):
|
||||
yield event
|
||||
|
||||
async def _stream(self, url: str, body: dict) -> AsyncIterator[tuple[str, str]]:
|
||||
"""Shared SSE plumbing for generate()/chat(): POST a streaming request
|
||||
and yield ("text" | "reasoning", chunk) pairs, logging the exchange."""
|
||||
log = debuglog.start_entry(url, self.model, body)
|
||||
received: list[str] = []
|
||||
try:
|
||||
|
||||
@@ -32,6 +32,8 @@ def me_payload(user: models.User, db: Session) -> dict:
|
||||
"id": user.id,
|
||||
"email": user.email,
|
||||
"is_guest": user.is_guest,
|
||||
# Trusted testers: unmetered demo turns, plus the AI Chat scratchpad.
|
||||
"power_user": auth.is_power_user(user),
|
||||
"demo": {
|
||||
"enabled": auth.demo_enabled(),
|
||||
"using_demo": cfg.using_demo,
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
"""AI Chat — a plain scratchpad for talking to a model directly.
|
||||
|
||||
Power users only (the AIDND_POWER_USERS email allowlist). Deliberately thin:
|
||||
no story context, no scripts, no world state, and nothing persisted — the
|
||||
conversation lives in the browser and is posted up whole on each turn. It
|
||||
exists to poke at models, prompts and endpoints without starting an adventure.
|
||||
|
||||
Model choice is free-form when the user brought their own API key. On the
|
||||
shared demo key it stays pinned to the AIDND_DEMO_MODELS whitelist, exactly as
|
||||
turns are: the server funds that key, so it must not be able to reach paid
|
||||
models by way of this page.
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from .. import auth, limits, models, schemas
|
||||
from ..database import get_db
|
||||
from ..providers import OpenAICompatibleProvider, ProviderError
|
||||
from .adventures import SSE_HEADERS, sse
|
||||
from .settings import get_settings, list_endpoint_models
|
||||
|
||||
router = APIRouter(prefix="/api/chat", tags=["chat"])
|
||||
|
||||
|
||||
def power_user(
|
||||
db: Session = Depends(get_db),
|
||||
user: models.User = Depends(auth.get_current_user),
|
||||
) -> models.User:
|
||||
"""Gate for the whole router. 404 rather than 403 so the feature simply
|
||||
doesn't appear to exist for everyone else."""
|
||||
if not auth.is_power_user(user):
|
||||
raise HTTPException(404, "Not found")
|
||||
return user
|
||||
|
||||
|
||||
PowerUser = Depends(power_user)
|
||||
|
||||
|
||||
def _resolve_model(
|
||||
settings: models.Settings, requested: str | None
|
||||
) -> tuple[auth.ProviderConfig, str | None]:
|
||||
"""Provider config for this chat, plus a note when the requested model was
|
||||
not honoured. The pinning rule itself lives in resolve_provider_config —
|
||||
this only reports the substitution it made, so there is exactly one place
|
||||
that decides what the demo key is allowed to talk to."""
|
||||
cfg = auth.resolve_provider_config(settings, model_override=requested)
|
||||
wanted = (requested or "").strip()
|
||||
if wanted and wanted != cfg.model:
|
||||
return cfg, (
|
||||
f"'{wanted}' isn't available on the shared demo key — using "
|
||||
f"{cfg.model}. Add your own API key in Settings to use any model."
|
||||
)
|
||||
return cfg, None
|
||||
|
||||
|
||||
@router.get("/config")
|
||||
async def chat_config(
|
||||
db: Session = Depends(get_db),
|
||||
user: models.User = PowerUser,
|
||||
):
|
||||
"""What this page can talk to: the resolved endpoint/model, whether model
|
||||
choice is pinned to the demo whitelist, and the endpoint's model listing
|
||||
(best effort — an unreachable endpoint just yields an empty list)."""
|
||||
settings = get_settings(db, user)
|
||||
cfg = auth.resolve_provider_config(settings)
|
||||
listing = await list_endpoint_models(cfg)
|
||||
return {
|
||||
"endpoint_url": cfg.endpoint_url,
|
||||
"model": cfg.model,
|
||||
"using_demo": cfg.using_demo,
|
||||
"api_mode": settings.api_mode,
|
||||
"temperature": settings.temperature,
|
||||
"max_tokens": settings.max_output_tokens,
|
||||
# On the demo key the whitelist IS the list of choices; otherwise it's
|
||||
# whatever the endpoint advertises (suggestions, not a restriction).
|
||||
"models": auth.DEMO_MODELS if cfg.using_demo else listing.get("models", []),
|
||||
"models_error": None if listing.get("ok") else listing.get("detail"),
|
||||
}
|
||||
|
||||
|
||||
async def run_chat(cfg: auth.ProviderConfig, settings: models.Settings, payload: schemas.ChatRequest,
|
||||
note: str | None, db: Session, user: models.User):
|
||||
"""SSE generator mirroring the turn stream's event shape: reasoning/chunk
|
||||
while generating, then done — so the frontend reuses the same plumbing."""
|
||||
if note:
|
||||
yield sse({"type": "note", "detail": note})
|
||||
provider = OpenAICompatibleProvider(
|
||||
cfg.endpoint_url, cfg.api_key, cfg.model, settings.api_mode,
|
||||
settings.reasoning_max_tokens,
|
||||
)
|
||||
messages = [m.model_dump() for m in payload.messages]
|
||||
chunks: list[str] = []
|
||||
reasoning_chunks: list[str] = []
|
||||
try:
|
||||
async for kind, chunk in provider.chat(
|
||||
messages,
|
||||
temperature=payload.temperature if payload.temperature is not None else settings.temperature,
|
||||
max_tokens=payload.max_tokens or settings.max_output_tokens,
|
||||
):
|
||||
if kind == "reasoning":
|
||||
reasoning_chunks.append(chunk)
|
||||
yield sse({"type": "reasoning", "text": chunk})
|
||||
else:
|
||||
chunks.append(chunk)
|
||||
yield sse({"type": "chunk", "text": chunk})
|
||||
except ProviderError as exc:
|
||||
yield sse({"type": "error", "detail": str(exc)})
|
||||
return
|
||||
|
||||
text = "".join(chunks).strip()
|
||||
if not text:
|
||||
detail = (
|
||||
"The model used its entire token budget on reasoning and returned no "
|
||||
"reply — raise max tokens, cap the reasoning budget in Settings, or "
|
||||
"use a non-reasoning model."
|
||||
if reasoning_chunks
|
||||
else "The AI returned an empty response."
|
||||
)
|
||||
yield sse({"type": "error", "detail": detail})
|
||||
return
|
||||
|
||||
if cfg.using_demo:
|
||||
# Unmetered for power users (count_demo_turn is a no-op for them), but
|
||||
# keep the call so the accounting stays right if the gate ever widens.
|
||||
auth.count_demo_turn(user)
|
||||
db.commit()
|
||||
yield sse({
|
||||
"type": "done",
|
||||
"text": text,
|
||||
"reasoning": "".join(reasoning_chunks).strip() or None,
|
||||
"model": cfg.model,
|
||||
})
|
||||
|
||||
|
||||
@router.post("/stream")
|
||||
def chat_stream(
|
||||
payload: schemas.ChatRequest,
|
||||
request: Request,
|
||||
db: Session = Depends(get_db),
|
||||
user: models.User = PowerUser,
|
||||
):
|
||||
total = sum(len(m.content) for m in payload.messages)
|
||||
if total > schemas.CHAT_TOTAL_MAX:
|
||||
raise HTTPException(
|
||||
413, f"This conversation is too long to send ({total:,} characters) — "
|
||||
"clear it or start a new one."
|
||||
)
|
||||
limits.rate_limit("chat", request, user)
|
||||
settings = get_settings(db, user)
|
||||
cfg, note = _resolve_model(settings, payload.model)
|
||||
if not cfg.model:
|
||||
raise HTTPException(400, "No model configured — set one in Settings or pick one here.")
|
||||
return StreamingResponse(
|
||||
run_chat(cfg, settings, payload, note, db, user),
|
||||
media_type="text/event-stream",
|
||||
headers=SSE_HEADERS,
|
||||
)
|
||||
@@ -62,18 +62,9 @@ def update_settings(
|
||||
return settings
|
||||
|
||||
|
||||
@router.post("/test")
|
||||
async def test_connection(
|
||||
request: Request,
|
||||
db: Session = Depends(get_db),
|
||||
user: models.User = Depends(auth.get_current_user),
|
||||
):
|
||||
"""Hit the endpoint's /models listing as a cheap connectivity check.
|
||||
Tests whatever the turn engine would actually use — including the shared
|
||||
demo endpoint when the user has no key of their own."""
|
||||
limits.rate_limit("connection-test", request, user)
|
||||
settings = get_settings(db, user)
|
||||
cfg = auth.resolve_provider_config(settings)
|
||||
async def list_endpoint_models(cfg: auth.ProviderConfig) -> dict:
|
||||
"""GET the endpoint's /models listing. Doubles as a connectivity check, so
|
||||
failures come back as {"ok": False, "detail": ...} rather than raising."""
|
||||
url = cfg.endpoint_url.rstrip("/") + "/models"
|
||||
headers = {}
|
||||
if cfg.api_key:
|
||||
@@ -94,3 +85,16 @@ async def test_connection(
|
||||
except (ValueError, AttributeError, TypeError):
|
||||
pass # non-JSON or unexpected shape — connectivity is still confirmed
|
||||
return {"ok": True, "models": models_available}
|
||||
|
||||
|
||||
@router.post("/test")
|
||||
async def test_connection(
|
||||
request: Request,
|
||||
db: Session = Depends(get_db),
|
||||
user: models.User = Depends(auth.get_current_user),
|
||||
):
|
||||
"""Cheap connectivity check against whatever the turn engine would actually
|
||||
use — including the shared demo endpoint when the user has no key."""
|
||||
limits.rate_limit("connection-test", request, user)
|
||||
settings = get_settings(db, user)
|
||||
return await list_endpoint_models(auth.resolve_provider_config(settings))
|
||||
|
||||
@@ -319,6 +319,28 @@ class SettingsOut(ORMModel):
|
||||
ScenarioOut.model_rebuild()
|
||||
|
||||
|
||||
# ---------- AI Chat (power users) ----------
|
||||
# A scratchpad for talking to a model directly, with no story framing. Nothing
|
||||
# is persisted server-side, so these caps are purely per-request abuse limits.
|
||||
|
||||
CHAT_MESSAGE_MAX = 100_000 # one message
|
||||
CHAT_TOTAL_MAX = 400_000 # whole conversation sent up per request
|
||||
CHAT_MESSAGES_MAX = 200 # turns per request
|
||||
|
||||
|
||||
class ChatMessage(BaseModel):
|
||||
role: Literal["system", "user", "assistant"]
|
||||
content: Annotated[str, Field(max_length=CHAT_MESSAGE_MAX)]
|
||||
|
||||
|
||||
class ChatRequest(BaseModel):
|
||||
messages: Annotated[list[ChatMessage], Field(min_length=1, max_length=CHAT_MESSAGES_MAX)]
|
||||
# Empty/omitted = fall back to the user's configured model.
|
||||
model: Name | None = None
|
||||
temperature: Annotated[float, Field(ge=0, le=5)] | None = None
|
||||
max_tokens: Annotated[int, Field(ge=1, le=100_000)] | None = None
|
||||
|
||||
|
||||
class SettingsUpdate(BaseModel):
|
||||
endpoint_url: Annotated[str, Field(max_length=500)] | None = None # VARCHAR(500)
|
||||
# Encryption expands the stored value ~4/3 into the same VARCHAR(500):
|
||||
|
||||
Reference in New Issue
Block a user