From 51c9b029817af87421f4e686b0512d19556db044 Mon Sep 17 00:00:00 2001 From: coderconnoisseur <142653086+coderconnoisseur@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:06:12 +0530 Subject: [PATCH 01/10] BYOK: model catalogue per provider and role; OpenAI and Anthropic adapters --- backend/llm/client.py | 5 +- backend/llm/providers.py | 54 +++++++++++++++++++--- backend/llm/registry.py | 13 ++++-- backend/llm/registry.yaml | 25 ++++++++++ backend/tests/test_byok_catalogue.py | 69 ++++++++++++++++++++++++++++ backend/tests/test_llm_providers.py | 2 +- 6 files changed, 155 insertions(+), 13 deletions(-) create mode 100644 backend/tests/test_byok_catalogue.py diff --git a/backend/llm/client.py b/backend/llm/client.py index d321882..996af8f 100644 --- a/backend/llm/client.py +++ b/backend/llm/client.py @@ -214,12 +214,13 @@ def _call_once(entry: Entry, messages, request, timeout, attempt: int = 0, key: if request[name] is not None: payload[name] = request[name] payload.update(provider.reasoning_params(entry.model, request["reasoning"])) + payload = provider.adapt(payload) slice_s = call_timeout(timeout) start = time.monotonic() try: resp = run_capped( - lambda: requests.post(provider.chat_url, headers=provider.headers(), json=payload, timeout=slice_s), + lambda: requests.post(provider.chat_url, headers=provider.headers(entry.key), json=payload, timeout=slice_s), seconds=slice_s, ) except requests.Timeout as e: @@ -256,7 +257,7 @@ def _parse(provider, entry: Entry, resp, latency: float, max_tokens: int) -> LLM kind = provider.kind_for_status(status) detail = f"{provider.name} HTTP {status} (model={entry.model}): {error_message(body) or resp.text[:200]}" if kind == LLMError.AUTH: - detail += provider.auth_hint() + detail += provider.auth_hint() if entry.key is None else " Check your key in Settings." logger.error("LLM call failed [%s]: %s after %.1fs", kind, detail, latency) err = LLMError(kind, detail, status=status) err.retry_after = _retry_after(resp) diff --git a/backend/llm/providers.py b/backend/llm/providers.py index b88dffd..a1985af 100644 --- a/backend/llm/providers.py +++ b/backend/llm/providers.py @@ -65,9 +65,9 @@ def chat_url(self) -> str: def api_key(self) -> str | None: return os.getenv(self.key_env) if self.key_env else None - def headers(self) -> dict: + def headers(self, key: str | None = None) -> dict: headers = {"Content-Type": "application/json"} - key = self.api_key() + key = key or self.api_key() if key: headers["Authorization"] = f"Bearer {key}" return headers @@ -83,10 +83,14 @@ def reasoning_params(self, model: str, reasoning: dict | None) -> dict: def auth_hint(self) -> str: return f" Check {self.key_env} in backend/.env." if self.key_env else "" - def check_key(self, timeout: float = 8) -> tuple[bool, str]: + def adapt(self, payload: dict) -> dict: + """Drop or rename request fields this provider rejects. Default: unchanged.""" + return payload + + def check_key(self, timeout: float = 8, key: str | None = None) -> tuple[bool, str]: """Ask the provider whether the key works. Costs no generation quota.""" try: - resp = requests.get(f"{self.base_url}/models", headers=self.headers(), timeout=timeout) + resp = requests.get(f"{self.base_url}/models", headers=self.headers(key), timeout=timeout) except requests.RequestException as e: return False, f"could not reach {self.name} ({type(e).__name__})." if resp.status_code != 200: @@ -114,11 +118,11 @@ def auth_hint(self) -> str: ) return " Check OPENROUTER_API_KEY in backend/.env." - def check_key(self, timeout: float = 8) -> tuple[bool, str]: - if not self.api_key(): + def check_key(self, timeout: float = 8, key: str | None = None) -> tuple[bool, str]: + if not (key or self.api_key()): return False, "OPENROUTER_API_KEY is not set." try: - resp = requests.get(f"{self.base_url}/key", headers=self.headers(), timeout=timeout) + resp = requests.get(f"{self.base_url}/key", headers=self.headers(key), timeout=timeout) except requests.RequestException as e: return False, f"could not reach OpenRouter to verify the key ({type(e).__name__})." if resp.status_code != 200: @@ -154,7 +158,43 @@ def reasoning_params(self, model, reasoning): return {"think": False} if reasoning == REASONING_OFF else {} +class _OpenAI(Provider): + # GPT-5+ are reasoning models: temperature other than 1 and `max_tokens` are rejected (checked 2026-10-16). + def adapt(self, payload): + payload.pop("temperature", None) + payload["max_completion_tokens"] = payload.pop("max_tokens") + return payload + + def reasoning_params(self, model, reasoning): + return {"reasoning_effort": "low"} if reasoning == REASONING_OFF else {} + + +class _Anthropic(Provider): + # OpenAI-compatibility layer: temperature < 1 is a 400 on Claude 4.7+; seed/response_format are ignored. + def adapt(self, payload): + payload.pop("temperature", None) + payload.pop("seed", None) + return payload + + def reasoning_params(self, model, reasoning): + always_on = "opus" in model or "fable" in model + return {"thinking": {"type": "disabled"}} if reasoning == REASONING_OFF and not always_on else {} + + def check_key(self, timeout: float = 8, key: str | None = None) -> tuple[bool, str]: + try: + resp = requests.get(f"{self.base_url}/models", timeout=timeout, + headers={"x-api-key": key or "", "anthropic-version": "2023-06-01"}) + except requests.RequestException as e: + return False, f"could not reach anthropic ({type(e).__name__})." + if resp.status_code != 200: + return False, f"anthropic rejected the key (HTTP {resp.status_code})." + return True, f"anthropic key OK ({len((resp.json() or {}).get('data') or [])} models available)." + + PROVIDERS: dict[str, Provider] = { + # No platform key for these two: users supply them (SPEC-BYOK). + "openai": _OpenAI("openai", "https://api.openai.com/v1", "OPENAI_BASE_URL", None), + "anthropic": _Anthropic("anthropic", "https://api.anthropic.com/v1", "ANTHROPIC_BASE_URL", None), "openrouter": _OpenRouter("openrouter", "https://openrouter.ai/api/v1", "OPENROUTER_BASE_URL", "OPENROUTER_API_KEY"), "groq": _Groq("groq", "https://api.groq.com/openai/v1", "GROQ_BASE_URL", "GROQ_API_KEY"), "ollama": _Ollama("ollama", "http://localhost:11434/v1", "OLLAMA_BASE_URL", None), diff --git a/backend/llm/registry.py b/backend/llm/registry.py index fdceeaf..f7d9883 100644 --- a/backend/llm/registry.py +++ b/backend/llm/registry.py @@ -2,13 +2,14 @@ from __future__ import annotations import os -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import Mapping import yaml -ROLES = ("text", "vision", "judge") +ROLES = ("text", "vision", "judge", "agent") # agent rows exist only under byok; the platform agent is `text` +BYOK_PROVIDERS = ("openai", "anthropic", "gemini", "groq", "openrouter") TIERS = ("openrouter", "groq", "ollama") _DEFAULT_PATH = Path(__file__).with_name("registry.yaml") @@ -28,6 +29,7 @@ class Entry: provider: str model: str family: str + key: str | None = field(default=None, repr=False, compare=False) # a user's key (SPEC-BYOK); None = platform env key def guess_family(model: str) -> str: @@ -74,6 +76,11 @@ def chain(self, tier: str, role: str, env: Mapping[str, str] | None = None) -> l unique.append(entry) return unique + def byok(self, provider: str, role: str) -> list[Entry]: + rows = ((self._data.get("byok") or {}).get(provider) or {}).get(role) or [] + models = [r if isinstance(r, str) else r["model"] for r in rows] + return [Entry(provider, m, guess_family(m)) for m in models] + def check(self, tier: str, env: Mapping[str, str] | None = None) -> list[str]: """Raise on a judge that shares the text model's family; return warnings.""" env = os.environ if env is None else env @@ -85,7 +92,7 @@ def check(self, tier: str, env: Mapping[str, str] | None = None) -> list[str]: f"{text[0].model!r}; a model must not grade its own family" ) warnings = [] - for role in ROLES: + for role in ("text", "vision", "judge"): for entry in self.chain(tier, role, env): if entry.model == "openrouter/free": warnings.append( diff --git a/backend/llm/registry.yaml b/backend/llm/registry.yaml index 5dae341..728ac37 100644 --- a/backend/llm/registry.yaml +++ b/backend/llm/registry.yaml @@ -41,3 +41,28 @@ ollama: - {model: "qwen2.5:3b", family: qwen} vision: [] judge: [] + +# BYOK catalogue (SPEC-BYOK): the only models a user's key may run, per role, best first. Users pick a provider; +# Lumen takes the first model (or the user's pick, if it is still listed). Ids checked 2026-10-16 against each +# provider's model docs and OpenRouter's /models. A retired id: delete the line. +byok: + openai: + agent: [gpt-6.1-sol, gpt-5.6-sol] + text: [gpt-6-luna, gpt-5.4-mini] + vision: [gpt-6-luna, gpt-5.4-mini] + anthropic: + agent: [claude-sonnet-5-5, claude-opus-5-5] # Sonnet first: its thinking can be switched off (Opus is always on) + text: [claude-haiku-5-5] + vision: [claude-haiku-5-5, claude-sonnet-5-5] + gemini: + agent: [gemini-4-argon, gemini-3.8-flash] + text: [gemini-3.5-flash-lite, gemini-3.8-flash] + vision: [gemini-3.1-flash-lite, gemini-3.8-flash] + groq: + agent: [openai/gpt-oss-120b] + text: [openai/gpt-oss-120b, openai/gpt-oss-20b] + vision: [] # no production vision model on Groq: vision falls back to another key, else the platform + openrouter: + agent: [anthropic/claude-sonnet-5.5, openai/gpt-6.1-sol] + text: [google/gemini-3.5-flash-lite, openai/gpt-oss-120b] + vision: [google/gemini-3.1-flash-lite, anthropic/claude-haiku-5.5] diff --git a/backend/tests/test_byok_catalogue.py b/backend/tests/test_byok_catalogue.py new file mode 100644 index 0000000..6aaf464 --- /dev/null +++ b/backend/tests/test_byok_catalogue.py @@ -0,0 +1,69 @@ +from llm import client +from llm.providers import PROVIDERS, REASONING_OFF +from llm.registry import BYOK_PROVIDERS, Registry + + +def test_catalogue_covers_every_provider_and_role(): + reg = Registry.load() + for p in BYOK_PROVIDERS: + assert reg.byok(p, "agent") and reg.byok(p, "text"), p + assert all(e.provider == p for e in reg.byok(p, "agent")) + assert reg.byok("groq", "vision") == [] # no Groq vision model: falls back + assert reg.byok("nope", "agent") == [] + + +def test_platform_chain_unchanged(): + reg = Registry.load() + # snapshot of today's chains (copy the real output of the current code into these literals before editing) + assert [(e.provider, e.model) for e in reg.chain("groq", "text", env={})] == [ + ("groq", "openai/gpt-oss-120b"), ("groq", "qwen/qwen3.8-27b"), ("groq", "openai/gpt-oss-20b"), + ("openrouter", "nvidia/nemotron-3-ultra-550b-a55b:free"), ("openrouter", "qwen/qwen3.8-27b:free")] + assert all(e.key is None for e in reg.chain("groq", "vision", env={})) + + +def test_entry_repr_hides_key(): + e = Registry.load().byok("openai", "agent")[0] + from dataclasses import replace + assert "sk-secret" not in repr(replace(e, key="sk-secret")) + + +def test_adapt_payloads(): + base = {"model": "m", "messages": [], "temperature": 0.0, "max_tokens": 300, "seed": 1} + oa = PROVIDERS["openai"].adapt(dict(base)) + assert "temperature" not in oa and "max_tokens" not in oa and oa["max_completion_tokens"] == 300 + an = PROVIDERS["anthropic"].adapt(dict(base)) + assert "temperature" not in an and "seed" not in an and an["max_tokens"] == 300 + assert PROVIDERS["groq"].adapt(dict(base)) == base + assert PROVIDERS["anthropic"].reasoning_params("claude-haiku-5-5", REASONING_OFF) == {"thinking": {"type": "disabled"}} + assert PROVIDERS["anthropic"].reasoning_params("claude-opus-5-5", REASONING_OFF) == {} + assert PROVIDERS["openai"].reasoning_params("gpt-6-luna", REASONING_OFF) == {"reasoning_effort": "low"} + + +def test_user_key_goes_in_the_header(monkeypatch): + sent = {} + + class R: + status_code, headers = 200, {} + def json(self): return {"choices": [{"message": {"content": "hi"}, "finish_reason": "stop"}]} + + monkeypatch.setattr(client.requests, "post", lambda url, headers, json, timeout: sent.update(h=headers, j=json, u=url) or R()) + from dataclasses import replace + entry = replace(Registry.load().byok("anthropic", "text")[0], key="sk-ant-user") + client.complete([{"role": "user", "content": "x"}], chain=[entry]) + assert sent["h"]["Authorization"] == "Bearer sk-ant-user" + assert sent["u"] == "https://api.anthropic.com/v1/chat/completions" + assert "temperature" not in sent["j"] + + +def test_check_key_uses_given_key(monkeypatch): + seen = {} + + class R: + status_code = 200 + def json(self): return {"data": [{}, {}]} + + monkeypatch.setattr("llm.providers.requests.get", lambda url, headers, timeout: seen.update(u=url, h=headers) or R()) + ok, _ = PROVIDERS["anthropic"].check_key(key="sk-ant-x") + assert ok and seen["h"]["x-api-key"] == "sk-ant-x" and seen["u"].endswith("/v1/models") + ok, _ = PROVIDERS["openai"].check_key(key="sk-x") + assert ok and seen["h"]["Authorization"] == "Bearer sk-x" diff --git a/backend/tests/test_llm_providers.py b/backend/tests/test_llm_providers.py index 9b7f363..c1c6145 100644 --- a/backend/tests/test_llm_providers.py +++ b/backend/tests/test_llm_providers.py @@ -19,7 +19,7 @@ def json(self): def test_providers_are_registered(): from llm.providers import PROVIDERS - assert set(PROVIDERS) == {"openrouter", "groq", "ollama", "gemini"} + assert set(PROVIDERS) == {"openrouter", "groq", "ollama", "gemini", "openai", "anthropic"} assert PROVIDERS["groq"].base_url == "https://api.groq.com/openai/v1" assert PROVIDERS["ollama"].base_url == "http://localhost:11434/v1" From 5ed66c6f6ba6a28db8db37e61fa5d0c0f21c7423 Mon Sep 17 00:00:00 2001 From: coderconnoisseur <142653086+coderconnoisseur@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:07:38 +0530 Subject: [PATCH 02/10] BYOK: OpenAI and Anthropic base-URL overrides use LUMEN_ prefix --- backend/llm/providers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/backend/llm/providers.py b/backend/llm/providers.py index a1985af..362c97d 100644 --- a/backend/llm/providers.py +++ b/backend/llm/providers.py @@ -193,8 +193,8 @@ def check_key(self, timeout: float = 8, key: str | None = None) -> tuple[bool, s PROVIDERS: dict[str, Provider] = { # No platform key for these two: users supply them (SPEC-BYOK). - "openai": _OpenAI("openai", "https://api.openai.com/v1", "OPENAI_BASE_URL", None), - "anthropic": _Anthropic("anthropic", "https://api.anthropic.com/v1", "ANTHROPIC_BASE_URL", None), + "openai": _OpenAI("openai", "https://api.openai.com/v1", "LUMEN_OPENAI_BASE_URL", None), + "anthropic": _Anthropic("anthropic", "https://api.anthropic.com/v1", "LUMEN_ANTHROPIC_BASE_URL", None), "openrouter": _OpenRouter("openrouter", "https://openrouter.ai/api/v1", "OPENROUTER_BASE_URL", "OPENROUTER_API_KEY"), "groq": _Groq("groq", "https://api.groq.com/openai/v1", "GROQ_BASE_URL", "GROQ_API_KEY"), "ollama": _Ollama("ollama", "http://localhost:11434/v1", "OLLAMA_BASE_URL", None), From b6f16833fa122b44e261b5d5f507a17bac7a4aa6 Mon Sep 17 00:00:00 2001 From: coderconnoisseur <142653086+coderconnoisseur@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:11:56 +0530 Subject: [PATCH 03/10] BYOK: encrypted user keys; Ask Lumen, invoices and email use the user's provider; rejected keys are reported --- backend/agent/graph.py | 4 +- backend/api/agent.py | 4 +- backend/api/documents.py | 4 +- backend/api/errors.py | 4 +- backend/llm/client.py | 19 +++- backend/llm/errors.py | 4 +- backend/llm/keys.py | 121 +++++++++++++++++++++ backend/models/__init__.py | 16 +++ backend/tests/test_byok_client.py | 174 ++++++++++++++++++++++++++++++ backend/utils/auth.py | 6 +- backend/utils/crypto.py | 16 +-- backend/utils/errors.py | 19 +++- 12 files changed, 371 insertions(+), 20 deletions(-) create mode 100644 backend/llm/keys.py create mode 100644 backend/tests/test_byok_client.py diff --git a/backend/agent/graph.py b/backend/agent/graph.py index d9d1d56..1213f52 100644 --- a/backend/agent/graph.py +++ b/backend/agent/graph.py @@ -93,8 +93,8 @@ def agent(state: AgentState) -> dict: out_of_calls = state["tool_calls"] >= max_tool_calls if out_of_calls: messages.append({"role": "user", "content": FORCE_ANSWER}) - result = complete(messages, tools=None if out_of_calls else schemas, temperature=0, max_tokens=900, - timeout=40, retries=1) + result = complete(messages, role="agent", tools=None if out_of_calls else schemas, temperature=0, + max_tokens=900, timeout=40, retries=1) calls = [] if out_of_calls else (result.tool_calls or []) for call in calls: call.setdefault("id", f"call_{uuid.uuid4().hex[:12]}") diff --git a/backend/api/agent.py b/backend/api/agent.py index 4b3e9bf..67712d2 100644 --- a/backend/api/agent.py +++ b/backend/api/agent.py @@ -19,6 +19,7 @@ from sqlalchemy.orm import Session from api.deps import current_user, rate_limit +from llm import keys router = APIRouter(prefix="/api/agent", tags=["agent"]) @@ -84,7 +85,8 @@ def ask(body: Question, claims: dict = Depends(current_user), deps: AgentDeps = ctx = ToolContext(user_id=claims["sub"], engine=deps.engine, sql_agent=deps.sql_agent, rag=deps.rag, today=deps.today or date.today(), profile=profile_for(deps.engine, claims["sub"]), question=body.question) - out = run_agent(body.question, ctx, complete=deps.complete) + with keys.for_user(claims["sub"], deps.engine): + out = run_agent(body.question, ctx, complete=deps.complete) return {"success": True, **out} diff --git a/backend/api/documents.py b/backend/api/documents.py index 26ccadc..e4beaa4 100644 --- a/backend/api/documents.py +++ b/backend/api/documents.py @@ -7,6 +7,7 @@ from api.deps import current_user, rate_limit from config import Config +from llm import keys from rag.service import DOC_TYPES, RagService, get_service from utils import files @@ -67,6 +68,7 @@ def search(body: Query, claims: dict = Depends(current_user), rag: RagService = @router.post("/ask", dependencies=[Depends(rate_limit("10 per minute"))]) def ask(body: Question, claims: dict = Depends(current_user), rag: RagService = Depends(get_service)): - out = rag.ask(claims["sub"], body.question) + with keys.for_user(claims["sub"], rag.store.engine): + out = rag.ask(claims["sub"], body.question) return {"success": True, "answer": out["answer"], "abstained": out["abstained"], "reason": out["reason"], "sources": out["sources"]} diff --git a/backend/api/errors.py b/backend/api/errors.py index 171e23b..e2ec408 100644 --- a/backend/api/errors.py +++ b/backend/api/errors.py @@ -9,7 +9,7 @@ from starlette.exceptions import HTTPException from api.deps import AuthFailed -from utils.errors import llm_error_status +from utils.errors import llm_error_status, llm_message from utils.llm import LLMError logger = logging.getLogger(__name__) @@ -44,7 +44,7 @@ async def invalid_request(_request: Request, e: RequestValidationError): async def llm_error(request: Request, e: LLMError): status, code, key = llm_error_status(e) logger.error("%s: LLM %s: %s", request.url.path, e.kind, e.detail) - return _error(LLM_MESSAGES[key], status, code) + return _error(llm_message(e, key, LLM_MESSAGES), status, code) @app.exception_handler(Exception) async def unexpected(request: Request, e: Exception): diff --git a/backend/llm/client.py b/backend/llm/client.py index 996af8f..b9e00da 100644 --- a/backend/llm/client.py +++ b/backend/llm/client.py @@ -25,7 +25,7 @@ import requests -from llm import cassette +from llm import cassette, keys from llm.deadline import MARGIN_SECONDS, MIN_CALL_SECONDS, call_timeout, remaining, run_capped from llm.errors import LLMError from llm.providers import PROVIDERS, REASONING_OFF, error_message @@ -89,13 +89,17 @@ def complete( reasoning: dict | None = REASONING_OFF, seed: int | None = None, ) -> LLMResult: - entries = chain if chain is not None else registry().chain(tier(), role) + if chain is not None: + entries = chain + else: + entries = keys.user_chain(role) or registry().chain(tier(), "text" if role == "agent" else role) if not entries: raise LLMError(LLMError.CONFIG, f"no models configured for role {role!r} in tier {tier()!r}") request = dict(tools=tools, response_format=response_format, temperature=temperature, max_tokens=max_tokens, reasoning=reasoning, seed=seed) dead_providers: set[str] = set() last: LLMError | None = None + last_entry: Entry | None = None first_miss: cassette.CassetteMiss | None = None if cassette.mode() == "record": index = _recorded_entry(entries, messages, request) @@ -113,7 +117,10 @@ def complete( first_miss = first_miss or miss continue except LLMError as e: - last = e + if entry.key: + e.detail = keys.mask(e.detail, entry.key) + e.args = (e.detail,) + last, last_entry = e, entry if e.kind in (LLMError.DEADLINE, LLMError.BAD_RESPONSE): raise if e.kind in (LLMError.AUTH, LLMError.CREDITS): @@ -123,7 +130,12 @@ def complete( calls = _collector.get() if calls is not None: calls.append(result) + if entry.key: + keys.note_ok(entry.provider) return result + if last is not None and last_entry.key and last.kind in keys.KEY_ERRORS: + last.user_key = last_entry.provider + keys.mark_failed(last_entry.provider, last.detail) raise first_miss or last @@ -258,6 +270,7 @@ def _parse(provider, entry: Entry, resp, latency: float, max_tokens: int) -> LLM detail = f"{provider.name} HTTP {status} (model={entry.model}): {error_message(body) or resp.text[:200]}" if kind == LLMError.AUTH: detail += provider.auth_hint() if entry.key is None else " Check your key in Settings." + detail = keys.mask(detail, entry.key or "") logger.error("LLM call failed [%s]: %s after %.1fs", kind, detail, latency) err = LLMError(kind, detail, status=status) err.retry_after = _retry_after(resp) diff --git a/backend/llm/errors.py b/backend/llm/errors.py index deb533c..d8c279f 100644 --- a/backend/llm/errors.py +++ b/backend/llm/errors.py @@ -12,16 +12,18 @@ class LLMError(RuntimeError): UNAVAILABLE = "unavailable" # timeout, network error, provider 5xx BAD_RESPONSE = "bad_response" # 2xx but nothing usable in it DEADLINE = "deadline" # the request's time budget ran out before the call + NEEDS_KEY = "needs_key" # free allowance used up and the user has no key of their own (SPEC-BYOK) # These fail every later call in the same request too, so callers must not # paper over them with a fallback. - FATAL = frozenset({AUTH, CREDITS, RATE_LIMITED, CONFIG, UNAVAILABLE, DEADLINE}) + FATAL = frozenset({AUTH, CREDITS, RATE_LIMITED, CONFIG, UNAVAILABLE, DEADLINE, NEEDS_KEY}) def __init__(self, kind: str, detail: str, status: int | None = None): super().__init__(detail) self.kind = kind self.detail = detail self.status = status + self.user_key: str | None = None # the provider, when the final failure was the user's own key @property def is_fatal(self) -> bool: diff --git a/backend/llm/keys.py b/backend/llm/keys.py new file mode 100644 index 0000000..4a5f1f7 --- /dev/null +++ b/backend/llm/keys.py @@ -0,0 +1,121 @@ +"""Users' own model keys (SPEC-BYOK): who the current request is for, and the chain their keys allow.""" +from __future__ import annotations + +import contextvars +import json +import logging +from contextlib import contextmanager +from dataclasses import replace +from datetime import datetime + +from sqlalchemy import select, update +from sqlalchemy.orm import Session + +from llm.errors import LLMError +from llm.registry import BYOK_PROVIDERS, Entry + +logger = logging.getLogger(__name__) +KEY_ENV = "KEYS_ENCRYPTION_KEY" +KEY_ERRORS = frozenset({LLMError.AUTH, LLMError.CREDITS, LLMError.CONFIG}) +PROVIDER_NAMES = {"openai": "OpenAI", "anthropic": "Anthropic", "gemini": "Gemini", "groq": "Groq", + "openrouter": "OpenRouter"} +_ctx: contextvars.ContextVar[dict | None] = contextvars.ContextVar("llm_user", default=None) + + +@contextmanager +def for_user(user_id: str, engine): + token = _ctx.set({"user_id": str(user_id), "engine": engine, "rows": None, "ok": set()}) + try: + yield + finally: + _ctx.reset(token) + + +def current(): + c = _ctx.get() + return (c["user_id"], c["engine"]) if c else None + + +def _rows() -> list[tuple[str, str, dict]]: + """(provider, decrypted key, {role: model}) for the current user; loaded once per request.""" + from models import UserKey + from utils.crypto import decrypt_secret + + c = _ctx.get() + if c is None: + return [] + if c["rows"] is None: + with Session(c["engine"]) as s: + rows = s.scalars(select(UserKey).where(UserKey.user_id == c["user_id"])).all() + out = [] + for r in rows: + key = decrypt_secret(r.key_enc, key_env=KEY_ENV) + if key.startswith("enc:"): # encryption key changed: unusable, say so instead of sending it + logger.error("Stored %s key for a user can't be decrypted", r.provider) + continue + out.append((r.provider, key, json.loads(r.models or "{}"))) + c["rows"] = sorted(out, key=lambda t: BYOK_PROVIDERS.index(t[0]) if t[0] in BYOK_PROVIDERS else 99) + return c["rows"] + + +def user_chain(role: str) -> list[Entry]: + from llm.client import registry + + chain: list[Entry] = [] + for provider, key, picks in _rows(): + listed = registry().byok(provider, role) + pick = picks.get(role) + listed.sort(key=lambda e: e.model != pick) # the user's pick first, if it is still listed + chain += [replace(e, key=key) for e in listed] + return chain + + +def has_key(role: str) -> bool: + return bool(user_chain(role)) + + +def last4(key: str) -> str: + return key[-4:] if key else "" + + +def mask(text: str, key: str) -> str: + if not text or not key: + return text + text = text.replace(key, f"****{last4(key)}") + # ponytail: O(n²) substring scan in the key length (keys are under 200 chars); use a regex if keys grow. + for n in range(len(key) - 1, 7, -1): # providers echo a prefix/suffix of the key + for i in range(0, len(key) - n + 1): + text = text.replace(key[i:i + n], "****") + return text + + +def _set(engine, user_id: str, provider: str, **values) -> None: + from models import UserKey + + with Session(engine) as s, s.begin(): + s.execute(update(UserKey).where(UserKey.user_id == user_id, UserKey.provider == provider).values(**values)) + + +def mark_failed(provider: str, detail: str) -> None: + c = _ctx.get() + if c: + _set(c["engine"], c["user_id"], provider, last_error=detail[:500]) + + +def mark_ok(engine, user_id: str, provider: str) -> None: + _set(engine, user_id, provider, last_ok_at=datetime.utcnow(), last_error=None) + + +def note_ok(provider: str) -> None: + """mark_ok at most once per provider per context, so an agent loop doesn't write on every call.""" + c = _ctx.get() + if c and provider not in c["ok"]: + c["ok"].add(provider) + mark_ok(c["engine"], c["user_id"], provider) + + +def failing(engine, user_id: str): + from models import UserKey + + with Session(engine) as s: + return s.scalars(select(UserKey).where(UserKey.user_id == user_id, UserKey.last_error.is_not(None))).all() diff --git a/backend/models/__init__.py b/backend/models/__init__.py index 3c7a906..926a539 100644 --- a/backend/models/__init__.py +++ b/backend/models/__init__.py @@ -346,3 +346,19 @@ class ProcessedEmail(db.Model): created_at = db.Column(db.DateTime, default=datetime.utcnow) __table_args__ = (db.UniqueConstraint("user_id", "message_id", name="u_processed_email"),) + + +class UserKey(db.Model): + """SPEC-BYOK: a user's own model key for one provider, encrypted with KEYS_ENCRYPTION_KEY.""" + __tablename__ = "user_keys" + + id = db.Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4())) + user_id = db.Column(String(36), nullable=False, index=True) + provider = db.Column(db.String, nullable=False) + key_enc = db.Column(Text, nullable=False) + models = db.Column(Text) # JSON {role: model}: the user's "advanced" pick + created_at = db.Column(db.DateTime, default=datetime.utcnow) + last_ok_at = db.Column(db.DateTime) + last_error = db.Column(Text) + + __table_args__ = (db.UniqueConstraint("user_id", "provider", name="u_user_key"),) diff --git a/backend/tests/test_byok_client.py b/backend/tests/test_byok_client.py new file mode 100644 index 0000000..9b11c7e --- /dev/null +++ b/backend/tests/test_byok_client.py @@ -0,0 +1,174 @@ +"""SPEC-BYOK: stored user keys build the user's chain in llm.client; key failures are flagged and masked.""" +import json +import logging + +import pytest +from sqlalchemy import create_engine, select +from sqlalchemy.orm import Session + +import models # noqa: F401 +from llm import client, keys +from llm.errors import LLMError +from models import UserKey +from models.database import db +from utils.crypto import encrypt_secret + +U = "b0b00000-0000-4000-8000-000000000001" +OK = {"model": "served/m", "choices": [{"message": {"content": "hello"}, "finish_reason": "stop"}]} +MSGS = [{"role": "user", "content": "hi"}] + + +class _Resp: + def __init__(self, status, body): + self.status_code, self._body, self.text, self.headers = status, body, str(body), {} + + def json(self): + return self._body + + +@pytest.fixture +def engine(tmp_path, monkeypatch): + monkeypatch.setenv("KEYS_ENCRYPTION_KEY", "k" * 32) + e = create_engine(f"sqlite:///{tmp_path / 'k.db'}") + db.metadata.create_all(e, tables=[db.metadata.tables["user_keys"]]) + return e + + +@pytest.fixture +def posts(monkeypatch): + """Every POST is recorded; replies come from `script` (a status/body pair repeated when only one is left).""" + calls, script = [], [] + + def post(url, headers, json, timeout): + calls.append({"url": url, "headers": headers, "json": json}) + return script.pop(0) if len(script) > 1 else script[0] + + monkeypatch.setattr(client.requests, "post", post) + monkeypatch.setattr(client, "_sleep", lambda s: None) + return calls, script + + +def add_key(engine, user, provider, key, models=None): + with Session(engine) as s, s.begin(): + s.add(UserKey(user_id=user, provider=provider, key_enc=encrypt_secret(key, key_env=keys.KEY_ENV), + models=json.dumps(models) if models else None)) + + +def row(engine, provider): + with Session(engine) as s: + return s.scalars(select(UserKey).where(UserKey.provider == provider)).one() + + +def test_no_user_context_uses_platform_chain(posts, monkeypatch): + calls, script = posts + script.append(_Resp(200, OK)) + monkeypatch.setenv("LUMEN_LLM_TIER", "ollama") + client.complete(MSGS, role="text") + assert calls[0]["url"].startswith("http://localhost:11434/v1") and calls[0]["json"]["model"] == "qwen2.5:3b" + + +def test_user_key_builds_user_chain(engine, posts): + calls, script = posts + script.append(_Resp(200, OK)) + add_key(engine, U, "openai", "sk-user-openai") + with keys.for_user(U, engine): + client.complete(MSGS, role="agent") + assert calls[0]["url"].startswith("https://api.openai.com") + assert calls[0]["json"]["model"] == "gpt-6.1-sol" + assert calls[0]["headers"]["Authorization"] == "Bearer sk-user-openai" + + +def test_vision_falls_back_for_groq_only_user(engine, posts): + calls, script = posts + script.append(_Resp(200, OK)) + add_key(engine, U, "groq", "gsk-user") + with keys.for_user(U, engine): + client.complete(MSGS, role="vision") + client.complete(MSGS, role="text") + assert calls[0]["headers"].get("Authorization") != "Bearer gsk-user" + assert calls[1]["url"].startswith("https://api.groq.com") + assert calls[1]["json"]["model"] == "openai/gpt-oss-120b" + assert calls[1]["headers"]["Authorization"] == "Bearer gsk-user" + + +def test_two_keys_split_roles(engine): + add_key(engine, U, "groq", "gsk-user") + add_key(engine, U, "gemini", "gem-user") + with keys.for_user(U, engine): + assert [e.model for e in keys.user_chain("agent")] == ["gemini-4-argon", "gemini-3.8-flash", "openai/gpt-oss-120b"] + assert [e.model for e in keys.user_chain("vision")] == [e.model for e in client.registry().byok("gemini", "vision")] + assert keys.has_key("vision") + + +@pytest.mark.parametrize("pick, first", [("gpt-4o", "gpt-6.1-sol"), ("gpt-5.6-sol", "gpt-5.6-sol")]) +def test_retired_preference_uses_first(engine, pick, first): + add_key(engine, U, "openai", "sk-user", models={"agent": pick}) + with keys.for_user(U, engine): + assert keys.user_chain("agent")[0].model == first + + +def test_user_key_auth_error_is_flagged_and_not_failed_over(engine, posts): + calls, script = posts + script.append(_Resp(401, {"error": {"message": "bad key"}})) + add_key(engine, U, "openai", "sk-user-openai") + with keys.for_user(U, engine), pytest.raises(LLMError) as err: + client.complete(MSGS, role="agent") + assert err.value.kind == "auth" and err.value.user_key == "openai" + assert len(calls) == 1 + r = row(engine, "openai") + assert r.last_error and r.last_ok_at is None + + +def test_failed_call_detail_is_masked(engine, posts, caplog): + calls, script = posts + secret = "sk-live-ABCDEFGH1234" + script.append(_Resp(401, {"error": {"message": f"Incorrect API key provided: {secret}"}})) + add_key(engine, U, "openai", secret) + caplog.set_level(logging.DEBUG) + with keys.for_user(U, engine), pytest.raises(LLMError) as err: + client.complete(MSGS, role="agent") + assert "ABCDEFGH" not in str(err.value) + assert "ABCDEFGH" not in row(engine, "openai").last_error + assert not any("ABCDEFGH" in r.getMessage() for r in caplog.records) + + +def test_success_sets_last_ok(engine, posts): + calls, script = posts + script.append(_Resp(200, OK)) + add_key(engine, U, "openai", "sk-user") + keys._set(engine, U, "openai", last_error="old") + with keys.for_user(U, engine): + client.complete(MSGS, role="agent") + r = row(engine, "openai") + assert r.last_ok_at is not None and r.last_error is None + + +def test_llm_error_status_codes(): + from utils.errors import llm_error_status + + e = LLMError(LLMError.AUTH, "x") + e.user_key = "openai" + assert llm_error_status(e) == (400, "key_rejected", "key_rejected") + assert llm_error_status(LLMError("needs_key", "y")) == (402, "needs_key", "needs_key") + assert LLMError.NEEDS_KEY in LLMError.FATAL + + +def test_require_auth_sets_context(monkeypatch): + from flask import Flask + + import utils.auth as auth + from models.database import db as flask_db + + monkeypatch.setattr(auth, "verify_token", lambda token: {"sub": U}) + app = Flask(__name__) + app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite://" + flask_db.init_app(app) + + @app.route("/who") + @auth.require_auth + def who(): + return keys.current()[0] + + resp = app.test_client().get("/who", headers={"Authorization": "Bearer t"}) + assert resp.get_data(as_text=True) == U + assert keys.current() is None diff --git a/backend/utils/auth.py b/backend/utils/auth.py index bc2d1ee..92012d3 100644 --- a/backend/utils/auth.py +++ b/backend/utils/auth.py @@ -218,6 +218,10 @@ def wrapper(*args, **kwargs): g.user_id = claims["sub"] g.user_email = claims.get("email") g.jwt_claims = claims - return fn(*args, **kwargs) + from llm import keys + from models.database import db + + with keys.for_user(claims["sub"], db.engine): + return fn(*args, **kwargs) return wrapper diff --git a/backend/utils/crypto.py b/backend/utils/crypto.py index b4e23a9..e77f132 100644 --- a/backend/utils/crypto.py +++ b/backend/utils/crypto.py @@ -13,8 +13,8 @@ _PREFIX = "enc:" -def _fernet() -> Fernet | None: - raw = os.getenv("EMAIL_ENCRYPTION_KEY") or os.getenv("SECRET_KEY") +def _fernet(key_env: str = "EMAIL_ENCRYPTION_KEY") -> Fernet | None: + raw = os.getenv(key_env) or os.getenv("SECRET_KEY") if not raw: return None # Fernet needs a 32-byte url-safe base64 key; derive deterministically. @@ -23,25 +23,25 @@ def _fernet() -> Fernet | None: return Fernet(key) -def encrypt_secret(plaintext: str) -> str: +def encrypt_secret(plaintext: str, *, key_env: str = "EMAIL_ENCRYPTION_KEY") -> str: if not plaintext: return plaintext - f = _fernet() + f = _fernet(key_env) if f is None: - logger.warning("EMAIL_ENCRYPTION_KEY/SECRET_KEY missing — storing IMAP password in plaintext") + logger.warning("%s/SECRET_KEY missing — storing secret in plaintext", key_env) return plaintext token = f.encrypt(plaintext.encode()).decode() return f"{_PREFIX}{token}" -def decrypt_secret(stored: str) -> str: +def decrypt_secret(stored: str, *, key_env: str = "EMAIL_ENCRYPTION_KEY") -> str: if not stored: return stored if not stored.startswith(_PREFIX): return stored - f = _fernet() + f = _fernet(key_env) if f is None: - logger.warning("Cannot decrypt IMAP password — encryption key missing") + logger.warning("Cannot decrypt secret — encryption key missing") return stored try: return f.decrypt(stored[len(_PREFIX) :].encode()).decode() diff --git a/backend/utils/errors.py b/backend/utils/errors.py index 7e79e15..9b44e85 100644 --- a/backend/utils/errors.py +++ b/backend/utils/errors.py @@ -36,13 +36,30 @@ def llm_api_error(e, messages: dict[str, str], *, context: str = "") -> tuple[An """ status, code, key = llm_error_status(e) logger.error("%sLLM %s: %s", f"{context}: " if context else "", e.kind, e.detail) - return api_error(messages[key], status=status, code=code) + return api_error(llm_message(e, key, messages), status=status, code=code) + + +KEY_MESSAGES = {"key_rejected": "Your {provider} key was rejected. Update it in Settings.", + "needs_key": "You've used Lumen's free allowance. Add your own model key in Settings to keep using Lumen."} + + +def llm_message(e, key: str, messages: dict[str, str]) -> str: + """User-facing text for `llm_error_status`'s message key; the BYOK texts are fixed here.""" + from llm.keys import PROVIDER_NAMES + + if key in KEY_MESSAGES: + return KEY_MESSAGES[key].format(provider=PROVIDER_NAMES.get(getattr(e, "user_key", None), "")) + return messages[key] def llm_error_status(e) -> tuple[int, str, str]: """(HTTP status, error code, message key) for an LLMError; shared with FastAPI.""" from utils.llm import LLMError + if getattr(e, "user_key", None): + return 400, "key_rejected", "key_rejected" + if e.kind == LLMError.NEEDS_KEY: + return 402, "needs_key", "needs_key" if e.kind == LLMError.RATE_LIMITED: return 429, "llm_rate_limited", "rate_limited" if e.kind == LLMError.BAD_RESPONSE: From e065b7dcc751d70986c51c2c09d118739a5c2d74 Mon Sep 17 00:00:00 2001 From: coderconnoisseur <142653086+coderconnoisseur@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:18:06 +0530 Subject: [PATCH 04/10] BYOK: an undecryptable stored key fails without sending anything; only a 403 config error blames the key --- backend/llm/client.py | 4 +++- backend/llm/keys.py | 14 ++++++++++---- backend/tests/test_byok_client.py | 26 ++++++++++++++++++++++++++ 3 files changed, 39 insertions(+), 5 deletions(-) diff --git a/backend/llm/client.py b/backend/llm/client.py index b9e00da..4ce8e6a 100644 --- a/backend/llm/client.py +++ b/backend/llm/client.py @@ -111,6 +111,8 @@ def complete( if entry.provider in dead_providers: continue try: + if entry.key == keys.UNREADABLE: + raise LLMError(LLMError.AUTH, keys.UNREADABLE_DETAIL) result = _complete_entry(entry, messages, request, timeout, retries) except cassette.CassetteMiss as miss: # While recording, this entry may have failed over to the next one: look there too. @@ -133,7 +135,7 @@ def complete( if entry.key: keys.note_ok(entry.provider) return result - if last is not None and last_entry.key and last.kind in keys.KEY_ERRORS: + if last is not None and last_entry.key and keys.is_key_error(last): last.user_key = last_entry.provider keys.mark_failed(last_entry.provider, last.detail) raise first_miss or last diff --git a/backend/llm/keys.py b/backend/llm/keys.py index 4a5f1f7..bff4d40 100644 --- a/backend/llm/keys.py +++ b/backend/llm/keys.py @@ -16,7 +16,8 @@ logger = logging.getLogger(__name__) KEY_ENV = "KEYS_ENCRYPTION_KEY" -KEY_ERRORS = frozenset({LLMError.AUTH, LLMError.CREDITS, LLMError.CONFIG}) +UNREADABLE = "" # stands in for a stored key that can't be decrypted: never sent anywhere +UNREADABLE_DETAIL = "Stored key can't be decrypted; re-enter it in Settings." PROVIDER_NAMES = {"openai": "OpenAI", "anthropic": "Anthropic", "gemini": "Gemini", "groq": "Groq", "openrouter": "OpenRouter"} _ctx: contextvars.ContextVar[dict | None] = contextvars.ContextVar("llm_user", default=None) @@ -50,9 +51,9 @@ def _rows() -> list[tuple[str, str, dict]]: out = [] for r in rows: key = decrypt_secret(r.key_enc, key_env=KEY_ENV) - if key.startswith("enc:"): # encryption key changed: unusable, say so instead of sending it + if key.startswith("enc:"): # encryption key changed: still the user's provider, but never sent logger.error("Stored %s key for a user can't be decrypted", r.provider) - continue + key = UNREADABLE out.append((r.provider, key, json.loads(r.models or "{}"))) c["rows"] = sorted(out, key=lambda t: BYOK_PROVIDERS.index(t[0]) if t[0] in BYOK_PROVIDERS else 99) return c["rows"] @@ -70,6 +71,11 @@ def user_chain(role: str) -> list[Entry]: return chain +def is_key_error(e: LLMError) -> bool: + """The user's key itself is at fault. A CONFIG error counts only on 403: other 400/404s are request problems.""" + return e.kind in (LLMError.AUTH, LLMError.CREDITS) or (e.kind == LLMError.CONFIG and e.status == 403) + + def has_key(role: str) -> bool: return bool(user_chain(role)) @@ -79,7 +85,7 @@ def last4(key: str) -> str: def mask(text: str, key: str) -> str: - if not text or not key: + if not text or not key or key == UNREADABLE: return text text = text.replace(key, f"****{last4(key)}") # ponytail: O(n²) substring scan in the key length (keys are under 200 chars); use a regex if keys grow. diff --git a/backend/tests/test_byok_client.py b/backend/tests/test_byok_client.py index 9b11c7e..8cb172f 100644 --- a/backend/tests/test_byok_client.py +++ b/backend/tests/test_byok_client.py @@ -119,6 +119,32 @@ def test_user_key_auth_error_is_flagged_and_not_failed_over(engine, posts): assert r.last_error and r.last_ok_at is None +def test_undecryptable_key_fails_without_sending_anything(engine, posts, monkeypatch): + calls, script = posts + script.append(_Resp(200, OK)) + monkeypatch.setenv("KEYS_ENCRYPTION_KEY", "other" * 8) + add_key(engine, U, "openai", "sk-user-openai") + monkeypatch.setenv("KEYS_ENCRYPTION_KEY", "k" * 32) + with keys.for_user(U, engine), pytest.raises(LLMError) as err: + assert keys.has_key("agent") + client.complete(MSGS, role="agent") + assert calls == [] + assert err.value.kind == "auth" and err.value.user_key == "openai" + assert row(engine, "openai").last_error == keys.UNREADABLE_DETAIL + + +@pytest.mark.parametrize("status, flagged", [(400, False), (403, True)]) +def test_only_a_403_config_error_blames_the_key(engine, posts, status, flagged): + calls, script = posts + script.append(_Resp(status, {"error": {"message": "nope"}})) + add_key(engine, U, "openai", "sk-user-openai") + with keys.for_user(U, engine), pytest.raises(LLMError) as err: + client.complete(MSGS, role="agent") + assert err.value.kind == "config" + assert (err.value.user_key == "openai") is flagged + assert (row(engine, "openai").last_error is not None) is flagged + + def test_failed_call_detail_is_masked(engine, posts, caplog): calls, script = posts secret = "sk-live-ABCDEFGH1234" From 94b5d77f55954a5e2e687a51a046d3ab05ab7db6 Mon Sep 17 00:00:00 2001 From: coderconnoisseur <142653086+coderconnoisseur@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:23:59 +0530 Subject: [PATCH 05/10] BYOK: add, test and remove model keys (masked, validated, rate-limited) --- backend/api/errors.py | 2 +- backend/api/keys.py | 135 +++++++++++++++++++++++ backend/asgi.py | 5 +- backend/tests/test_api_keys.py | 196 +++++++++++++++++++++++++++++++++ 4 files changed, 335 insertions(+), 3 deletions(-) create mode 100644 backend/api/keys.py create mode 100644 backend/tests/test_api_keys.py diff --git a/backend/api/errors.py b/backend/api/errors.py index e2ec408..ab3780b 100644 --- a/backend/api/errors.py +++ b/backend/api/errors.py @@ -37,7 +37,7 @@ async def http_error(_request: Request, e: HTTPException): @app.exception_handler(RequestValidationError) async def invalid_request(_request: Request, e: RequestValidationError): - logger.info("invalid request: %s", e.errors()) + logger.info("invalid request: %s", [{k: v for k, v in err.items() if k != "input"} for err in e.errors()]) # input may be a secret return _error("Invalid request", 422, "invalid_request") @app.exception_handler(LLMError) diff --git a/backend/api/keys.py b/backend/api/keys.py new file mode 100644 index 0000000..604056e --- /dev/null +++ b/backend/api/keys.py @@ -0,0 +1,135 @@ +"""SPEC-BYOK: a user's own model keys: add (validated, encrypted), list (last 4 only), test and remove.""" +from __future__ import annotations + +import json +from datetime import datetime + +from fastapi import APIRouter, Depends, HTTPException +from fastapi.responses import JSONResponse +from pydantic import BaseModel, ConfigDict, Field +from sqlalchemy import select +from sqlalchemy.engine import Engine +from sqlalchemy.orm import Session + +from api.deps import current_user, rate_limit +from api.profile import get_engine +from llm import keys +from llm.registry import BYOK_PROVIDERS +from utils.crypto import decrypt_secret, encrypt_secret + +router = APIRouter(prefix="/api/keys", tags=["keys"]) +ROLES = ("agent", "text", "vision") + + +class KeyIn(BaseModel): + model_config = ConfigDict(str_strip_whitespace=True) + key: str = Field(min_length=8, max_length=400) + models: dict[str, str] | None = None + + +def _catalogue(provider: str) -> dict[str, list[str]]: + from llm.client import registry + + return {role: [e.model for e in registry().byok(provider, role)] for role in ROLES} + + +def _provider(provider: str) -> str: + if provider not in BYOK_PROVIDERS: + raise HTTPException(404, "Unknown provider.") + return provider + + +def _row(session: Session, user_id: str, provider: str): + from models import UserKey + + return session.scalars(select(UserKey).where(UserKey.user_id == user_id, UserKey.provider == provider)).first() + + +def _key(row) -> str | None: + """The decrypted key, or None when it can't be decrypted (the encryption key changed).""" + key = decrypt_secret(row.key_enc, key_env=keys.KEY_ENV) + return None if key.startswith("enc:") else key + + +def _public(row) -> dict: + """What the UI may see about a stored key: never the key itself.""" + return {"last4": keys.last4(_key(row) or ""), "created_at": row.created_at, "last_ok_at": row.last_ok_at, + "last_error": row.last_error, "models": json.loads(row.models or "{}")} + + +def _fail(status: int, message: str, code: str) -> JSONResponse: + return JSONResponse({"success": False, "error": message, "code": code}, status) + + +@router.get("", dependencies=[Depends(rate_limit())]) +def list_keys(claims: dict = Depends(current_user), engine: Engine = Depends(get_engine)): + from models import UserKey + + with Session(engine) as s: + rows = {r.provider: r for r in s.scalars(select(UserKey).where(UserKey.user_id == claims["sub"]))} + providers = [{"provider": p, "name": keys.PROVIDER_NAMES[p], "models": _catalogue(p), + "key": _public(rows[p]) if p in rows else None} for p in BYOK_PROVIDERS] + with keys.for_user(claims["sub"], engine): + in_use = {} + for role in ROLES: + chain = keys.user_chain(role) + in_use[role] = f"{chain[0].provider}:{chain[0].model}" if chain else "platform" + return {"success": True, "providers": providers, "in_use": in_use} + + +@router.put("/{provider}", dependencies=[Depends(rate_limit("5 per minute"))]) +def put_key(provider: str, body: KeyIn, claims: dict = Depends(current_user), engine: Engine = Depends(get_engine)): + from llm.providers import PROVIDERS + from models import UserKey + + if claims.get("is_anonymous"): + raise HTTPException(403, "Demo accounts can't save keys. Sign up to add yours.") + catalogue = _catalogue(_provider(provider)) + picks = body.models or {} + if any(model not in catalogue.get(role, []) for role, model in picks.items()): + raise HTTPException(422, "That model isn't offered for this provider.") + key_enc = encrypt_secret(body.key, key_env=keys.KEY_ENV) + if not key_enc.startswith("enc:"): + return _fail(503, "Saving keys isn't available right now.", "keys_unavailable") + ok, message = PROVIDERS[provider].check_key(key=body.key) + if not ok: + return _fail(400, keys.mask(message, body.key), "key_invalid") + with Session(engine) as s: + row = _row(s, claims["sub"], provider) + if row is None: + row = UserKey(user_id=claims["sub"], provider=provider) + s.add(row) + row.key_enc, row.models = key_enc, json.dumps(picks) if picks else None + row.last_ok_at, row.last_error = datetime.utcnow(), None + s.commit() + return {"success": True, "key": _public(row)} + + +@router.post("/{provider}/test", dependencies=[Depends(rate_limit("5 per minute"))]) +def test_key(provider: str, claims: dict = Depends(current_user), engine: Engine = Depends(get_engine)): + from llm.providers import PROVIDERS + + with Session(engine) as s: + row = _row(s, claims["sub"], _provider(provider)) + if row is None: + raise HTTPException(404, "No key saved for this provider.") + key = _key(row) + ok, message = PROVIDERS[provider].check_key(key=key) if key else (False, keys.UNREADABLE_DETAIL) + message = keys.mask(message, key) + if not ok: + row.last_error = message[:500] + s.commit() + if ok: + keys.mark_ok(engine, claims["sub"], provider) + return {"success": ok, "message": message} + + +@router.delete("/{provider}", dependencies=[Depends(rate_limit())]) +def delete_key(provider: str, claims: dict = Depends(current_user), engine: Engine = Depends(get_engine)): + with Session(engine) as s: + row = _row(s, claims["sub"], _provider(provider)) + if row is None: + raise HTTPException(404, "No key saved for this provider.") + s.delete(row) + s.commit() + return {"success": True} diff --git a/backend/asgi.py b/backend/asgi.py index 48e3552..026b657 100644 --- a/backend/asgi.py +++ b/backend/asgi.py @@ -20,7 +20,7 @@ from fastapi.middleware.cors import CORSMiddleware import app as flask_module -from api import agent, demo, documents, errors, files, profile, review, statements +from api import agent, demo, documents, errors, files, keys, profile, review, statements from api.deps import llm_deadline from config import Config @@ -89,4 +89,5 @@ def create_app(*routers: APIRouter) -> FastAPI: return app -app = create_app(documents.router, agent.router, demo.router, review.router, files.router, profile.router, statements.router) +app = create_app(documents.router, agent.router, demo.router, review.router, files.router, profile.router, statements.router, + keys.router) diff --git a/backend/tests/test_api_keys.py b/backend/tests/test_api_keys.py new file mode 100644 index 0000000..50e4f49 --- /dev/null +++ b/backend/tests/test_api_keys.py @@ -0,0 +1,196 @@ +"""SPEC-BYOK: the keys API (add, list, test and remove a user's model keys).""" +import logging + +import jwt +import pytest +from fastapi.testclient import TestClient +from sqlalchemy import create_engine, select +from sqlalchemy.orm import Session + +USER = "7a1c0000-0000-4000-8000-000000000001" +OTHER = "7a1c0000-0000-4000-8000-000000000002" +KEY = "sk-test-1234567890abcd" + + +def bearer(sub): + return {"Authorization": "Bearer " + jwt.encode({"sub": sub, "email": "me@example.com"}, "k" * 32, + algorithm="HS256")} + + +class _Resp: + def __init__(self, status, body): + self.status_code, self._body, self.text = status, body, str(body) + + def json(self): + return self._body + + +@pytest.fixture +def status(monkeypatch): + """The HTTP status every provider key check gets; 401 replies echo the key, as OpenAI's do.""" + import llm.providers + + s = {"code": 200} + + def get(url, headers=None, timeout=None): + if s["code"] == 200: + return _Resp(200, {"data": [{"id": "m"}]}) + return _Resp(s["code"], {"error": {"message": f"Incorrect API key provided: {KEY}"}}) + + monkeypatch.setattr(llm.providers.requests, "get", get) + return s + + +@pytest.fixture +def client(tmp_path, monkeypatch, status): + import models # noqa: F401 + import utils.auth + from api import keys as keys_api + from api import profile + from asgi import create_app + from models.database import db as flask_db + from utils.limiter import limiter + + monkeypatch.setenv("KEYS_ENCRYPTION_KEY", "k" * 32) + monkeypatch.setattr(utils.auth, "verify_token", lambda t: jwt.decode(t, options={"verify_signature": False})) + monkeypatch.setattr(limiter, "enabled", False) + engine = create_engine(f"sqlite:///{tmp_path / 'k.db'}") + flask_db.metadata.create_all(engine) + app = create_app(keys_api.router) + app.dependency_overrides[profile.get_engine] = lambda: engine + limiter.reset() + return TestClient(app), engine + + +def rows(engine): + from models import UserKey + + with Session(engine) as s: + return s.scalars(select(UserKey)).all() + + +def test_put_validates_encrypts_and_masks(client): + c, engine = client + r = c.put("/api/keys/openai", headers=bearer(USER), json={"key": f" {KEY} "}) + assert r.status_code == 200 and r.json()["key"]["last4"] == "abcd" + (row,) = rows(engine) + assert row.key_enc.startswith("enc:") and KEY not in row.key_enc + g = c.get("/api/keys", headers=bearer(USER)) + openai = next(p for p in g.json()["providers"] if p["provider"] == "openai") + assert openai["key"]["last4"] == "abcd" and openai["name"] == "OpenAI" and "gpt-6-luna" in openai["models"]["text"] + assert g.json()["in_use"]["agent"] == "openai:gpt-6.1-sol" + assert "sk-test" not in g.text and "sk-test" not in r.text + + +def test_put_rejected_key_not_saved(client, status): + c, engine = client + status["code"] = 401 + r = c.put("/api/keys/openai", headers=bearer(USER), json={"key": KEY}) + assert r.status_code == 400 and r.json()["code"] == "key_invalid" and "1234567890" not in r.text + assert rows(engine) == [] + + +def test_put_unknown_provider_or_unlisted_model(client): + c, engine = client + assert c.put("/api/keys/mistral", headers=bearer(USER), json={"key": KEY}).status_code == 404 + r = c.put("/api/keys/openai", headers=bearer(USER), json={"key": KEY, "models": {"agent": "gpt-4o"}}) + assert r.status_code == 422 and rows(engine) == [] + + +def test_put_again_replaces_key_and_models(client): + c, engine = client + c.put("/api/keys/openai", headers=bearer(USER), json={"key": KEY}) + r = c.put("/api/keys/openai", headers=bearer(USER), + json={"key": "sk-other-key-9999", "models": {"agent": "gpt-5.6-sol"}}) + assert r.json()["key"]["last4"] == "9999" and r.json()["key"]["models"] == {"agent": "gpt-5.6-sol"} + assert len(rows(engine)) == 1 + assert c.get("/api/keys", headers=bearer(USER)).json()["in_use"]["agent"] == "openai:gpt-5.6-sol" + + +def test_test_endpoint_updates_status(client, status): + c, engine = client + c.put("/api/keys/openai", headers=bearer(USER), json={"key": KEY}) + status["code"] = 401 + r = c.post("/api/keys/openai/test", headers=bearer(USER)).json() + assert r["success"] is False and "1234567890" not in r["message"] + assert rows(engine)[0].last_error and "1234567890" not in rows(engine)[0].last_error + status["code"] = 200 + r = c.post("/api/keys/openai/test", headers=bearer(USER)).json() + (row,) = rows(engine) + assert r["success"] is True and row.last_ok_at is not None and row.last_error is None + + +def test_delete_and_isolation(client): + c, engine = client + c.put("/api/keys/openai", headers=bearer(USER), json={"key": KEY}) + other = c.get("/api/keys", headers=bearer(OTHER)).json() + assert all(p["key"] is None for p in other["providers"]) and other["in_use"]["agent"] == "platform" + assert c.post("/api/keys/openai/test", headers=bearer(OTHER)).status_code == 404 + assert c.delete("/api/keys/openai", headers=bearer(OTHER)).status_code == 404 + assert len(rows(engine)) == 1 + assert c.delete("/api/keys/openai", headers=bearer(USER)).json() == {"success": True} + assert rows(engine) == [] + + +def test_key_never_in_logs_or_responses(client, status, caplog): + c, _ = client + caplog.set_level(logging.DEBUG) + bodies = [c.put("/api/keys/openai", headers=bearer(USER), json={"key": KEY}).text, + c.put("/api/keys/openai", headers=bearer(USER), json={"key": KEY + "x" * 400}).text, # 422 + c.post("/api/keys/openai/test", headers=bearer(USER)).text] + status["code"] = 401 + bodies.append(c.post("/api/keys/openai/test", headers=bearer(USER)).text) + bodies.append(c.get("/api/keys", headers=bearer(USER)).text) + for secret in (KEY, KEY[:12], KEY[-12:]): + assert not any(secret in b for b in bodies) + assert not any(secret in r.getMessage() for r in caplog.records) + + +def test_validation_is_rate_limited(client, monkeypatch): + from utils.limiter import limiter + + c, _ = client + monkeypatch.setattr(limiter, "enabled", True) + codes = [c.put("/api/keys/openai", headers=bearer(USER), json={"key": KEY}).status_code for _ in range(6)] + assert codes == [200] * 5 + [429] + + +def test_no_encryption_key_refuses(client, monkeypatch): + c, engine = client + monkeypatch.delenv("KEYS_ENCRYPTION_KEY") + monkeypatch.delenv("SECRET_KEY", raising=False) + r = c.put("/api/keys/openai", headers=bearer(USER), json={"key": KEY}) + assert r.status_code == 503 and r.json()["code"] == "keys_unavailable" + assert rows(engine) == [] + + +def test_chat_with_rejected_key_says_so(authed_client, monkeypatch): + import json as _json + + import llm.client + import routes.chat as chat + from app import app + from llm import keys + from models import UserKey + from models.database import db + from utils.crypto import encrypt_secret + + monkeypatch.setenv("KEYS_ENCRYPTION_KEY", "k" * 32) + monkeypatch.setattr(chat, "_save_exchange", lambda *a: None) + monkeypatch.setattr(llm.client, "_sleep", lambda s: None) + monkeypatch.setattr(llm.client.requests, "post", + lambda *a, **k: type("R", (_Resp,), {"headers": {}})(401, {"error": {"message": "bad key"}})) + with app.app_context(): + db.create_all() + db.session.query(UserKey).filter_by(user_id="user-1").delete() + db.session.add(UserKey(user_id="user-1", provider="openai", key_enc=encrypt_secret(KEY, key_env=keys.KEY_ENV), + models=_json.dumps({}))) + db.session.commit() + try: + r = authed_client.post("/chat", json={"query": "hi"}, headers={"Authorization": "Bearer x.y.z"}) + assert r.status_code == 400 and r.get_json()["code"] == "key_rejected" + assert "OpenAI" in r.get_json()["error"] + finally: + with app.app_context(): + db.session.query(UserKey).filter_by(user_id="user-1").delete() + db.session.commit() From 454c4a324c3a02f962d851e95fd23ae26188dd0b Mon Sep 17 00:00:00 2001 From: coderconnoisseur <142653086+coderconnoisseur@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:29:03 +0530 Subject: [PATCH 06/10] Email poller: a rejected user key or used-up allowance stops retrying and shows on the inbox status --- backend/tests/test_email_poller.py | 58 ++++++++++++++++++++++++++++++ backend/utils/email_poller.py | 26 ++++++++++++-- 2 files changed, 82 insertions(+), 2 deletions(-) diff --git a/backend/tests/test_email_poller.py b/backend/tests/test_email_poller.py index d5d0391..9a5ab4c 100644 --- a/backend/tests/test_email_poller.py +++ b/backend/tests/test_email_poller.py @@ -133,6 +133,64 @@ def parse(subject, *a, **k): assert r["success"] and r["transactions_added"] == 1 from models import ProcessedEmail assert [p.message_id for p in ProcessedEmail.query.all()] == [""] + assert env.config(U1).last_error is None # an outage is not a key problem + + +def _add_groq_key(env): + from llm import keys + from models import UserKey + from utils.crypto import encrypt_secret + env.db.session.add(UserKey(user_id=U1, provider="groq", key_enc=encrypt_secret("gsk_test", key_env=keys.KEY_ENV))) + env.db.session.commit() + + +def test_poller_stops_on_user_key_error_and_resumes_after_fix(env, monkeypatch): + from llm import keys + from models import ProcessedEmail, UserKey + _add_groq_key(env) + calls = [] + + def parse(*a, **k): + calls.append(1) + e = LLMError("auth", "bad") + e.user_key = "groq" + keys.mark_failed("groq", "bad") + raise e + monkeypatch.setattr(env.poller, "parse_money_email", parse) + env.services_[U1] = Fake([mail("<1>"), mail("<2>")]) + env.services_[U2] = Fake([]) + before = env.config(U1).last_poll_time + env.poller.poll_all_users() + assert len(calls) == 1 and ProcessedEmail.query.count() == 0 + assert "Groq key was rejected" in env.config(U1).last_error + assert env.config(U1).last_poll_time == before + + calls.clear() + monkeypatch.setattr(env.poller, "parse_money_email", lambda *a, **k: calls.append(1) or parsed()) + results = env.poller.poll_all_users() + assert calls == [] and [r["user_id"] for r in results] == [U2] # U1 skipped, U2 still polled + + keys.mark_ok(env.db.engine, U1, "groq") + env.db.session.expire_all() + env.poller.poll_all_users() + assert len(calls) == 2 and ProcessedEmail.query.count() == 2 + assert env.db.session.query(UserKey).filter_by(user_id=U1).one().last_error is None + + +def test_needs_key_stops_poll(env, monkeypatch): + from models import ProcessedEmail + from utils.errors import KEY_MESSAGES + calls = [] + + def parse(*a, **k): + calls.append(1) + raise LLMError(LLMError.NEEDS_KEY, "allowance used") + monkeypatch.setattr(env.poller, "parse_money_email", parse) + env.services_[U1] = Fake([mail("<1>"), mail("<2>")]) + env.poller.process_user_emails(env.config(U1)) + assert len(calls) == 1 and ProcessedEmail.query.count() == 0 + assert env.config(U1).last_error == KEY_MESSAGES["needs_key"] + assert env.config(U1).last_poll_time is None def test_default_currency_is_users_then_config(env, monkeypatch): diff --git a/backend/utils/email_poller.py b/backend/utils/email_poller.py index faede89..f506556 100644 --- a/backend/utils/email_poller.py +++ b/backend/utils/email_poller.py @@ -5,6 +5,7 @@ from datetime import date, datetime, timedelta from email.utils import parsedate_to_datetime from config import Config +from llm import keys from inbox.filter import is_candidate from inbox.gmail import GmailClient from inbox.parse import parse_money_email @@ -36,6 +37,16 @@ def _readable(day) -> bool: return False +def _key_problem(e: LLMError) -> str | None: + """A failure only the user can fix (their key, or the free allowance): stop retrying until they do.""" + from llm.keys import PROVIDER_NAMES + from utils.errors import KEY_MESSAGES + + if e.user_key: + return KEY_MESSAGES["key_rejected"].format(provider=PROVIDER_NAMES.get(e.user_key, e.user_key)) + return KEY_MESSAGES["needs_key"] if e.kind == LLMError.NEEDS_KEY else None + + _COUNT = {'imported': 'transactions_added', 'credit': 'credits', 'duplicate': 'duplicates', 'matched': 'duplicates'} @@ -59,6 +70,7 @@ def process_user_emails(config: EmailConfig) -> dict: # SINCE is day-granular, so re-scan 2 days back; Message-ID dedup makes that safe since_date = (config.last_poll_time - timedelta(days=2) if config.last_poll_time else datetime.utcnow() - timedelta(days=30)) + key_problem = None # set when only the user can fix the failure: the poll stops and says why retry = False # an email left for the next poll keeps last_poll_time where it was total = 0 # fetch_new_emails yields one message at a time for email_data in service.fetch_new_emails(since_date): @@ -90,12 +102,15 @@ def process_user_emails(config: EmailConfig) -> dict: except LLMError as e: logger.warning(f"Invoice read unavailable ({e.kind}); will retry next poll") llm_down = True + key_problem = _key_problem(e) break except Exception as e: logger.error(f"Error processing attachment {attachment['filename']}: {e}") error_count += 1 if llm_down: # provider outage: leave it unprocessed so the next poll retries retry = True + if key_problem: + break continue elif is_candidate(email_data['subject'], email_data['text'], False): try: @@ -104,6 +119,9 @@ def process_user_emails(config: EmailConfig) -> dict: # provider outage: leave it unprocessed so the next poll retries logger.warning(f"Email parse unavailable ({e.kind}); will retry next poll") retry = True + key_problem = _key_problem(e) + if key_problem: + break continue outcome = "skipped" if parsed is not None: @@ -141,7 +159,7 @@ def process_user_emails(config: EmailConfig) -> dict: config.last_poll_time = datetime.utcnow() config.last_successful_poll = datetime.utcnow() config.emails_processed = (config.emails_processed or 0) + processed_count + counts['transactions_added'] - config.last_error = None + config.last_error = key_problem db.session.commit() logger.info(f"Polling complete: {processed_count} invoices, {counts}, {error_count} errors") @@ -278,8 +296,12 @@ def poll_all_users(): results = [] for config in configs: + if keys.failing(db.engine, config.user_id): + logger.info(f"Skipping user {config.user_id}: user key needs attention") + continue try: - result = process_user_emails(config) + with keys.for_user(config.user_id, db.engine): + result = process_user_emails(config) results.append({ 'user_id': config.user_id, 'email': config.email_address, From c7587951455ce4131f9312e6389e17986a27d08c Mon Sep 17 00:00:00 2001 From: coderconnoisseur <142653086+coderconnoisseur@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:35:29 +0530 Subject: [PATCH 07/10] Free allowance on the platform key for users without their own key (20 questions/day, 20 invoice reads) --- backend/api/agent.py | 30 +++-- backend/api/review.py | 12 ++ backend/config.py | 4 + backend/routes/chat.py | 9 +- backend/tests/test_api_agent.py | 2 +- backend/tests/test_api_demo.py | 11 +- backend/tests/test_ask_lumen.py | 4 +- backend/tests/test_byok_allowance.py | 168 +++++++++++++++++++++++++++ 8 files changed, 209 insertions(+), 31 deletions(-) create mode 100644 backend/tests/test_byok_allowance.py diff --git a/backend/api/agent.py b/backend/api/agent.py index 67712d2..c468e59 100644 --- a/backend/api/agent.py +++ b/backend/api/agent.py @@ -20,6 +20,7 @@ from api.deps import current_user, rate_limit from llm import keys +from utils.errors import KEY_MESSAGES router = APIRouter(prefix="/api/agent", tags=["agent"]) @@ -54,30 +55,22 @@ class Question(BaseModel): question: str = Field(min_length=1, max_length=1000) -DEMO_QUESTIONS_PER_DAY = 20 # all visitors share one free Groq quota (~1K requests/day; SPEC-DEPLOY) - - -def count_demo_question(claims: dict, authorization: str, client: str) -> bool: - """One daily allowance per demo (anonymous) account, shared by /chat and /api/agent/ask. False = used up.""" +def count_platform_question(claims: dict, authorization: str, client: str) -> bool: + """Daily free allowance on the platform key for any user without their own agent key (demo or signed up), + shared by /chat and /api/agent/ask. Needs the `keys.for_user` context. False = used up.""" from limits import parse + from config import Config from utils.limiter import limiter, rate_limit_key - if not claims.get("is_anonymous") or not limiter.enabled: + if keys.has_key("agent") or not limiter.enabled: return True - return limiter.limiter.hit(parse(f"{DEMO_QUESTIONS_PER_DAY} per day"), rate_limit_key(authorization, client), - "demo-questions") - - -def demo_quota(request: Request, claims: dict = Depends(current_user)) -> None: - """Anonymous demo visitors get a daily question cap; signed-up accounts don't.""" - if not count_demo_question(claims, request.headers.get("Authorization", ""), - request.client.host if request.client else "127.0.0.1"): - raise HTTPException(429, f"The demo allows {DEMO_QUESTIONS_PER_DAY} questions a day. Sign up to keep going.") + return limiter.limiter.hit(parse(f"{Config.PLATFORM_QUESTIONS_PER_DAY} per day"), + rate_limit_key(authorization, client), "platform-questions") -@router.post("/ask", dependencies=[Depends(rate_limit("10 per minute")), Depends(demo_quota)]) -def ask(body: Question, claims: dict = Depends(current_user), deps: AgentDeps = Depends(get_agent_deps)): +@router.post("/ask", dependencies=[Depends(rate_limit("10 per minute"))]) +def ask(body: Question, request: Request, claims: dict = Depends(current_user), deps: AgentDeps = Depends(get_agent_deps)): from agent.graph import run_agent from agent.tools import ToolContext from api.profile import profile_for @@ -86,6 +79,9 @@ def ask(body: Question, claims: dict = Depends(current_user), deps: AgentDeps = today=deps.today or date.today(), profile=profile_for(deps.engine, claims["sub"]), question=body.question) with keys.for_user(claims["sub"], deps.engine): + if not count_platform_question(claims, request.headers.get("Authorization", ""), + request.client.host if request.client else "127.0.0.1"): + raise HTTPException(402, KEY_MESSAGES["needs_key"]) out = run_agent(body.question, ctx, complete=deps.complete) return {"success": True, **out} diff --git a/backend/api/review.py b/backend/api/review.py index 1a1deb4..2ff3a22 100644 --- a/backend/api/review.py +++ b/backend/api/review.py @@ -177,7 +177,19 @@ def read_with_feedback(engine: Engine, user_id: str, image_base64: str, media_ty """Read an invoice; if the user has corrected this vendor's invoices before, read it again with those corrections as hints (one extra call, only for vendors with a correction history). The hints used are kept on the invoice (`feedback_hints`) so the reviewer can see them.""" + from sqlalchemy import func + + from config import Config from extract import read + from llm import keys + from models import ReviewItem + from utils.llm import LLMError + + if not keys.has_key("vision"): + with Session(engine) as session: + used = session.scalar(select(func.count()).select_from(ReviewItem).where(ReviewItem.user_id == user_id)) + if used >= Config.PLATFORM_INVOICE_READS: + raise LLMError(LLMError.NEEDS_KEY, "free invoice reads used up") inv = read.read_invoice(image_base64, media_type) hints = corrections_for(engine, user_id, inv.get("vendor_name")) diff --git a/backend/config.py b/backend/config.py index 4cd6a17..b7169a1 100644 --- a/backend/config.py +++ b/backend/config.py @@ -207,6 +207,10 @@ def get_llm_vision_fallback_models(cls) -> list[str]: MAX_UPLOAD_MB = int(os.getenv("MAX_UPLOAD_MB", "10")) MAX_UPLOAD_BYTES = MAX_UPLOAD_MB * 1024 * 1024 + # === Free allowance on the platform key, for users without their own key (BYOK) === + PLATFORM_QUESTIONS_PER_DAY = int(os.getenv("PLATFORM_QUESTIONS_PER_DAY", "20")) + PLATFORM_INVOICE_READS = int(os.getenv("PLATFORM_INVOICE_READS", "20")) + # === Localization === DEFAULT_CURRENCY = os.getenv("DEFAULT_CURRENCY", "INR") diff --git a/backend/routes/chat.py b/backend/routes/chat.py index b2c25ad..f69090a 100644 --- a/backend/routes/chat.py +++ b/backend/routes/chat.py @@ -7,7 +7,7 @@ from models import ChatMessage from models.database import db from utils.auth import require_auth -from utils.errors import api_error, llm_api_error +from utils.errors import KEY_MESSAGES, api_error, llm_api_error from utils.limiter import limiter from utils.llm import LLMError @@ -109,11 +109,10 @@ def chat(): user_id = str(g.user_id) - from api.agent import DEMO_QUESTIONS_PER_DAY, count_demo_question + from api.agent import count_platform_question - if not count_demo_question(g.jwt_claims, request.headers.get("Authorization", ""), request.remote_addr or ""): - return api_error(f"The demo allows {DEMO_QUESTIONS_PER_DAY} questions a day. Sign up to keep going.", - status=429, code="demo_limit") + if not count_platform_question(g.jwt_claims, request.headers.get("Authorization", ""), request.remote_addr or ""): + return api_error(KEY_MESSAGES["needs_key"], status=402, code="needs_key") try: logger.info("Processing chat query for user=%s", user_id) diff --git a/backend/tests/test_api_agent.py b/backend/tests/test_api_agent.py index 673690d..a1dff19 100644 --- a/backend/tests/test_api_agent.py +++ b/backend/tests/test_api_agent.py @@ -33,7 +33,7 @@ def setup(tmp_path, monkeypatch): path = tmp_path / "w.db" engine = create_engine(f"sqlite:///{path}") load_world(engine, world) - flask_db.metadata.create_all(engine, tables=[flask_db.metadata.tables[t] for t in ("proposals", "audit_events")]) + flask_db.metadata.create_all(engine, tables=[flask_db.metadata.tables[t] for t in ("proposals", "audit_events", "user_keys")]) u1, u2 = (u["id"] for u in world["users"]) invoice = next(t for t in world["transactions"] if t["user_id"] == u1 and t["category"] == "Groceries") diff --git a/backend/tests/test_api_demo.py b/backend/tests/test_api_demo.py index e429e4c..916e0e8 100644 --- a/backend/tests/test_api_demo.py +++ b/backend/tests/test_api_demo.py @@ -71,15 +71,14 @@ def test_real_accounts_never_get_demo_data(app): assert _count(engine, MEMBER) == 0 -def test_visitors_have_a_daily_question_cap_members_do_not(app, monkeypatch): - from api import agent as agent_api +def test_keyless_users_have_a_daily_question_cap(app, monkeypatch): + from config import Config - assert agent_api.DEMO_QUESTIONS_PER_DAY == 20 - monkeypatch.setattr(agent_api, "DEMO_QUESTIONS_PER_DAY", 3) # stay under the 10-per-minute limit + assert Config.PLATFORM_QUESTIONS_PER_DAY == 20 + monkeypatch.setattr(Config, "PLATFORM_QUESTIONS_PER_DAY", 3) # stay under the 10-per-minute limit client, _, _ = app ask = lambda who, anon: client.post("/api/agent/ask", headers=bearer(who, anon), json={"question": "hi"}).status_code - assert [ask(VISITOR, True) for _ in range(4)] == [200, 200, 200, 429] - assert [ask(MEMBER, False) for _ in range(5)] == [200] * 5 + assert [ask(VISITOR, True) for _ in range(4)] == [200, 200, 200, 402] @pytest.mark.skipif(not os.getenv("LUMEN_TEST_POSTGRES_URL"), diff --git a/backend/tests/test_ask_lumen.py b/backend/tests/test_ask_lumen.py index fe6fd99..aa4ec52 100644 --- a/backend/tests/test_ask_lumen.py +++ b/backend/tests/test_ask_lumen.py @@ -305,10 +305,10 @@ def test_chat_shares_the_demo_question_cap(authed_client, monkeypatch): import api.agent import routes.chat as chat - monkeypatch.setattr(api.agent, "count_demo_question", lambda *a: False) + monkeypatch.setattr(api.agent, "count_platform_question", lambda *a: False) monkeypatch.setattr(chat, "_ask", lambda *a: pytest.fail("the agent must not run past the cap")) resp = authed_client.post("/chat", json={"query": "hi"}, headers={"Authorization": "Bearer x.y.z"}) - assert resp.status_code == 429 and resp.get_json()["code"] == "demo_limit" + assert resp.status_code == 402 and resp.get_json()["code"] == "needs_key" def test_chat_sources_carry_the_passage_and_the_original_file(authed_client, monkeypatch): diff --git a/backend/tests/test_byok_allowance.py b/backend/tests/test_byok_allowance.py new file mode 100644 index 0000000..9a985c8 --- /dev/null +++ b/backend/tests/test_byok_allowance.py @@ -0,0 +1,168 @@ +"""SPEC-BYOK: the free allowance on the platform key (questions per day, lifetime invoice reads) for keyless users.""" +import io + +import jwt +import pytest +from fastapi.testclient import TestClient +from sqlalchemy import create_engine + +from llm.client import LLMResult + +AUTH = {"Authorization": "Bearer x.y.z"} +USER = "user-1" # what `authed_client` signs in as + + +@pytest.fixture +def flask_db(authed_client, monkeypatch): + """Signed-in Flask client, a real limiter, and no leftover keys or review rows for the user.""" + from app import app + from models import ReviewItem, UserKey + from models.database import db + from utils.limiter import limiter + + monkeypatch.setenv("KEYS_ENCRYPTION_KEY", "k" * 32) + monkeypatch.setattr(limiter, "enabled", True) + limiter.reset() + + def clean(): + with app.app_context(): + db.session.query(UserKey).filter_by(user_id=USER).delete() + db.session.query(ReviewItem).filter_by(user_id=USER).delete() + db.session.commit() + + clean() + yield authed_client + clean() + + +def add_key(provider): + from app import app + from llm import keys + from models import UserKey + from models.database import db + from utils.crypto import encrypt_secret + + with app.app_context(): + db.session.add(UserKey(user_id=USER, provider=provider, key_enc=encrypt_secret("sk-test-1234567890", key_env=keys.KEY_ENV))) + db.session.commit() + + +def add_read(): + from app import app + from models import ReviewItem + from models.database import db + + with app.app_context(): + db.session.add(ReviewItem(user_id=USER, invoice="{}", flags="[]", confidence="high", status="approved")) + db.session.commit() + + +def chat(client): + return client.post("/chat", json={"query": "hi"}, headers=AUTH) + + +@pytest.fixture +def fake_ask(monkeypatch): + import routes.chat as chat_routes + + monkeypatch.setattr(chat_routes, "_ask", lambda *a, **k: {"response": "ok"}) + monkeypatch.setattr(chat_routes, "_save_exchange", lambda *a: None) + + +def test_keyless_user_hits_question_allowance(flask_db, fake_ask, monkeypatch): + from config import Config + + monkeypatch.setattr(Config, "PLATFORM_QUESTIONS_PER_DAY", 2) + assert [chat(flask_db).status_code for _ in range(2)] == [200, 200] + r = chat(flask_db) + assert r.status_code == 402 and r.get_json()["code"] == "needs_key" + assert "Add your own model key" in r.get_json()["error"] + + +def test_user_with_key_not_counted(flask_db, fake_ask, monkeypatch): + from config import Config + + monkeypatch.setattr(Config, "PLATFORM_QUESTIONS_PER_DAY", 2) + add_key("openai") + assert [chat(flask_db).status_code for _ in range(5)] == [200] * 5 + + +def test_demo_user_same_allowance(tmp_path, monkeypatch): + import models # noqa: F401 + import utils.auth + from ai.sql_agent import SQLAgent + from api import agent as agent_api + from asgi import create_app + from config import Config + from models.database import db as flask_db + from rag.embed import FakeEmbedder + from rag.service import RagService + from rag.store import ChunkStore + from utils.limiter import limiter + + monkeypatch.setattr(utils.auth, "verify_token", lambda t: jwt.decode(t, options={"verify_signature": False})) + monkeypatch.setattr(Config, "PLATFORM_QUESTIONS_PER_DAY", 2) + path = tmp_path / "a.db" + engine = create_engine(f"sqlite:///{path}") + flask_db.metadata.create_all(engine) + store = ChunkStore(engine) + store.ensure_schema() + deps = agent_api.AgentDeps(engine=engine, sql_agent=SQLAgent(db_path=str(path)), + rag=RagService(store, FakeEmbedder(), reranker=None), + complete=lambda *a, **k: LLMResult(text="Done.", tool_calls=None, provider="f", model="f")) + app = create_app(agent_api.router) + app.dependency_overrides[agent_api.get_agent_deps] = lambda: deps + limiter.reset() + c = TestClient(app) + h = {"Authorization": "Bearer " + jwt.encode({"sub": "5b0c2d6e-0000-4000-8000-000000000001", "is_anonymous": True}, + "k" * 32, algorithm="HS256")} + ask = lambda: c.post("/api/agent/ask", headers=h, json={"question": "hi"}) + assert [ask().status_code for _ in range(2)] == [200, 200] + r = ask() + assert r.status_code == 402 and "Add your own model key" in r.text + + +def test_invoice_reads_capped_without_vision_key(flask_db, monkeypatch): + from app import app + from config import Config + from extract import read + from models.database import db + from api.review import read_with_feedback + from utils.llm import LLMError + from llm import keys + + monkeypatch.setattr(Config, "PLATFORM_INVOICE_READS", 1) + monkeypatch.setattr(read, "read_invoice", lambda *a, **k: {"vendor_name": None}) + add_read() + + def run(): + with app.app_context(), keys.for_user(USER, db.engine): + return read_with_feedback(db.engine, USER, "aGk=", "image/png") + + with pytest.raises(LLMError) as e: + run() + assert e.value.kind == LLMError.NEEDS_KEY + + r = flask_db.post("/extract", data={"file": (io.BytesIO(b"x" * 20), "i.png")}, headers=AUTH, + content_type="multipart/form-data") + assert r.status_code == 402 and r.get_json()["code"] == "needs_key" + + add_key("groq") # no vision model on Groq: still the platform's + with pytest.raises(LLMError): + run() + + +def test_vision_key_is_not_capped(flask_db, monkeypatch): + from app import app + from api.review import read_with_feedback + from config import Config + from extract import read + from llm import keys + from models.database import db + + monkeypatch.setattr(Config, "PLATFORM_INVOICE_READS", 1) + monkeypatch.setattr(read, "read_invoice", lambda *a, **k: {"vendor_name": None}) + add_read() + add_key("gemini") + with app.app_context(), keys.for_user(USER, db.engine): + assert read_with_feedback(db.engine, USER, "aGk=", "image/png") == {"vendor_name": None} From fa85f807c8fb382a50310fd9ffa5cfdbeec9154a Mon Sep 17 00:00:00 2001 From: coderconnoisseur <142653086+coderconnoisseur@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:39:16 +0530 Subject: [PATCH 08/10] BYOK: users can add their own Tavily key for web search (no daily cap with their key) --- backend/agent/tools.py | 53 +++++++++++++++++++++++------- backend/api/keys.py | 21 ++++++++---- backend/llm/keys.py | 11 +++++-- backend/tests/test_agent_tools.py | 54 +++++++++++++++++++++++++++++++ backend/tests/test_api_keys.py | 32 ++++++++++++++++++ backend/tests/test_byok_client.py | 11 +++++++ 6 files changed, 162 insertions(+), 20 deletions(-) diff --git a/backend/agent/tools.py b/backend/agent/tools.py index 1b90211..0ffe0a7 100644 --- a/backend/agent/tools.py +++ b/backend/agent/tools.py @@ -458,19 +458,44 @@ def _post_tavily(body: dict, key: str) -> requests.Response: return r +def check_tavily_key(key: str) -> tuple[bool, str]: + """One basic search (costs one Tavily credit). Only the status goes in the message, never the reply.""" + try: + r = requests.post(TAVILY, json={"query": "test", "search_depth": "basic", "max_results": 1, + "include_answer": False}, + headers={"Authorization": f"Bearer {key}"}, timeout=8) + except requests.RequestException as e: + return False, f"could not reach Tavily ({type(e).__name__})." + if r.status_code != 200: + return False, f"Tavily rejected the key (HTTP {r.status_code})." + return True, "Tavily key OK." + + def _web_search(ctx: ToolContext, args: WebSearchArgs) -> dict: - def used(key: str) -> int: - day, n = _web_searches.get(key, (ctx.today, 0)) - return n if day == ctx.today else 0 - mine, everyone = used(ctx.user_id), used("*") # "*": every account (user ids are uuids) - if mine >= WEB_SEARCH_PER_DAY: - return {"error": f"The daily web-search limit ({WEB_SEARCH_PER_DAY}) is reached; answer without it and say so."} - if everyone >= WEB_SEARCH_GLOBAL_PER_DAY: - return {"error": "Lumen's daily web-search limit is reached; answer without it and say so."} - _web_searches[ctx.user_id], _web_searches["*"] = (ctx.today, mine + 1), (ctx.today, everyone + 1) + from llm import keys + + rejected = {"error": "Your Tavily key was rejected; update it in Settings. Answer without web search and say so."} + key = keys.user_key("tavily") + if key == keys.UNREADABLE: + keys.mark_failed("tavily", keys.UNREADABLE_DETAIL) + return rejected + if not key: # the platform key: capped per account and overall + def used(k: str) -> int: + day, n = _web_searches.get(k, (ctx.today, 0)) + return n if day == ctx.today else 0 + mine, everyone = used(ctx.user_id), used("*") # "*": every account (user ids are uuids) + if mine >= WEB_SEARCH_PER_DAY: + return {"error": f"The daily web-search limit ({WEB_SEARCH_PER_DAY}) is reached; answer without it and " + "say so."} + if everyone >= WEB_SEARCH_GLOBAL_PER_DAY: + return {"error": "Lumen's daily web-search limit is reached; answer without it and say so."} + _web_searches[ctx.user_id], _web_searches["*"] = (ctx.today, mine + 1), (ctx.today, everyone + 1) try: r = _post_tavily({"query": args.query, "search_depth": "basic", "max_results": 5, "include_answer": False}, - os.environ.get("TAVILY_API_KEY", "")) + key or os.environ.get("TAVILY_API_KEY", "")) + if key and r.status_code in (401, 403): + keys.mark_failed("tavily", f"Tavily rejected the key (HTTP {r.status_code}).") + return rejected r.raise_for_status() results = [{"id": f"web:{ctx.web_ids + i}", "title": str(x.get("title") or "")[:120], "url": x["url"], "site": urlparse(x["url"]).hostname or "", "published": x.get("published_date"), @@ -483,6 +508,8 @@ def used(key: str) -> int: while len(json.dumps(out, default=str)) >= MAX_TOOL_RESULT_CHARS: out["results"].pop() # long URLs: drop the last result, never cut mid-JSON (a cut one could be cited) ctx.web_ids += len(out["results"]) + if key: + keys.note_ok("tavily") return out except (requests.RequestException, KeyError, ValueError, TypeError, AttributeError) as e: status = getattr(getattr(e, "response", None), "status_code", None) @@ -533,7 +560,11 @@ def _income_tax_india(ctx: ToolContext, args: IncomeTaxArgs) -> dict: def _available(name: str) -> bool: - return name in TOOLS and (name != "web_search" or bool(os.environ.get("TAVILY_API_KEY"))) + if name != "web_search": + return name in TOOLS + from llm import keys + + return bool(os.environ.get("TAVILY_API_KEY") or keys.user_key("tavily")) def _clean(schema: dict) -> dict: diff --git a/backend/api/keys.py b/backend/api/keys.py index 604056e..f9a0e67 100644 --- a/backend/api/keys.py +++ b/backend/api/keys.py @@ -19,6 +19,7 @@ router = APIRouter(prefix="/api/keys", tags=["keys"]) ROLES = ("agent", "text", "vision") +PROVIDERS_ALL = (*BYOK_PROVIDERS, "tavily") # tavily: web search only, no models class KeyIn(BaseModel): @@ -30,11 +31,13 @@ class KeyIn(BaseModel): def _catalogue(provider: str) -> dict[str, list[str]]: from llm.client import registry + if provider == "tavily": + return {} return {role: [e.model for e in registry().byok(provider, role)] for role in ROLES} def _provider(provider: str) -> str: - if provider not in BYOK_PROVIDERS: + if provider not in PROVIDERS_ALL: raise HTTPException(404, "Unknown provider.") return provider @@ -57,6 +60,13 @@ def _public(row) -> dict: "last_error": row.last_error, "models": json.loads(row.models or "{}")} +def _check(provider: str, key: str) -> tuple[bool, str]: + from agent.tools import check_tavily_key + from llm.providers import PROVIDERS + + return check_tavily_key(key) if provider == "tavily" else PROVIDERS[provider].check_key(key=key) + + def _fail(status: int, message: str, code: str) -> JSONResponse: return JSONResponse({"success": False, "error": message, "code": code}, status) @@ -68,7 +78,7 @@ def list_keys(claims: dict = Depends(current_user), engine: Engine = Depends(get with Session(engine) as s: rows = {r.provider: r for r in s.scalars(select(UserKey).where(UserKey.user_id == claims["sub"]))} providers = [{"provider": p, "name": keys.PROVIDER_NAMES[p], "models": _catalogue(p), - "key": _public(rows[p]) if p in rows else None} for p in BYOK_PROVIDERS] + "key": _public(rows[p]) if p in rows else None} for p in PROVIDERS_ALL] with keys.for_user(claims["sub"], engine): in_use = {} for role in ROLES: @@ -79,7 +89,6 @@ def list_keys(claims: dict = Depends(current_user), engine: Engine = Depends(get @router.put("/{provider}", dependencies=[Depends(rate_limit("5 per minute"))]) def put_key(provider: str, body: KeyIn, claims: dict = Depends(current_user), engine: Engine = Depends(get_engine)): - from llm.providers import PROVIDERS from models import UserKey if claims.get("is_anonymous"): @@ -91,7 +100,7 @@ def put_key(provider: str, body: KeyIn, claims: dict = Depends(current_user), en key_enc = encrypt_secret(body.key, key_env=keys.KEY_ENV) if not key_enc.startswith("enc:"): return _fail(503, "Saving keys isn't available right now.", "keys_unavailable") - ok, message = PROVIDERS[provider].check_key(key=body.key) + ok, message = _check(provider, body.key) if not ok: return _fail(400, keys.mask(message, body.key), "key_invalid") with Session(engine) as s: @@ -107,14 +116,12 @@ def put_key(provider: str, body: KeyIn, claims: dict = Depends(current_user), en @router.post("/{provider}/test", dependencies=[Depends(rate_limit("5 per minute"))]) def test_key(provider: str, claims: dict = Depends(current_user), engine: Engine = Depends(get_engine)): - from llm.providers import PROVIDERS - with Session(engine) as s: row = _row(s, claims["sub"], _provider(provider)) if row is None: raise HTTPException(404, "No key saved for this provider.") key = _key(row) - ok, message = PROVIDERS[provider].check_key(key=key) if key else (False, keys.UNREADABLE_DETAIL) + ok, message = _check(provider, key) if key else (False, keys.UNREADABLE_DETAIL) message = keys.mask(message, key) if not ok: row.last_error = message[:500] diff --git a/backend/llm/keys.py b/backend/llm/keys.py index bff4d40..ffc29f8 100644 --- a/backend/llm/keys.py +++ b/backend/llm/keys.py @@ -19,7 +19,7 @@ UNREADABLE = "" # stands in for a stored key that can't be decrypted: never sent anywhere UNREADABLE_DETAIL = "Stored key can't be decrypted; re-enter it in Settings." PROVIDER_NAMES = {"openai": "OpenAI", "anthropic": "Anthropic", "gemini": "Gemini", "groq": "Groq", - "openrouter": "OpenRouter"} + "openrouter": "OpenRouter", "tavily": "Tavily"} _ctx: contextvars.ContextVar[dict | None] = contextvars.ContextVar("llm_user", default=None) @@ -71,6 +71,11 @@ def user_chain(role: str) -> list[Entry]: return chain +def user_key(provider: str) -> str | None: + """The current user's key for a non-model service (Tavily), or None; UNREADABLE if it can't be decrypted.""" + return next((key for p, key, _ in _rows() if p == provider), None) + + def is_key_error(e: LLMError) -> bool: """The user's key itself is at fault. A CONFIG error counts only on 403: other 400/404s are request problems.""" return e.kind in (LLMError.AUTH, LLMError.CREDITS) or (e.kind == LLMError.CONFIG and e.status == 403) @@ -121,7 +126,9 @@ def note_ok(provider: str) -> None: def failing(engine, user_id: str): + """Model keys with an error. A broken Tavily key never stops the email poller.""" from models import UserKey with Session(engine) as s: - return s.scalars(select(UserKey).where(UserKey.user_id == user_id, UserKey.last_error.is_not(None))).all() + return s.scalars(select(UserKey).where(UserKey.user_id == user_id, UserKey.last_error.is_not(None), + UserKey.provider.in_(BYOK_PROVIDERS))).all() diff --git a/backend/tests/test_agent_tools.py b/backend/tests/test_agent_tools.py index cce8409..d536b75 100644 --- a/backend/tests/test_agent_tools.py +++ b/backend/tests/test_agent_tools.py @@ -637,3 +637,57 @@ def test_web_search_description_and_prompt(): assert "Never put the user's data" in d and "untrusted data" in d and "[web:n]" in d text = system_prompt(date(2026, 10, 17)) assert "web_search when available" in text and "currency" not in text + + +@pytest.fixture +def tavily_user(tmp_path, monkeypatch): + """A user with their own Tavily key, inside keys.for_user, and no platform key.""" + from llm import keys + from models import UserKey + from models.database import db as flask_db + from sqlalchemy.orm import Session + from utils.crypto import encrypt_secret + + monkeypatch.setenv("KEYS_ENCRYPTION_KEY", "k" * 32) + monkeypatch.delenv("TAVILY_API_KEY", raising=False) + engine = create_engine(f"sqlite:///{tmp_path / 'k.db'}") + flask_db.metadata.create_all(engine, tables=[flask_db.metadata.tables["user_keys"]]) + key = "tvly-user-1234567890abcd" + with Session(engine) as s, s.begin(): + s.add(UserKey(user_id="tav-user", provider="tavily", key_enc=encrypt_secret(key, key_env=keys.KEY_ENV))) + with keys.for_user("tav-user", engine): + yield key, engine + + +def test_web_search_uses_user_tavily_key_without_caps(db, world, monkeypatch, tavily_user): + from agent import tools + from agent.tools import WEB_SEARCH_PER_DAY, run_tool, tool_schemas + + key, _ = tavily_user + sent = [] + + def post(url, json=None, headers=None, timeout=None): + sent.append(headers) + return _Tavily(body={"results": []}) + monkeypatch.setattr(tools.requests, "post", post) + assert "web_search" in [s["function"]["name"] for s in tool_schemas()] + me = _web_ctx(db, world) + for _ in range(WEB_SEARCH_PER_DAY + 2): + assert "results" in run_tool(me, "web_search", {"query": "repo rate"}) + assert all(h == {"Authorization": f"Bearer {key}"} for h in sent) and len(sent) == WEB_SEARCH_PER_DAY + 2 + assert me.user_id not in tools._web_searches + + +def test_web_search_user_key_rejected(db, world, monkeypatch, tavily_user): + from agent import tools + from agent.tools import run_tool + from models import UserKey + from sqlalchemy.orm import Session + + key, engine = tavily_user + monkeypatch.setattr(tools.requests, "post", lambda *a, **k: _Tavily(401, {"detail": f"bad key {key}"})) + out = run_tool(_web_ctx(db, world), "web_search", {"query": "repo rate"}) + assert "Your Tavily key was rejected" in out["error"] + with Session(engine) as s: + err = s.query(UserKey).one().last_error + assert err == "Tavily rejected the key (HTTP 401)." and not any(key[i:i + 8] in err for i in range(len(key) - 7)) diff --git a/backend/tests/test_api_keys.py b/backend/tests/test_api_keys.py index 50e4f49..24a155b 100644 --- a/backend/tests/test_api_keys.py +++ b/backend/tests/test_api_keys.py @@ -194,3 +194,35 @@ def test_chat_with_rejected_key_says_so(authed_client, monkeypatch): with app.app_context(): db.session.query(UserKey).filter_by(user_id="user-1").delete() db.session.commit() + + +def test_tavily_key_put_test_delete(client, monkeypatch): + from agent import tools + + c, engine = client + seen, code = [], {"n": 200} + + def post(url, json=None, headers=None, timeout=None): + seen.append({"url": url, "json": json, "headers": headers, "timeout": timeout}) + return _Resp(code["n"], {"detail": {"error": f"Unauthorized: {KEY}"}} if code["n"] != 200 else {"results": []}) + monkeypatch.setattr(tools.requests, "post", post) + r = c.put("/api/keys/tavily", headers=bearer(USER), json={"key": KEY}) + assert r.status_code == 200 and r.json()["key"]["last4"] == "abcd" + assert seen[0] == {"url": "https://api.tavily.com/search", "headers": {"Authorization": f"Bearer {KEY}"}, + "json": {"query": "test", "search_depth": "basic", "max_results": 1, "include_answer": False}, + "timeout": 8} + g = c.get("/api/keys", headers=bearer(USER)).json() + tavily = next(p for p in g["providers"] if p["provider"] == "tavily") + assert tavily["name"] == "Tavily" and tavily["models"] == {} and tavily["key"]["last4"] == "abcd" + assert all("tavily" not in v for v in g["in_use"].values()) + r = c.put("/api/keys/tavily", headers=bearer(USER), json={"key": KEY, "models": {"agent": "x"}}) + assert r.status_code == 422 + assert c.post("/api/keys/tavily/test", headers=bearer(USER)).json()["success"] is True + c.delete("/api/keys/tavily", headers=bearer(USER)) + code["n"] = 401 + r = c.put("/api/keys/tavily", headers=bearer(USER), json={"key": KEY}) + assert r.status_code == 400 and r.json()["code"] == "key_invalid" and "1234567890" not in r.text + assert "Tavily rejected the key (HTTP 401)" in r.json()["error"] and rows(engine) == [] + code["n"] = 200 + c.put("/api/keys/tavily", headers=bearer(USER), json={"key": KEY}) + assert c.delete("/api/keys/tavily", headers=bearer(USER)).json() == {"success": True} and rows(engine) == [] diff --git a/backend/tests/test_byok_client.py b/backend/tests/test_byok_client.py index 8cb172f..e9d49f3 100644 --- a/backend/tests/test_byok_client.py +++ b/backend/tests/test_byok_client.py @@ -198,3 +198,14 @@ def who(): resp = app.test_client().get("/who", headers={"Authorization": "Bearer t"}) assert resp.get_data(as_text=True) == U assert keys.current() is None + + +def test_failing_ignores_a_broken_tavily_key(engine): + add_key(engine, U, "tavily", "tvly-user-1234567890") + with keys.for_user(U, engine): + keys.mark_failed("tavily", "Tavily rejected the key (HTTP 401).") + assert row(engine, "tavily").last_error and keys.failing(engine, U) == [] + add_key(engine, U, "openai", "sk-user-1234567890") + with keys.for_user(U, engine): + keys.mark_failed("openai", "bad") + assert [r.provider for r in keys.failing(engine, U)] == ["openai"] From e0ae9d68c6065b20d55885bc1d5bc01b0ef8ca96 Mon Sep 17 00:00:00 2001 From: coderconnoisseur <142653086+coderconnoisseur@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:45:31 +0530 Subject: [PATCH 09/10] Settings: add your own model key; prompts when the free allowance or a key runs out; BYOK docs --- TODO.md | 1 + backend/.env.example | 10 + docs/CODEBASE-GUIDE.md | 16 ++ docs/DEPLOY.md | 5 +- docs/direction/PROGRESS.md | 3 + docs/direction/STORY.md | 13 ++ .../src/app/(app)/chatbot/chatbotContent.tsx | 16 +- .../app/(app)/documents/documentsContent.tsx | 4 +- frontend/src/app/(app)/settings/page.tsx | 8 + .../app/(app)/settings/settingsContent.tsx | 195 ++++++++++++++++++ frontend/src/components/app-sidebar.tsx | 6 + frontend/src/components/emailStatusCard.tsx | 9 + frontend/src/lib/api/client.ts | 27 +++ render.yaml | 8 + 14 files changed, 318 insertions(+), 3 deletions(-) create mode 100644 frontend/src/app/(app)/settings/page.tsx create mode 100644 frontend/src/app/(app)/settings/settingsContent.tsx diff --git a/TODO.md b/TODO.md index fda61fb..60ebf0c 100644 --- a/TODO.md +++ b/TODO.md @@ -52,6 +52,7 @@ What `/` promises today (`frontend/src/app/page.tsx`) that the app can't do yet: - [ ] Suggested questions ignore `use_case`; no test that another user can't fetch your avatar. ## 5. AI engineering (see `docs/direction/PROGRESS.md` → Deferred) +- [x] Bring your own model keys: Settings page, encrypted per-user keys, free allowance on the platform key (SPEC-BYOK, 2026-10-20). - [x] Live exchange-rate tool and follow-up suggestions in Ask Lumen (SPEC-CHAT-TOOLS, 2026-10-15). - [x] The Ask Lumen audit as a 30-question eval with an offline grader (`ask_lumen`, 2026-10-17). - [ ] **Owner: record the `ask_lumen` eval** with a real Groq key and `TAVILY_API_KEY` unset (command in diff --git a/backend/.env.example b/backend/.env.example index ae4201b..3d71231 100644 --- a/backend/.env.example +++ b/backend/.env.example @@ -36,6 +36,16 @@ SUPABASE_JWT_AUD=authenticated # Email encryption (optional — falls back to SECRET_KEY) # EMAIL_ENCRYPTION_KEY=generate-a-long-random-string +# Users' own model keys (Settings page). Encrypts them at rest: a random 32+ character value. Changing it makes +# stored keys unreadable and users must re-add them. Without it, saving a key returns 503. +# KEYS_ENCRYPTION_KEY=generate-a-long-random-string +# Free allowance on Lumen's keys for users who haven't added their own +# PLATFORM_QUESTIONS_PER_DAY=20 +# PLATFORM_INVOICE_READS=20 +# Point the OpenAI / Anthropic providers at a test server (the tests do this; leave unset otherwise) +# LUMEN_OPENAI_BASE_URL= +# LUMEN_ANTHROPIC_BASE_URL= + # Upload limits MAX_UPLOAD_MB=10 diff --git a/docs/CODEBASE-GUIDE.md b/docs/CODEBASE-GUIDE.md index 4cfeff1..066941c 100644 --- a/docs/CODEBASE-GUIDE.md +++ b/docs/CODEBASE-GUIDE.md @@ -375,6 +375,22 @@ embeds the demo documents once per server and uploads the files in the backgroun that vendor has corrections, reads again with them as JSON hints in the prompt; the hints are kept on the invoice as `feedback_hints`. Not yet measured: the synthetic invoices are read perfectly, so there's nothing to correct. +## Model keys (BYOK, SPEC-BYOK) +- Users add their own OpenAI / Anthropic / Gemini / Groq / OpenRouter key (and a Tavily key for web search) on + `/settings`. `api/keys.py` validates it against the provider (`check_key`), encrypts it with + `KEYS_ENCRYPTION_KEY` (`utils/crypto.py`) into `user_keys`, and only ever returns the last 4. Test and Remove + are on the same router; the key never appears in a response, log line or error (`keys.mask`). +- Flow: `require_auth` calls `llm.keys.for_user(user, engine)` (a contextvar) -> `keys.user_chain(role)` builds the + models that user's keys allow, from the BYOK catalogue in `llm/registry.py` (the `agent` role exists only there) -> + `llm.client.complete` tries that chain. If the user has keys but none fits a role, or a key is rejected, it fails + with `key_rejected` (400); there is no silent fallback to the platform's models. +- No key at all: the platform's chain, capped by the free allowance (`PLATFORM_QUESTIONS_PER_DAY` questions a day, + `PLATFORM_INVOICE_READS` invoice reads in total). Past it: `needs_key` (402). The chat, the documents question box + and the inbox status card then link to Settings. +- Email poller (no request, so no contextvar): it enters `for_user` for each config's owner. A rejected key or a + used-up allowance stops the retries and is written to `last_error`, which the inbox card shows. +- `LUMEN_OPENAI_BASE_URL` / `LUMEN_ANTHROPIC_BASE_URL` move those providers to a test server (used by the tests). + ## Statement import (SPEC-STATEMENTS) - `backend/statements/read.py`: decodes the CSV, finds the header row, guesses the column mapping (date, description, debit, credit, amount, Dr/Cr marker), parses lines with signs. `vendor.py` cleans bank descriptions diff --git a/docs/DEPLOY.md b/docs/DEPLOY.md index 4f26a37..d95118a 100644 --- a/docs/DEPLOY.md +++ b/docs/DEPLOY.md @@ -57,7 +57,10 @@ free tier). Delete it: open it → **Settings** → bottom of the page → **Del `5432`, user `postgres.`). Not the "Direct connection" one. - `SUPABASE_SERVICE_ROLE_KEY` (1.4, the **service_role** key: server only, never in the frontend). - `GROQ_API_KEY`, `GEMINI_API_KEY`, `OPENROUTER_API_KEY`: same values as your local `backend/.env`. - - `SECRET_KEY` and `EMAIL_ENCRYPTION_KEY`: click **Generate** for each. + - `SECRET_KEY`, `EMAIL_ENCRYPTION_KEY` and `KEYS_ENCRYPTION_KEY`: click **Generate** for each. The last one + encrypts users' own model keys; changing it later makes stored keys unreadable and users must re-add them. + - `PLATFORM_QUESTIONS_PER_DAY` and `PLATFORM_INVOICE_READS` (optional, both default to 20): the free allowance + for users without their own key. - `TAVILY_API_KEY` (optional): from app.tavily.com; without it Ask Lumen has no web search. → **Save, rebuild, and deploy**. diff --git a/docs/direction/PROGRESS.md b/docs/direction/PROGRESS.md index 74cdd80..459c0ca 100644 --- a/docs/direction/PROGRESS.md +++ b/docs/direction/PROGRESS.md @@ -18,6 +18,9 @@ Tests: `cd backend && env -u OPENROUTER_API_KEY python -m pytest -q` and `... py | — | SPEC-AGENT / SPEC-EXTRACT / SPEC-UX one-pagers | ✅ approved 2026-10-04 (LangGraph; auto-approve high confidence; UX: ship, don't polish) | | | 7 | LLM-02 Groq SQL bench | ✅ done | PR #10 | | 8 | SPEC-AGENT: tools, loop + API, evals, test page | 🔄 PRs #11, #12 open; evals recording | PR #11, #12 | +| 9 | SPEC-BYOK: users' own model keys, free allowance, Settings page | ✅ done 2026-10-20 (backend suite 857 -> 898 passed) | branch `feat/byok` | + +BYOK note: the agent eval cassettes still need their one re-record after SPEC-AGENT-TOOLS; BYOK does not re-record them. Last commit: EVAL-03 in progress (see the branch history). Suite: 422 passed, 1 skipped (the Postgres dialect check; it runs when `LUMEN_TEST_POSTGRES_URL` is set, and in diff --git a/docs/direction/STORY.md b/docs/direction/STORY.md index f5572d0..18e9880 100644 --- a/docs/direction/STORY.md +++ b/docs/direction/STORY.md @@ -274,3 +274,16 @@ with the numbers and the files that prove them. Newest last. Raw numbers live in - **R:** Re-run on production: 20 right, 7 partial, 3 wrong (landing questions 9 / 5 / 0, was 1 / 4 / 9); demo setup 8.2 s (was 65-72 s). The 3 left wrong are named with their fixes (tax tool or web search, vendor lookup by kind, empty-result honesty). Story point: measure the product against its own marketing, not just against unit tests. + +## 24. Letting users bring their own model keys (BYOK) +- **S:** Every visitor shared Lumen's free model keys, so quotas ran out for everyone and the app could not be + promoted without paying for all of them. +- **T:** Let each user use their own provider key, without silently falling back to weak free models and without + ever exposing a key. +- **A:** A catalogue of models per provider, keys encrypted at rest and shown as the last 4 only, a contextvar that + builds each request's model chain from that user's keys, and a hard error (never a silent fallback) when a key is + rejected. Users with no key get a small daily allowance on the platform key. Settings page, plus prompts in chat + and the inbox card. A Tavily key covers web search the same way. +- **R:** Backend suite 857 -> 898 passed. Zero-leak tests (log capture and responses) in `test_api_keys.py` and + `test_byok_client.py` show no key, or part of one, reaches a log line or response. Story point: failing loudly + on a bad key beats quietly answering with a worse model. diff --git a/frontend/src/app/(app)/chatbot/chatbotContent.tsx b/frontend/src/app/(app)/chatbot/chatbotContent.tsx index 054c769..62a48da 100644 --- a/frontend/src/app/(app)/chatbot/chatbotContent.tsx +++ b/frontend/src/app/(app)/chatbot/chatbotContent.tsx @@ -2,6 +2,7 @@ // Ask Lumen: one chat with the agent. Every answer shows its sources, the steps behind it, and any change it proposes. import { useEffect, useRef, useState } from "react"; +import Link from "next/link"; import { Paperclip, Plus } from "lucide-react"; import { Answer } from "@/components/ask/answer"; import { Button } from "@/components/ask/button"; @@ -26,6 +27,8 @@ type Message = { elapsedMs?: number; pending?: boolean; error?: boolean; + /** the free allowance ran out, or the user's key was rejected: point at Settings */ + keyIssue?: boolean; }; // Tools whose results are the user's own data; an answer built on them has evidence even without a cited passage. @@ -36,6 +39,12 @@ function greeting() { return hour < 12 ? "Good morning" : hour < 17 ? "Good afternoon" : "Good evening"; } +const KEY_CODES = ["needs_key", "key_rejected"]; + +function errorCode(error: unknown): string | undefined { + return (error as { response?: { data?: { code?: string } } })?.response?.data?.code; +} + function errorMessage(error: unknown): string { return ( (error as { response?: { data?: { error?: string } } })?.response?.data?.error ?? @@ -95,7 +104,7 @@ export default function ChatbotContent() { ); } catch (error) { setMessages((list) => - list.map((m) => (m.id === id ? { ...m, content: errorMessage(error), pending: false, error: true } : m)) + list.map((m) => (m.id === id ? { ...m, content: errorMessage(error), pending: false, error: true, keyIssue: KEY_CODES.includes(errorCode(error) ?? "") } : m)) ); } } @@ -190,6 +199,11 @@ export default function ChatbotContent() { ) : m.error ? (
{m.content} + {m.keyIssue && ( + + Open Settings + + )} {m.question && ( + + + ) : ( + <> + + {saved && ( + <> + + + + )} + + )} +
+ {p.provider === "tavily" && ( +

+ Web search for Ask Lumen. Without your key, a small daily allowance on Lumen's key. +

+ )} + {saved && hasModels && ( +
+ Advanced +
+ {ROLES.map(([role, label]) => ( + + ))} +
+

Model choices are saved with the key: pick, then Replace.

+
+ )} + + ); +} + +export default function SettingsContent() { + const [data, setData] = useState> | null>(null); + + const load = useCallback(() => { + keysApi.list().then(setData).catch((e) => toast.error(apiError(e, "Couldn't load your keys."))); + }, []); + useEffect(load, [load]); + + return ( + +
+ {data?.providers.map((p) => ( + + ))} + {data && ( +

+ Models in use:{" "} + {ROLES.map(([role, label]) => `${label} → ${data.in_use[role] ?? "platform"}`).join(" · ")} +

+ )} +
+ + ); +} diff --git a/frontend/src/components/app-sidebar.tsx b/frontend/src/components/app-sidebar.tsx index a7116a9..c2ac24e 100644 --- a/frontend/src/components/app-sidebar.tsx +++ b/frontend/src/components/app-sidebar.tsx @@ -7,6 +7,7 @@ import { FileText, PieChart, LayoutDashboard, + Settings, Upload, } from "lucide-react"; import Image from "next/image"; @@ -78,6 +79,11 @@ const data = { url: "/import", icon: Upload, }, + { + name: "Settings", + url: "/settings", + icon: Settings, + }, { name: "Review (test)", url: "/review", diff --git a/frontend/src/components/emailStatusCard.tsx b/frontend/src/components/emailStatusCard.tsx index 0f34517..571a9ca 100644 --- a/frontend/src/components/emailStatusCard.tsx +++ b/frontend/src/components/emailStatusCard.tsx @@ -1,5 +1,6 @@ "use client"; +import Link from "next/link"; import { useState, useEffect } from "react"; import { Card, @@ -173,6 +174,14 @@ export function EmailStatusCard() { Last error: {status.last_error} + {/key|allowance/i.test(status.last_error) && ( + <> + {" "} + + Open Settings + + + )} )} diff --git a/frontend/src/lib/api/client.ts b/frontend/src/lib/api/client.ts index 992aca4..f49c6fb 100644 --- a/frontend/src/lib/api/client.ts +++ b/frontend/src/lib/api/client.ts @@ -437,6 +437,33 @@ export const chatApi = { }, }; +/** A user's own model / search keys (SPEC-BYOK). The key itself never comes back, only its last 4. */ +export interface ProviderKey { + provider: string; + name: string; + /** role -> models this provider offers for it; empty for Tavily */ + models: Record; + key: { + last4: string; + created_at: string; + last_ok_at: string | null; + last_error: string | null; + models: Record; + } | null; +} + +export const keysApi = { + list: async (): Promise<{ providers: ProviderKey[]; in_use: Record }> => + (await apiClient.get("/api/keys")).data, + save: async (provider: string, key: string, models?: Record): Promise => + (await apiClient.put(`/api/keys/${provider}`, { key, models })).data.key, + test: async (provider: string): Promise<{ success: boolean; message: string }> => + (await apiClient.post(`/api/keys/${provider}/test`)).data, + remove: async (provider: string): Promise => { + await apiClient.delete(`/api/keys/${provider}`); + }, +}; + /** Analytics — backend derives user from JWT. */ export const analyticsApi = { getTimeRangeAnalytics: async ( diff --git a/render.yaml b/render.yaml index 0d65908..7a2cb1d 100644 --- a/render.yaml +++ b/render.yaml @@ -28,6 +28,14 @@ services: generateValue: true - key: EMAIL_ENCRYPTION_KEY generateValue: true + # Encrypts users' own model keys. Changing it makes stored keys unreadable (users re-add them). + - key: KEYS_ENCRYPTION_KEY + generateValue: true + # Free allowance for users without their own key. + - key: PLATFORM_QUESTIONS_PER_DAY + value: "20" + - key: PLATFORM_INVOICE_READS + value: "20" # Supabase: the session pooler connection string (IPv4; Render's free tier has no IPv6). - key: DATABASE_URL sync: false From 88c47fec911032608653e273ea061e711960c06c Mon Sep 17 00:00:00 2001 From: coderconnoisseur <142653086+coderconnoisseur@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:59:57 +0530 Subject: [PATCH 10/10] BYOK: dedicated key-encryption key; poller skips only invoice emails over the allowance; documents questions count; review fixes --- backend/api/agent.py | 6 ++-- backend/api/documents.py | 10 +++++-- backend/llm/providers.py | 3 +- backend/tests/test_api_documents.py | 17 +++++++++++- backend/tests/test_api_keys.py | 4 +-- backend/tests/test_byok_allowance.py | 41 ++++++++++++++++++++++++---- backend/tests/test_email_poller.py | 11 +++++--- backend/tests/test_llm_providers.py | 8 ++++++ backend/utils/crypto.py | 3 +- backend/utils/email_poller.py | 11 ++++---- docs/CODEBASE-GUIDE.md | 19 +++++++------ 11 files changed, 100 insertions(+), 33 deletions(-) diff --git a/backend/api/agent.py b/backend/api/agent.py index c468e59..8f82171 100644 --- a/backend/api/agent.py +++ b/backend/api/agent.py @@ -20,7 +20,7 @@ from api.deps import current_user, rate_limit from llm import keys -from utils.errors import KEY_MESSAGES +from llm.errors import LLMError router = APIRouter(prefix="/api/agent", tags=["agent"]) @@ -63,7 +63,7 @@ def count_platform_question(claims: dict, authorization: str, client: str) -> bo from config import Config from utils.limiter import limiter, rate_limit_key - if keys.has_key("agent") or not limiter.enabled: + if not limiter.enabled or keys.has_key("agent"): return True return limiter.limiter.hit(parse(f"{Config.PLATFORM_QUESTIONS_PER_DAY} per day"), rate_limit_key(authorization, client), "platform-questions") @@ -81,7 +81,7 @@ def ask(body: Question, request: Request, claims: dict = Depends(current_user), with keys.for_user(claims["sub"], deps.engine): if not count_platform_question(claims, request.headers.get("Authorization", ""), request.client.host if request.client else "127.0.0.1"): - raise HTTPException(402, KEY_MESSAGES["needs_key"]) + raise LLMError(LLMError.NEEDS_KEY, "free question allowance used") # 402 needs_key (api/errors.py) out = run_agent(body.question, ctx, complete=deps.complete) return {"success": True, **out} diff --git a/backend/api/documents.py b/backend/api/documents.py index e4beaa4..4954293 100644 --- a/backend/api/documents.py +++ b/backend/api/documents.py @@ -2,12 +2,13 @@ limits and error bodies come from api/deps.py and api/errors.py (API-01 parity).""" from __future__ import annotations -from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile +from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile from pydantic import BaseModel, Field from api.deps import current_user, rate_limit from config import Config from llm import keys +from llm.errors import LLMError from rag.service import DOC_TYPES, RagService, get_service from utils import files @@ -67,8 +68,13 @@ def search(body: Query, claims: dict = Depends(current_user), rag: RagService = @router.post("/ask", dependencies=[Depends(rate_limit("10 per minute"))]) -def ask(body: Question, claims: dict = Depends(current_user), rag: RagService = Depends(get_service)): +def ask(body: Question, request: Request, claims: dict = Depends(current_user), rag: RagService = Depends(get_service)): + from api.agent import count_platform_question + with keys.for_user(claims["sub"], rag.store.engine): + if not count_platform_question(claims, request.headers.get("Authorization", ""), + request.client.host if request.client else "127.0.0.1"): + raise LLMError(LLMError.NEEDS_KEY, "free question allowance used") # 402 needs_key (api/errors.py) out = rag.ask(claims["sub"], body.question) return {"success": True, "answer": out["answer"], "abstained": out["abstained"], "reason": out["reason"], "sources": out["sources"]} diff --git a/backend/llm/providers.py b/backend/llm/providers.py index 362c97d..f392b58 100644 --- a/backend/llm/providers.py +++ b/backend/llm/providers.py @@ -126,7 +126,8 @@ def check_key(self, timeout: float = 8, key: str | None = None) -> tuple[bool, s except requests.RequestException as e: return False, f"could not reach OpenRouter to verify the key ({type(e).__name__})." if resp.status_code != 200: - return False, f"OpenRouter rejected the key (HTTP {resp.status_code}: {_safe_message(resp)}).{self.auth_hint()}" + hint = self.auth_hint() if key is None else "" # a user's key: never name the server's env var + return False, f"OpenRouter rejected the key (HTTP {resp.status_code}: {_safe_message(resp)}).{hint}" data = resp.json().get("data") or {} free = bool(data.get("is_free_tier")) # `limit_remaining` is the key's dollar credit, not a request count. The ~50/day cap on `:free` models diff --git a/backend/tests/test_api_documents.py b/backend/tests/test_api_documents.py index 7ba80ca..78035d0 100644 --- a/backend/tests/test_api_documents.py +++ b/backend/tests/test_api_documents.py @@ -46,7 +46,7 @@ def llm(prompt, **_): app = create_app(documents.router) app.dependency_overrides[get_service] = lambda: rag c = TestClient(app) - c.prompts = prompts + c.prompts, c.engine = prompts, store.engine return c @@ -89,6 +89,21 @@ def test_users_never_see_or_delete_each_others_documents(client): assert r.status_code == 404 and r.json()["code"] == "http_error" +def test_ask_counts_against_the_question_allowance(client, monkeypatch): + from config import Config + from models.database import db as flask_db + from utils.limiter import limiter + + flask_db.metadata.create_all(client.engine, tables=[flask_db.metadata.tables["user_keys"]]) + monkeypatch.setattr(Config, "PLATFORM_QUESTIONS_PER_DAY", 1) + monkeypatch.setattr(limiter, "enabled", True) + limiter.reset() + ask = lambda: client.post("/api/documents/ask", headers=bearer("carol"), json={"question": "NetLink fee?"}) + assert ask().status_code == 200 + r = ask() + assert r.status_code == 402 and r.json()["code"] == "needs_key" and "Add your own model key" in r.json()["error"] + + def test_rejects_non_pdfs_bad_types_and_missing_auth(client): assert _upload(client, data=b"hello", name="x.txt").status_code == 415 assert _upload(client, doc_type="weird").status_code == 400 diff --git a/backend/tests/test_api_keys.py b/backend/tests/test_api_keys.py index 24a155b..dce5e8d 100644 --- a/backend/tests/test_api_keys.py +++ b/backend/tests/test_api_keys.py @@ -157,8 +157,8 @@ def test_validation_is_rate_limited(client, monkeypatch): def test_no_encryption_key_refuses(client, monkeypatch): c, engine = client - monkeypatch.delenv("KEYS_ENCRYPTION_KEY") - monkeypatch.delenv("SECRET_KEY", raising=False) + monkeypatch.delenv("KEYS_ENCRYPTION_KEY") # SECRET_KEY stays set: it must not stand in + monkeypatch.setenv("SECRET_KEY", "test-secret-key-for-ci-only") r = c.put("/api/keys/openai", headers=bearer(USER), json={"key": KEY}) assert r.status_code == 503 and r.json()["code"] == "keys_unavailable" assert rows(engine) == [] diff --git a/backend/tests/test_byok_allowance.py b/backend/tests/test_byok_allowance.py index 9a985c8..713a111 100644 --- a/backend/tests/test_byok_allowance.py +++ b/backend/tests/test_byok_allowance.py @@ -87,7 +87,7 @@ def test_user_with_key_not_counted(flask_db, fake_ask, monkeypatch): assert [chat(flask_db).status_code for _ in range(5)] == [200] * 5 -def test_demo_user_same_allowance(tmp_path, monkeypatch): +def _agent_client(tmp_path, monkeypatch): import models # noqa: F401 import utils.auth from ai.sql_agent import SQLAgent @@ -101,7 +101,7 @@ def test_demo_user_same_allowance(tmp_path, monkeypatch): from utils.limiter import limiter monkeypatch.setattr(utils.auth, "verify_token", lambda t: jwt.decode(t, options={"verify_signature": False})) - monkeypatch.setattr(Config, "PLATFORM_QUESTIONS_PER_DAY", 2) + monkeypatch.setattr(limiter, "enabled", True) path = tmp_path / "a.db" engine = create_engine(f"sqlite:///{path}") flask_db.metadata.create_all(engine) @@ -113,13 +113,42 @@ def test_demo_user_same_allowance(tmp_path, monkeypatch): app = create_app(agent_api.router) app.dependency_overrides[agent_api.get_agent_deps] = lambda: deps limiter.reset() - c = TestClient(app) - h = {"Authorization": "Bearer " + jwt.encode({"sub": "5b0c2d6e-0000-4000-8000-000000000001", "is_anonymous": True}, - "k" * 32, algorithm="HS256")} + return TestClient(app), engine + + +def _bearer(sub, **claims): + return {"Authorization": "Bearer " + jwt.encode({"sub": sub, **claims}, "k" * 32, algorithm="HS256")} + + +def test_demo_user_same_allowance(tmp_path, monkeypatch): + from config import Config + + monkeypatch.setattr(Config, "PLATFORM_QUESTIONS_PER_DAY", 2) + c, _ = _agent_client(tmp_path, monkeypatch) + h = _bearer("5b0c2d6e-0000-4000-8000-000000000001", is_anonymous=True) ask = lambda: c.post("/api/agent/ask", headers=h, json={"question": "hi"}) assert [ask().status_code for _ in range(2)] == [200, 200] r = ask() - assert r.status_code == 402 and "Add your own model key" in r.text + assert r.status_code == 402 and r.json()["code"] == "needs_key" and "Add your own model key" in r.text + + +def test_fastapi_key_holder_not_counted(tmp_path, monkeypatch): + from sqlalchemy.orm import Session + + from config import Config + from llm import keys + from models import UserKey + from utils.crypto import encrypt_secret + + monkeypatch.setattr(Config, "PLATFORM_QUESTIONS_PER_DAY", 1) + monkeypatch.setenv("KEYS_ENCRYPTION_KEY", "k" * 32) + c, engine = _agent_client(tmp_path, monkeypatch) + sub = "5b0c2d6e-0000-4000-8000-000000000002" + with Session(engine) as s: + s.add(UserKey(user_id=sub, provider="openai", key_enc=encrypt_secret("sk-test-1234567890", key_env=keys.KEY_ENV))) + s.commit() + assert [c.post("/api/agent/ask", headers=_bearer(sub), json={"question": "hi"}).status_code + for _ in range(3)] == [200] * 3 def test_invoice_reads_capped_without_vision_key(flask_db, monkeypatch): diff --git a/backend/tests/test_email_poller.py b/backend/tests/test_email_poller.py index 9a5ab4c..8d11c1b 100644 --- a/backend/tests/test_email_poller.py +++ b/backend/tests/test_email_poller.py @@ -177,18 +177,21 @@ def parse(*a, **k): assert env.db.session.query(UserKey).filter_by(user_id=U1).one().last_error is None -def test_needs_key_stops_poll(env, monkeypatch): +def test_needs_key_skips_only_that_email(env, monkeypatch): from models import ProcessedEmail from utils.errors import KEY_MESSAGES calls = [] def parse(*a, **k): calls.append(1) - raise LLMError(LLMError.NEEDS_KEY, "allowance used") + if len(calls) == 1: + raise LLMError(LLMError.NEEDS_KEY, "allowance used") + return parsed() monkeypatch.setattr(env.poller, "parse_money_email", parse) env.services_[U1] = Fake([mail("<1>"), mail("<2>")]) - env.poller.process_user_emails(env.config(U1)) - assert len(calls) == 1 and ProcessedEmail.query.count() == 0 + r = env.poller.process_user_emails(env.config(U1)) + assert len(calls) == 2 and r["transactions_added"] == 1 + assert [p.message_id for p in ProcessedEmail.query.all()] == ["<2>"] assert env.config(U1).last_error == KEY_MESSAGES["needs_key"] assert env.config(U1).last_poll_time is None diff --git a/backend/tests/test_llm_providers.py b/backend/tests/test_llm_providers.py index c1c6145..93a7e35 100644 --- a/backend/tests/test_llm_providers.py +++ b/backend/tests/test_llm_providers.py @@ -107,6 +107,14 @@ def test_key_check_reports_rejection(monkeypatch): assert not ok and "401" in summary +def test_openrouter_user_key_rejection_hides_server_env_var(monkeypatch): + from llm import providers + + monkeypatch.setattr(providers.requests, "get", lambda *a, **k: _Resp(401, {"error": {"message": "User not found."}})) + assert "OPENROUTER_API_KEY" not in providers.PROVIDERS["openrouter"].check_key(key="sk-or-user-1234")[1] + assert "OPENROUTER_API_KEY" in providers.PROVIDERS["openrouter"].check_key()[1] + + # --- registry --------------------------------------------------------------------- diff --git a/backend/utils/crypto.py b/backend/utils/crypto.py index e77f132..005d9a1 100644 --- a/backend/utils/crypto.py +++ b/backend/utils/crypto.py @@ -14,7 +14,8 @@ def _fernet(key_env: str = "EMAIL_ENCRYPTION_KEY") -> Fernet | None: - raw = os.getenv(key_env) or os.getenv("SECRET_KEY") + # Users' model keys get their own key, never SECRET_KEY (SPEC-BYOK); IMAP passwords may fall back to it. + raw = os.getenv(key_env) or (None if key_env == "KEYS_ENCRYPTION_KEY" else os.getenv("SECRET_KEY")) if not raw: return None # Fernet needs a 32-byte url-safe base64 key; derive deterministically. diff --git a/backend/utils/email_poller.py b/backend/utils/email_poller.py index f506556..7b208b0 100644 --- a/backend/utils/email_poller.py +++ b/backend/utils/email_poller.py @@ -81,7 +81,7 @@ def process_user_emails(config: EmailConfig) -> dict: continue if email_data['has_invoice_attachment']: - llm_down = False + llm_down = stop = False for attachment in email_data['attachments']: try: # saved by an earlier, partly failed pass over this email: don't submit it twice @@ -102,14 +102,15 @@ def process_user_emails(config: EmailConfig) -> dict: except LLMError as e: logger.warning(f"Invoice read unavailable ({e.kind}); will retry next poll") llm_down = True - key_problem = _key_problem(e) + key_problem = _key_problem(e) or key_problem + stop = bool(e.user_key) # a rejected key stops the poll; the allowance skips only this email break except Exception as e: logger.error(f"Error processing attachment {attachment['filename']}: {e}") error_count += 1 if llm_down: # provider outage: leave it unprocessed so the next poll retries retry = True - if key_problem: + if stop: break continue elif is_candidate(email_data['subject'], email_data['text'], False): @@ -119,8 +120,8 @@ def process_user_emails(config: EmailConfig) -> dict: # provider outage: leave it unprocessed so the next poll retries logger.warning(f"Email parse unavailable ({e.kind}); will retry next poll") retry = True - key_problem = _key_problem(e) - if key_problem: + key_problem = _key_problem(e) or key_problem + if e.user_key: # a rejected key stops the poll; the allowance skips only this email break continue outcome = "skipped" diff --git a/docs/CODEBASE-GUIDE.md b/docs/CODEBASE-GUIDE.md index 066941c..2481949 100644 --- a/docs/CODEBASE-GUIDE.md +++ b/docs/CODEBASE-GUIDE.md @@ -378,17 +378,20 @@ embeds the demo documents once per server and uploads the files in the backgroun ## Model keys (BYOK, SPEC-BYOK) - Users add their own OpenAI / Anthropic / Gemini / Groq / OpenRouter key (and a Tavily key for web search) on `/settings`. `api/keys.py` validates it against the provider (`check_key`), encrypts it with - `KEYS_ENCRYPTION_KEY` (`utils/crypto.py`) into `user_keys`, and only ever returns the last 4. Test and Remove + `KEYS_ENCRYPTION_KEY` (`utils/crypto.py`; no `SECRET_KEY` fallback: unset = 503 `keys_unavailable`) into `user_keys`, and only ever returns the last 4. Test and Remove are on the same router; the key never appears in a response, log line or error (`keys.mask`). - Flow: `require_auth` calls `llm.keys.for_user(user, engine)` (a contextvar) -> `keys.user_chain(role)` builds the - models that user's keys allow, from the BYOK catalogue in `llm/registry.py` (the `agent` role exists only there) -> - `llm.client.complete` tries that chain. If the user has keys but none fits a role, or a key is rejected, it fails - with `key_rejected` (400); there is no silent fallback to the platform's models. + models that user's keys allow, from the BYOK catalogue in `llm/registry.yaml` (loaded by `llm/registry.py`; the + `agent` role exists only there) -> `llm.client.complete` tries that chain. A role none of the user's keys covers + (e.g. vision for a Groq-only user) uses the platform chain. `key_rejected` (400) is only for a failing user key; + then there is no silent fallback to the platform's models. - No key at all: the platform's chain, capped by the free allowance (`PLATFORM_QUESTIONS_PER_DAY` questions a day, - `PLATFORM_INVOICE_READS` invoice reads in total). Past it: `needs_key` (402). The chat, the documents question box - and the inbox status card then link to Settings. -- Email poller (no request, so no contextvar): it enters `for_user` for each config's owner. A rejected key or a - used-up allowance stops the retries and is written to `last_error`, which the inbox card shows. + `PLATFORM_INVOICE_READS` invoice reads in total). `/chat`, `/api/agent/ask` and `/api/documents/ask` share the + question allowance. Past it: `needs_key` (402). The chat and the inbox status card link to Settings; the documents + question box appends a Settings hint (not a link). +- Email poller (no request, so no contextvar): it enters `for_user` for each config's owner. A rejected key stops + that user's poll; a used-up allowance skips only the invoice email that needed it (retried next poll). Either is + written to `last_error`, which the inbox card shows. - `LUMEN_OPENAI_BASE_URL` / `LUMEN_ANTHROPIC_BASE_URL` move those providers to a test server (used by the tests). ## Statement import (SPEC-STATEMENTS)