diff --git a/README.md b/README.md index 474913f..dfe04da 100644 --- a/README.md +++ b/README.md @@ -124,6 +124,7 @@ because it inserts directly into the in-process store. | `RAG_MAX_QUESTION_CHARS` | `2000` | Maximum query length. | | `RAG_MAX_REQUEST_BYTES` | `4194304` | Maximum HTTP request body size (4 MiB), enforced before JSON parsing. | | `RAG_CLAUDE_TIMEOUT_SECONDS` | `30` | Claude request timeout; automatic SDK retries are disabled. | +| `RAG_SERVICE_TIMEOUT_MS` | `30000` | Next.js proxy timeout for calls to the Python service. | | `RAG_CORS_ORIGINS` | empty | Exact comma-separated allowed origins. Empty installs no CORS middleware; `*` restores wildcard access explicitly. | Chunking accepts `max_words` from 1 through 1000 and `overlap` from 0 diff --git a/app/app/lib/limits.ts b/app/app/lib/limits.ts index 9bbfe36..084c824 100644 --- a/app/app/lib/limits.ts +++ b/app/app/lib/limits.ts @@ -25,18 +25,43 @@ export class InputRequestError extends Error { } export async function readLimitedJson(request: Request): Promise { - const declaredLength = Number(request.headers.get("content-length")); - if (Number.isFinite(declaredLength) && declaredLength > INPUT_LIMITS.requestBytes) { - throw new InputRequestError("request body is too large", 413); + const contentLength = request.headers.get("content-length"); + if (contentLength !== null) { + const declaredLength = Number(contentLength); + if (!Number.isInteger(declaredLength) || declaredLength < 0) { + throw new InputRequestError("content-length header must be a non-negative integer", 400); + } + if (declaredLength > INPUT_LIMITS.requestBytes) { + throw new InputRequestError("request body is too large", 413); + } } - const bytes = await request.arrayBuffer(); - if (bytes.byteLength > INPUT_LIMITS.requestBytes) { - throw new InputRequestError("request body is too large", 413); + if (!request.body) { + throw new InputRequestError("request body must be valid JSON", 400); + } + + const reader = request.body.getReader(); + const chunks: Uint8Array[] = []; + let size = 0; + while (true) { + const { done, value } = await reader.read(); + if (done) break; + size += value.byteLength; + if (size > INPUT_LIMITS.requestBytes) { + await reader.cancel(); + throw new InputRequestError("request body is too large", 413); + } + chunks.push(value); } try { - return JSON.parse(new TextDecoder().decode(bytes)) as T; + const bytes = new Uint8Array(size); + let offset = 0; + for (const chunk of chunks) { + bytes.set(chunk, offset); + offset += chunk.byteLength; + } + return JSON.parse(new TextDecoder("utf-8", { fatal: true }).decode(bytes)) as T; } catch { throw new InputRequestError("request body must be valid JSON", 400); } diff --git a/app/app/lib/rag.ts b/app/app/lib/rag.ts index a9af27e..aebde6d 100644 --- a/app/app/lib/rag.ts +++ b/app/app/lib/rag.ts @@ -4,6 +4,7 @@ import { DEFAULT_RAG_SERVICE_URL } from "./config"; const SERVICE_URL = process.env.RAG_SERVICE_URL ?? DEFAULT_RAG_SERVICE_URL; +const SERVICE_TIMEOUT_MS = 30_000; export interface Source { chunk_id: number; @@ -57,6 +58,11 @@ export class RagServiceError extends Error { } } +function serviceTimeoutMs(): number { + const value = Number(process.env.RAG_SERVICE_TIMEOUT_MS); + return Number.isInteger(value) && value > 0 ? value : SERVICE_TIMEOUT_MS; +} + function errorMessage(body: string, status: number): string { try { const parsed = JSON.parse(body) as { detail?: unknown; error?: unknown }; @@ -82,11 +88,20 @@ function errorMessage(body: string, status: number): string { } async function call(path: string, init?: RequestInit): Promise { - const response = await fetch(`${SERVICE_URL}${path}`, { - ...init, - headers: { "Content-Type": "application/json", ...(init?.headers ?? {}) }, - cache: "no-store", - }); + let response: Response; + try { + response = await fetch(`${SERVICE_URL}${path}`, { + ...init, + headers: { "Content-Type": "application/json", ...(init?.headers ?? {}) }, + cache: "no-store", + signal: init?.signal ?? AbortSignal.timeout(serviceTimeoutMs()), + }); + } catch (error) { + if (error instanceof Error && (error.name === "TimeoutError" || error.name === "AbortError")) { + throw new RagServiceError("RAG service request timed out", 504); + } + throw error; + } if (!response.ok) { const body = await response.text(); throw new RagServiceError(errorMessage(body, response.status), response.status); diff --git a/service/rag_service/app.py b/service/rag_service/app.py index 9a7e479..766f04e 100644 --- a/service/rag_service/app.py +++ b/service/rag_service/app.py @@ -16,9 +16,10 @@ import contextlib import logging +import math import os import pathlib -from typing import AsyncIterator, List, Optional +from typing import AsyncIterator, List from fastapi import Depends, FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware @@ -61,6 +62,17 @@ def _positive_env_int(name: str, default: int) -> int: return value +def _finite_env_float(name: str, default: float) -> float: + raw = os.environ.get(name, str(default)) + try: + value = float(raw) + except ValueError as exc: + raise RuntimeError(f"{name} must be a number") from exc + if not math.isfinite(value): + raise RuntimeError(f"{name} must be finite") + return value + + def _add_cors_middleware(target: FastAPI, origins: List[str]) -> None: if not origins: return @@ -94,11 +106,18 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: raw_length = headers.get(b"content-length") if raw_length is not None: try: - if int(raw_length) > self.max_bytes: + content_length = int(raw_length) + if content_length < 0: + raise ValueError + if content_length > self.max_bytes: await self._reject(scope, receive, send) return except ValueError: - pass + response = JSONResponse( + {"detail": "invalid Content-Length header"}, status_code=400 + ) + await response(scope, receive, send) + return messages: List[Message] = [] total = 0 @@ -140,7 +159,7 @@ async def _reject(self, scope: Scope, receive: Receive, send: Send) -> None: _embedder = get_embedder(os.environ.get("RAG_EMBEDDER", "auto")) _store = DocumentStore( embedder=_embedder, - min_score=float(os.environ.get("RAG_MIN_SCORE", "0.09")), + min_score=_finite_env_float("RAG_MIN_SCORE", 0.09), ) SAMPLE_DOCS = pathlib.Path(__file__).resolve().parent.parent / "sample_docs" diff --git a/service/rag_service/generation.py b/service/rag_service/generation.py index 08c5878..6875f84 100644 --- a/service/rag_service/generation.py +++ b/service/rag_service/generation.py @@ -8,6 +8,7 @@ from __future__ import annotations +import math import os import re from dataclasses import dataclass @@ -49,7 +50,7 @@ def _positive_env_float(name: str, default: float) -> float: value = float(raw) except ValueError as exc: raise RuntimeError(f"{name} must be a number") from exc - if value <= 0: + if not math.isfinite(value) or value <= 0: raise RuntimeError(f"{name} must be greater than zero") return value diff --git a/service/rag_service/store.py b/service/rag_service/store.py index 86339dd..cbd2cf4 100644 --- a/service/rag_service/store.py +++ b/service/rag_service/store.py @@ -82,6 +82,10 @@ def add_document( return doc vectors = self.embedder.embed([c.text for c in chunks]) + if len(vectors) != len(chunks): + raise ValueError( + f"embedder returned {len(vectors)} vectors for {len(chunks)} chunks" + ) with self._lock: if self._index is None: @@ -95,6 +99,10 @@ def add_document( doc_id = self._next_doc_id self._next_doc_id += 1 ids = self._index.insert_batch(vectors) + if len(ids) != len(chunks): + raise RuntimeError( + f"index returned {len(ids)} ids for {len(chunks)} chunks" + ) for chunk, cid in zip(chunks, ids): self._chunks[cid] = StoredChunk( id=cid, @@ -109,33 +117,34 @@ def add_document( def retrieve(self, query: str, k: int = 5, ef_search: int = 100) -> List[RetrievedChunk]: """Embed the query and return the k most relevant chunks.""" + qvec = self.embedder.embed([query])[0] with self._lock: if self._index is None or len(self._index) == 0: return [] - index = self._index - - qvec = self.embedder.embed([query])[0] - hits = index.search(qvec, k=k, ef_search=ef_search) - out: List[RetrievedChunk] = [] - for cid, distance in hits: - sc = self._chunks[cid] - # Cosine distance is 1 - similarity; report similarity so higher - # is more relevant, which is what a reader expects from a score. - score = 1.0 - distance if self.metric == "cosine" else -distance - if self.min_score is not None and score < self.min_score: - # Hits are closest-first, so their relevance scores only - # decrease. Weak nearest neighbors are not useful grounding. - break - out.append( - RetrievedChunk( - id=sc.id, - text=sc.text, - doc_id=sc.doc_id, - doc_title=sc.doc_title, - ordinal=sc.ordinal, - score=score, + # The Python wrapper exposes insert as a mutable operation. Keep + # search and its metadata lookup in the same critical section so + # concurrent uploads cannot mutate the graph beneath a query. + hits = self._index.search(qvec, k=k, ef_search=ef_search) + out: List[RetrievedChunk] = [] + for cid, distance in hits: + sc = self._chunks[cid] + # Cosine distance is 1 - similarity; report similarity so higher + # is more relevant, which is what a reader expects from a score. + score = 1.0 - distance if self.metric == "cosine" else -distance + if self.min_score is not None and score < self.min_score: + # Hits are closest-first, so their relevance scores only + # decrease. Weak nearest neighbors are not useful grounding. + break + out.append( + RetrievedChunk( + id=sc.id, + text=sc.text, + doc_id=sc.doc_id, + doc_title=sc.doc_title, + ordinal=sc.ordinal, + score=score, + ) ) - ) return out def stats(self) -> dict: diff --git a/service/tests/test_e2e.py b/service/tests/test_e2e.py index 881c9cc..eca9045 100644 --- a/service/tests/test_e2e.py +++ b/service/tests/test_e2e.py @@ -136,6 +136,17 @@ def test_oversized_http_body_returns_413_before_json_parsing(client): assert "exceeds" in response.json()["detail"] +@pytest.mark.parametrize("content_length", ["not-a-number", "-1"]) +def test_invalid_content_length_returns_400(client, content_length): + response = client.post( + "/query", + content=b'{}', + headers={"Content-Length": content_length, "Content-Type": "application/json"}, + ) + assert response.status_code == 400 + assert response.json()["detail"] == "invalid Content-Length header" + + def test_claude_timeout_is_exposed_as_504(client, monkeypatch): app_module = importlib.import_module("rag_service.app")