parththakkar106andClaude Opus 4.8
c500203270
The per-IP rate limits could be bypassed entirely: uvicorn ran with --forwarded-allow-ips "*", which trusts the leftmost X-Forwarded-For value (client-controlled), and Render forwards the inbound header rather than stripping it. Rotating the header handed out a fresh rate-limit bucket per request, so the login/register limit (10/5min) and guest-minting limit (30/5min) were no throttle at all — unbounded password guessing and guest-row creation. Confirmed live: fixed IP -> 429 after 10; rotating spoofed header -> no 429 across 14 attempts. Two-layer fix: - limits._client_ip now derives the client IP from the hop the trusted edge appends (rightmost of X-Forwarded-For), which a client can't spoof past; tunable via AIDND_TRUSTED_PROXY_HOPS. Dropped --forwarded-allow-ips "*". - New per-account login throttle (email-keyed, 8 fails / 15 min, cleared on success): stops distributed guessing against one account that a per-IP limit can't, since it can't be diluted across many source addresses. Also close an SSRF on the BYOK endpoint_url (hosted mode only): the connection test and turn/chat streams now refuse a URL that resolves to a non-public address (private/loopback/link-local metadata/reserved), checked at request time so it resists a DNS record flipping to a private IP. No-op locally, where reaching localhost Ollama is intended. Tests: test_ratelimit_hardening.py (8), test_netguard.py (13). 172 pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CYEJKobJ2Re4Dv7qUoSA7
342 lines
15 KiB
Python
342 lines
15 KiB
Python
import json
|
||
from typing import AsyncIterator
|
||
|
||
import httpx
|
||
|
||
from .. import debuglog, netguard
|
||
from .base import PromptParts, Provider, ProviderError
|
||
|
||
# Framing appended after the story text in chat mode, so chat-tuned models keep
|
||
# continuing prose instead of replying conversationally.
|
||
CHAT_CONTINUE_HINT = "\n\n[Continue the story directly. Output only story text.]"
|
||
|
||
# Completion endpoints have no roles, so a plain chat has to be flattened into
|
||
# one labelled transcript that trails off on "Assistant:" for the model to continue.
|
||
_ROLE_LABELS = {"system": "System", "user": "User", "assistant": "Assistant"}
|
||
|
||
|
||
def flatten_messages(messages: list[dict]) -> str:
|
||
turns = "\n\n".join(
|
||
f"{_ROLE_LABELS.get(m['role'], m['role'])}: {m['content']}" for m in messages
|
||
)
|
||
return f"{turns}\n\nAssistant:"
|
||
|
||
|
||
class OpenAICompatibleProvider(Provider):
|
||
"""Adapter for any /v1-style endpoint: Ollama, LM Studio, OpenAI, OpenRouter, vLLM, Groq…"""
|
||
|
||
def __init__(
|
||
self,
|
||
endpoint_url: str,
|
||
api_key: str,
|
||
model: str,
|
||
api_mode: str = "chat",
|
||
reasoning_max_tokens: int = 0,
|
||
):
|
||
self.base_url = endpoint_url.rstrip("/")
|
||
self.api_key = api_key
|
||
self.model = model
|
||
self.api_mode = api_mode # "chat" | "completion"
|
||
# Thinking budget for reasoning models, on top of max_tokens. 0 = the
|
||
# `reasoning` param is not sent (endpoints that don't know it may
|
||
# reject unknown fields); negative = explicitly ask the endpoint to
|
||
# turn reasoning off.
|
||
self.reasoning_max_tokens = reasoning_max_tokens
|
||
|
||
def _headers(self) -> dict:
|
||
headers = {"Content-Type": "application/json"}
|
||
if self.api_key:
|
||
headers["Authorization"] = f"Bearer {self.api_key}"
|
||
return headers
|
||
|
||
def _apply_reasoning_budget(self, body: dict) -> None:
|
||
"""Give reasoning models their own thinking budget (OpenRouter-style),
|
||
raising max_tokens so the actual output keeps its full budget.
|
||
|
||
A negative budget means the opposite: send `effort: "none"` to switch
|
||
reasoning off on models that do it by default (DeepSeek V4 Flash, say).
|
||
That's distinct from `exclude: true`, which still thinks — and bills —
|
||
but hides the trace. Zero stays "send nothing at all" so endpoints that
|
||
reject unknown fields (Ollama) keep working."""
|
||
if self.api_mode != "chat":
|
||
return
|
||
if self.reasoning_max_tokens < 0:
|
||
body["reasoning"] = {"effort": "none"}
|
||
elif self.reasoning_max_tokens > 0:
|
||
body["reasoning"] = {"max_tokens": self.reasoning_max_tokens}
|
||
body["max_tokens"] += self.reasoning_max_tokens
|
||
|
||
def _request(self, parts: PromptParts, temperature: float, max_tokens: int) -> tuple[str, dict]:
|
||
if self.api_mode == "completion":
|
||
url = f"{self.base_url}/completions"
|
||
body = {
|
||
"model": self.model,
|
||
"prompt": f"{parts.system}\n\n{parts.story}",
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens,
|
||
"stream": True,
|
||
}
|
||
else:
|
||
url = f"{self.base_url}/chat/completions"
|
||
body = {
|
||
"model": self.model,
|
||
"messages": [
|
||
{"role": "system", "content": parts.system},
|
||
{"role": "user", "content": parts.story + CHAT_CONTINUE_HINT},
|
||
],
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens,
|
||
"stream": True,
|
||
}
|
||
self._apply_reasoning_budget(body)
|
||
return url, body
|
||
|
||
@staticmethod
|
||
def _extract_chunk(payload: dict) -> str:
|
||
choices = payload.get("choices") or []
|
||
if not choices:
|
||
return ""
|
||
choice = choices[0]
|
||
# chat stream → delta.content; completion stream → text;
|
||
# non-stream fallbacks → message.content / text
|
||
delta = choice.get("delta") or {}
|
||
return (
|
||
delta.get("content")
|
||
or choice.get("text")
|
||
or (choice.get("message") or {}).get("content")
|
||
or ""
|
||
)
|
||
|
||
@staticmethod
|
||
def _extract_reasoning(payload: dict) -> str:
|
||
"""Reasoning-model thinking: OpenRouter normalizes to `reasoning`;
|
||
DeepSeek-style servers use `reasoning_content`."""
|
||
choices = payload.get("choices") or []
|
||
if not choices:
|
||
return ""
|
||
choice = choices[0]
|
||
delta = choice.get("delta") or {}
|
||
message = choice.get("message") or {}
|
||
return (
|
||
delta.get("reasoning")
|
||
or delta.get("reasoning_content")
|
||
or message.get("reasoning")
|
||
or message.get("reasoning_content")
|
||
or ""
|
||
)
|
||
|
||
async def generate(
|
||
self, parts: PromptParts, *, temperature: float, max_tokens: int
|
||
) -> AsyncIterator[tuple[str, str]]:
|
||
"""Yields ("text" | "reasoning", chunk) pairs."""
|
||
if not self.model:
|
||
raise ProviderError("No model configured — set one in Settings.")
|
||
url, body = self._request(parts, temperature, max_tokens)
|
||
async for event in self._stream(url, body):
|
||
yield event
|
||
|
||
async def chat(
|
||
self, messages: list[dict], *, temperature: float, max_tokens: int
|
||
) -> AsyncIterator[tuple[str, str]]:
|
||
"""Plain multi-turn chat — no story framing, no context assembly. Takes
|
||
[{"role", "content"}, ...] straight to the endpoint. Used by the AI Chat
|
||
scratchpad; the turn engine uses generate()."""
|
||
if not self.model:
|
||
raise ProviderError("No model configured — set one in Settings.")
|
||
if self.api_mode == "completion":
|
||
url = f"{self.base_url}/completions"
|
||
body = {
|
||
"model": self.model,
|
||
"prompt": flatten_messages(messages),
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens,
|
||
"stream": True,
|
||
}
|
||
else:
|
||
url = f"{self.base_url}/chat/completions"
|
||
body = {
|
||
"model": self.model,
|
||
"messages": messages,
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens,
|
||
"stream": True,
|
||
}
|
||
self._apply_reasoning_budget(body)
|
||
async for event in self._stream(url, body):
|
||
yield event
|
||
|
||
async def _stream(self, url: str, body: dict) -> AsyncIterator[tuple[str, str]]:
|
||
"""Shared SSE plumbing for generate()/chat(): POST a streaming request
|
||
and yield ("text" | "reasoning", chunk) pairs, logging the exchange."""
|
||
# SSRF guard (hosted mode): a user-supplied endpoint_url must not point
|
||
# at an internal/metadata address. No-op for local installs.
|
||
reason = netguard.endpoint_block_reason(url)
|
||
if reason:
|
||
raise ProviderError(f"This endpoint can't be used — {reason}.")
|
||
log = debuglog.start_entry(url, self.model, body)
|
||
received: list[str] = []
|
||
try:
|
||
async with httpx.AsyncClient(timeout=httpx.Timeout(120, connect=10)) as client:
|
||
async with client.stream("POST", url, json=body, headers=self._headers()) as resp:
|
||
if resp.status_code != 200:
|
||
detail = (await resp.aread()).decode(errors="replace")[:500]
|
||
raise ProviderError(self._friendly_http_error(resp.status_code, detail))
|
||
# Some servers ignore stream=true and return one plain JSON
|
||
# body; buffer non-SSE lines so we can fall back to it.
|
||
saw_sse = False
|
||
raw_lines: list[str] = []
|
||
async for line in resp.aiter_lines():
|
||
if not line.startswith("data:"):
|
||
if not saw_sse:
|
||
raw_lines.append(line)
|
||
continue
|
||
saw_sse = True
|
||
data = line[5:].strip()
|
||
if data == "[DONE]":
|
||
debuglog.finish_entry(log, response="".join(received))
|
||
return
|
||
try:
|
||
payload = json.loads(data)
|
||
except ValueError:
|
||
continue
|
||
reasoning = self._extract_reasoning(payload)
|
||
if reasoning:
|
||
yield "reasoning", reasoning
|
||
chunk = self._extract_chunk(payload)
|
||
if chunk:
|
||
received.append(chunk)
|
||
yield "text", chunk
|
||
if not saw_sse:
|
||
body_text = "\n".join(raw_lines).strip()
|
||
try:
|
||
payload = json.loads(body_text)
|
||
except ValueError:
|
||
raise ProviderError(
|
||
"AI endpoint returned neither an SSE stream nor JSON: "
|
||
+ body_text[:200]
|
||
)
|
||
reasoning = self._extract_reasoning(payload)
|
||
if reasoning:
|
||
yield "reasoning", reasoning
|
||
chunk = self._extract_chunk(payload)
|
||
if chunk:
|
||
received.append(chunk)
|
||
yield "text", chunk
|
||
if not received:
|
||
raise ProviderError(
|
||
"AI endpoint returned a response with no text: "
|
||
+ body_text[:200]
|
||
)
|
||
debuglog.finish_entry(log, response="".join(received))
|
||
except httpx.ConnectError as exc:
|
||
error = f"Could not connect to {self.base_url} — is the AI server running?"
|
||
debuglog.finish_entry(log, response="".join(received), error=error)
|
||
raise ProviderError(error) from exc
|
||
except httpx.TimeoutException as exc:
|
||
debuglog.finish_entry(log, response="".join(received), error="Timed out")
|
||
raise ProviderError("The AI endpoint timed out.") from exc
|
||
except httpx.HTTPError as exc:
|
||
debuglog.finish_entry(log, response="".join(received), error=str(exc))
|
||
raise ProviderError(f"Request to AI endpoint failed: {exc}") from exc
|
||
except (ProviderError, GeneratorExit, BaseException) as exc:
|
||
status = "cancelled" if isinstance(exc, GeneratorExit) else str(exc)
|
||
debuglog.finish_entry(log, response="".join(received), error=status)
|
||
raise
|
||
|
||
async def complete(
|
||
self, system: str, user: str, *, temperature: float = 0.3, max_tokens: int = 400
|
||
) -> str:
|
||
"""Single non-streaming completion for background calls (summarization).
|
||
Unlike generate(), no story-continuation framing is added."""
|
||
if not self.model:
|
||
raise ProviderError("No model configured — set one in Settings.")
|
||
if self.api_mode == "completion":
|
||
url = f"{self.base_url}/completions"
|
||
body = {
|
||
"model": self.model,
|
||
"prompt": f"{system}\n\n{user}",
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens,
|
||
"stream": False,
|
||
}
|
||
else:
|
||
url = f"{self.base_url}/chat/completions"
|
||
body = {
|
||
"model": self.model,
|
||
"messages": [
|
||
{"role": "system", "content": system},
|
||
{"role": "user", "content": user},
|
||
],
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens,
|
||
"stream": False,
|
||
}
|
||
self._apply_reasoning_budget(body)
|
||
|
||
log = debuglog.start_entry(url, self.model, body)
|
||
try:
|
||
async with httpx.AsyncClient(timeout=httpx.Timeout(120, connect=10)) as client:
|
||
resp = await client.post(url, json=body, headers=self._headers())
|
||
except httpx.HTTPError as exc:
|
||
debuglog.finish_entry(log, error=str(exc))
|
||
raise ProviderError(f"Request to AI endpoint failed: {exc}") from exc
|
||
if resp.status_code != 200:
|
||
error = self._friendly_http_error(resp.status_code, resp.text[:500])
|
||
debuglog.finish_entry(log, error=error)
|
||
raise ProviderError(error)
|
||
try:
|
||
text = self._extract_chunk(resp.json())
|
||
except ValueError as exc:
|
||
debuglog.finish_entry(log, error="Invalid JSON response")
|
||
raise ProviderError("AI endpoint returned invalid JSON.") from exc
|
||
debuglog.finish_entry(log, response=text)
|
||
return text.strip()
|
||
|
||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||
"""POST /v1/embeddings; self.model is the embedding model here."""
|
||
if not self.model:
|
||
raise ProviderError("No embedding model configured — set one in Settings.")
|
||
url = f"{self.base_url}/embeddings"
|
||
body = {"model": self.model, "input": texts}
|
||
log = debuglog.start_entry(url, self.model, body)
|
||
try:
|
||
async with httpx.AsyncClient(timeout=httpx.Timeout(60, connect=10)) as client:
|
||
resp = await client.post(url, json=body, headers=self._headers())
|
||
except httpx.HTTPError as exc:
|
||
debuglog.finish_entry(log, error=str(exc))
|
||
raise ProviderError(f"Embedding request failed: {exc}") from exc
|
||
if resp.status_code != 200:
|
||
error = self._friendly_http_error(resp.status_code, resp.text[:500])
|
||
debuglog.finish_entry(log, error=error)
|
||
raise ProviderError(error)
|
||
try:
|
||
data = resp.json().get("data", [])
|
||
vectors = [item["embedding"] for item in sorted(data, key=lambda d: d.get("index", 0))]
|
||
except (ValueError, KeyError, TypeError) as exc:
|
||
debuglog.finish_entry(log, error="Malformed embeddings response")
|
||
raise ProviderError("AI endpoint returned malformed embeddings.") from exc
|
||
if len(vectors) != len(texts):
|
||
debuglog.finish_entry(log, error="Embedding count mismatch")
|
||
raise ProviderError("AI endpoint returned the wrong number of embeddings.")
|
||
debuglog.finish_entry(log, response=f"{len(vectors)} vectors × {len(vectors[0]) if vectors else 0} dims")
|
||
return vectors
|
||
|
||
def _friendly_http_error(self, status: int, detail: str) -> str:
|
||
if status == 401:
|
||
return "Authentication failed — check your API key in Settings."
|
||
if status == 404:
|
||
return (
|
||
f"Endpoint or model not found (HTTP 404). Check the endpoint URL and that "
|
||
f"model '{self.model}' exists. {detail}"
|
||
)
|
||
if status == 429:
|
||
# OpenRouter's shared free tier has a per-day cap; distinguish it
|
||
# from a short-term burst limit so the message is actionable.
|
||
if "free-models-per-day" in detail:
|
||
return (
|
||
"The free demo has hit its daily request limit (resets at "
|
||
"00:00 UTC). Please try again later."
|
||
)
|
||
return "The AI is getting too many requests right now — wait a moment and try again."
|
||
return f"AI endpoint returned HTTP {status}: {detail}"
|