From 9730b8cc14b058a564d78afed5488eb2df50b4e5 Mon Sep 17 00:00:00 2001 From: drag0sd0g Date: Thu, 30 Jul 2026 05:56:59 +0100 Subject: [PATCH] Performance Uplifts --- .env.example | 43 +++ README.ja.md | 28 +- README.md | 28 +- docker-compose.yml | 70 +++-- docs/design-decisions.ja.md | 9 +- docs/design-decisions.md | 9 +- docs/evaluation-results.ja.md | 2 +- docs/evaluation-results.md | 2 +- docs/technical-design-document.ja.md | 4 +- docs/technical-design-document.md | 4 +- examples/README.ja.md | 2 +- examples/README.md | 2 +- helm/findoc-rag/templates/configmaps.yaml | 1 + helm/findoc-rag/tests/configmaps_test.yaml | 18 ++ helm/findoc-rag/values.yaml | 3 + services/embedding-worker/src/chunker.py | 60 ++++- .../embedding-worker/src/drift_detector.py | 12 +- .../embedding-worker/tests/test_chunker.py | 90 +++++++ .../tests/test_drift_detector.py | 50 ++++ services/ingestion/src/kafka_producer.py | 93 +++++-- services/ingestion/src/main.py | 43 ++- services/ingestion/tests/test_edgar_client.py | 209 ++++++++++++++- services/query-api/src/llm/ollama_backend.py | 42 +++ services/query-api/src/main.py | 25 +- services/query-api/src/rag/generator.py | 17 +- services/query-api/src/rag/retriever.py | 99 ++++++- services/query-api/tests/test_generator.py | 118 +++++++++ services/query-api/tests/test_llm_backends.py | 85 ++++++ services/query-api/tests/test_main.py | 32 ++- services/query-api/tests/test_retriever.py | 250 +++++++++++++++++- 30 files changed, 1347 insertions(+), 103 deletions(-) diff --git a/.env.example b/.env.example index 82cd827..14a9e6c 100644 --- a/.env.example +++ b/.env.example @@ -18,8 +18,28 @@ KAFKA_BOOTSTRAP_SERVERS=kafka:9092 LLM_BACKEND=ollama # --- Ollama (if LLM_BACKEND=ollama) --- +# Containerised Ollama on the compose network (default). NOTE: on macOS and +# Windows this runs CPU-only — Docker's Linux VM cannot reach the host GPU. OLLAMA_URL=http://ollama:11434 +# For GPU-accelerated local inference, run Ollama natively on the host, start +# the stack WITHOUT the local-llm profile, and use the host address instead: +# OLLAMA_URL=http://host.docker.internal:11434 +# Some runtimes (OrbStack) forward host.docker.internal to the host loopback, +# so a server bound to 127.0.0.1 works as-is. On Docker Desktop the native +# server must listen on all interfaces instead, which exposes it to the LAN: +# OLLAMA_HOST=0.0.0.0:11434 ollama serve OLLAMA_MODEL=mistral:7b +# Host port for the containerised Ollama. Deliberately NOT 11434 by default: +# the container binds all interfaces (including IPv6), so publishing 11434 +# would shadow a natively installed Ollama whenever you use "localhost". +# OLLAMA_HOST_PORT=127.0.0.1:11435 +# Context window Ollama allocates for a request. The whole RAG prompt must fit: +# roughly top_k x 512 tokens of chunks, plus injected XBRL facts, plus the +# system prompt, plus the 1024-token answer budget. Ollama's own default (4096) +# truncates the START of the prompt — silently dropping the system prompt and +# the citation instruction. 8192 covers top_k up to ~10; raise it for larger +# top_k (costs RAM for the KV cache). +OLLAMA_NUM_CTX=8192 # --- OpenAI (if LLM_BACKEND=openai) --- OPENAI_API_KEY= @@ -69,5 +89,28 @@ EDGAR_RATE_LIMIT_RPS=10 # EVAL_OLLAMA_URL=http://localhost:11434 # EVAL_OLLAMA_MODEL=mistral:7b +# --- Resource limits / performance tuning --- +# Defaults are conservative so the stack starts on a modest machine. Scale +# these to your hardware. A memory limit larger than the container runtime's +# own VM allowance is silently ignored — check `docker info` first. +# POSTGRES_CPUS=4.0 +# POSTGRES_MEMORY=8G +# POSTGRES_SHARED_BUFFERS=2GB +# POSTGRES_WORK_MEM=256MB +# POSTGRES_MAINTENANCE_WORK_MEM=1GB +# POSTGRES_EFFECTIVE_CACHE_SIZE=4GB +# POSTGRES_PARALLEL_WORKERS_PER_GATHER=2 +# KAFKA_CPUS=2.0 +# KAFKA_MEMORY=4G +# INGESTION_CPUS=2.0 +# INGESTION_MEMORY=2G +# EMBEDDING_WORKER_CPUS=12.0 +# EMBEDDING_WORKER_MEMORY=6G +# Retrieval runs in a thread pool, so this caps concurrent query throughput. +# QUERY_API_CPUS=4.0 +# QUERY_API_MEMORY=4G +# OLLAMA_CPUS=8.0 +# OLLAMA_MEMORY=8G + # --- Logging --- LOG_LEVEL=INFO \ No newline at end of file diff --git a/README.ja.md b/README.ja.md index 63a93df..5dcbe80 100644 --- a/README.ja.md +++ b/README.ja.md @@ -30,7 +30,9 @@ API Client <--> Query API <--> Ollama / OpenAI / Claude | - Python 3.12 以降（ローカル開発および Docker 外でのテスト実行用） - GNU Make（任意、便利なターゲット用） -**ハードウェア（Ollama / ローカル LLM パスの場合のみ）：** `make run` は `mistral:7b`（約 4 GB のダウンロード）を取得し、実行時に少なくとも **8 GB の空き RAM** が必要です（モデルの重み + Docker オーバーヘッド）。RAM が不足しているマシンでは Ollama コンテナが OOM で強制終了されます。リモート LLM パス（`LLM_BACKEND=claude` または `openai` を指定した `make run-remote`）では、ベーススタック（約 2 GB）以外に GPU や RAM の要件はありません。 +**ハードウェア：** 本スタックは実行するマシンに合わせてサイズを調整できます。コンテナの CPU およびメモリ上限、埋め込みのバッチサイズ、Ollama のコンテキストウィンドウはすべて環境変数です（[設定](#設定)を参照）。手持ちのハードウェアに合わせてスケールしてください。必要なリソースは、選択する LLM バックエンドにほぼ完全に依存します。`LLM_BACKEND=claude` または `openai` を指定した `make run-remote` はローカル推論を一切行いません。一方、ローカル LLM パスは選択したモデルによって決まります。 + +ローカルでベンチマークを取る前に知っておくべき注意点が 1 つあります。**macOS および Windows では、コンテナ化された Ollama は CPU のみで動作します**。Docker の Linux VM はホストの GPU にアクセスできないためです。ローカル生成を大幅に高速化するには、ホスト上で Ollama をネイティブに実行し、`OLLAMA_URL` をそこに向けてください。[LLM バックエンド](#llm-バックエンド)を参照してください。 ## クイックスタート @@ -64,7 +66,7 @@ make run-remote 初回起動時は、コンテナイメージのプル、3 つのサービスのビルド、データベースマイグレーションの実行が行われます。スキーマは `db/migrations/001_initial_schema.sql` から自動的に適用されます（`pgvector` 拡張、`ingestion_log` テーブル、`document_chunks` テーブル、HNSW ベクターインデックスを作成します）。マイグレーションを手動で実行するには：`make migrate` -Ollama を使用する場合（`make run`）、`mistral:7b` モデル（約 4 GB）は初回起動時に自動的にダウンロードされます。2 回目以降の起動ではキャッシュされたレイヤーとボリュームが再利用されます。 +Ollama を使用する場合（`make run`）、`OLLAMA_MODEL` で指定したモデルが初回起動時に自動的にダウンロードされます。2 回目以降の起動ではキャッシュされたレイヤーとボリュームが再利用されます。 4. サービスが起動していることを確認します。 @@ -350,7 +352,20 @@ make helm-teardown | `OPENAI_MODEL` | `gpt-4o-mini` | OpenAI モデル名 | | `ANTHROPIC_API_KEY` | （空） | `LLM_BACKEND=claude` 時に必須 | | `CLAUDE_MODEL` | `claude-opus-4-6` | Claude モデル名 | -| `EMBEDDING_MODEL` | `sentence-transformers/all-MiniLM-L6-v2` | 埋め込み用 sentence-transformers モデル | +| `EMBEDDING_MODEL` | `nomic-ai/nomic-embed-text-v1.5` | 埋め込み用 sentence-transformers モデル（768 次元） | + +**パフォーマンスチューニング** + +実行するマシンに合わせてスケールしてください。デフォルト値はどの環境でも起動できるよう控えめに設定されています。 + +| 変数 | デフォルト | 説明 | +| ------------------------- | ---------- | ---------------------------------------------------------------------------------------------------------------- | +| `OLLAMA_NUM_CTX` | `8192` | Ollama に要求するコンテキストウィンドウ。RAG プロンプト全体が収まらない場合、先頭から切り捨てられシステムプロンプトが失われます。 | +| `EMBEDDING_BATCH_SIZE` | `64` | 埋め込みワーカーが 1 バッチで処理するチャンク数。大きくするとスループットが向上しますが、メモリを消費します。 | +| `POSTGRES_SHARED_BUFFERS` | `2GB` | PostgreSQL のページキャッシュ。コーパスがこれを超えてから引き上げる価値があります。 | +| `POSTGRES_WORK_MEM` | `256MB` | PostgreSQL の操作あたりのソート / ハッシュ用メモリ。 | +| `QUERY_API_CPUS` | `4.0` | Query API コンテナの CPU 上限。検索はスレッドプールで実行されるため、同時クエリのスループットを制限します。 | +| `EMBEDDING_WORKER_CPUS` | `12.0` | 埋め込みワーカーコンテナの CPU 上限。 | **セキュリティと API** @@ -389,7 +404,12 @@ make helm-teardown Query API は `LLM_BACKEND` 環境変数で選択可能な 3 つの LLM バックエンドをサポートします。 -**Ollama（デフォルト）** -- Docker Compose スタック内でローカルに実行されます。API キー不要。モデルは初回起動時に自動的に取得されます。`make run`（または `docker compose --profile local-llm up`）で起動します。開発環境およびセルフホスト型デプロイに適しています。`mistral:7b` のモデルウェイトには約 6 GB の RAM が必要です（Docker オーバーヘッドを含む合計では約 8 GB）。 +**Ollama（デフォルト）** -- ローカルで実行され、API キーは不要です。`OLLAMA_MODEL` で指定したモデルは初回起動時に自動的に取得されます。開発環境およびセルフホスト型デプロイに適しています。実行方法は 2 通りあります。 + +- *コンテナ化*（`make run`、または `docker compose --profile local-llm up`）— セットアップ不要で完全にポータブルです。ただし macOS および Windows では推論が **CPU のみ** で実行されます。Docker の Linux VM はホストの GPU にアクセスできないため、コンテナに CPU やメモリをどれだけ割り当てても生成は低速です。 +- *ホストネイティブ* — ホスト上で `ollama serve` を実行してホストの GPU（macOS では Metal、Linux では CUDA / ROCm）を利用し、`local-llm` プロファイル **なし** でスタックを起動して（`make run-remote`）、`OLLAMA_URL=http://host.docker.internal:11434` を設定します。この方法は通常、コンテナ化パスより一桁高速です。一部のランタイム（OrbStack）は `host.docker.internal` をホストのループバックに転送するため、`127.0.0.1` にバインドされたサーバーにそのまま到達できます。Docker Desktop の場合はネイティブサーバーをすべてのインターフェースで待ち受けさせる必要があります（`OLLAMA_HOST=0.0.0.0:11434`）。これはローカルネットワークに公開されることを意味します。 + +`OLLAMA_NUM_CTX` はモデルとマシンに合わせて設定してください。RAG プロンプト全体がその中に収まらない場合、Ollama はプロンプトを先頭から静かに切り捨て、最初にシステム指示が失われます。 **OpenAI** -- OpenAI のチャット補完 API を呼び出します。`LLM_BACKEND=openai` を設定し、有効な `OPENAI_API_KEY` を提供します。デフォルトでは `gpt-4o-mini` を使用します。`make run-remote` で起動します。より高品質な回答や評価比較に有用です。 diff --git a/README.md b/README.md index e6d7395..5f1680b 100644 --- a/README.md +++ b/README.md @@ -30,7 +30,9 @@ API Client <--> Query API <--> Ollama / OpenAI / Claude | - Python 3.12+ (for local development and running tests outside of Docker) - GNU Make (optional, for convenience targets) -**Hardware (Ollama / local LLM path only):** `make run` pulls `mistral:7b` (~4 GB download) and requires at least **8 GB of free RAM** at runtime (model weights + Docker overhead). On machines with less RAM the Ollama container will be OOM-killed. The remote-LLM path (`make run-remote` with `LLM_BACKEND=claude` or `openai`) has no GPU or RAM requirements beyond the base stack (~2 GB). +**Hardware:** The stack sizes itself to the machine it runs on — container CPU and memory limits, the embedding batch size, and the Ollama context window are all environment variables (see [Configuration](#configuration)), so scale them to the hardware you have. Requirements depend almost entirely on which LLM backend you choose: `make run-remote` with `LLM_BACKEND=claude` or `openai` performs no local inference at all, while the local-LLM path is bounded by the model you select. + +One caveat worth knowing before benchmarking locally: **containerised Ollama on macOS and Windows runs on CPU only**, because Docker's Linux VM cannot reach the host GPU. For substantially faster local generation, run Ollama natively on the host and point `OLLAMA_URL` at it — see [LLM Backends](#llm-backends). ## Quick Start @@ -64,7 +66,7 @@ make run-remote On the first run this will pull container images, build the three services, and run database migrations. The schema is applied automatically from `db/migrations/001_initial_schema.sql` (creates the `pgvector` extension, `ingestion_log` and `document_chunks` tables, and the HNSW vector index). To run migrations manually at any time: `make migrate`. -If using Ollama (`make run`), the `mistral:7b` model (~4 GB) is downloaded automatically on first start. Subsequent starts reuse cached layers and volumes. +If using Ollama (`make run`), the model named by `OLLAMA_MODEL` is downloaded automatically on first start. Subsequent starts reuse cached layers and volumes. 4. Verify the services are running: @@ -350,7 +352,20 @@ All configuration is driven by environment variables. See `.env.example` for the | `OPENAI_MODEL` | `gpt-4o-mini` | OpenAI model name | | `ANTHROPIC_API_KEY` | (empty) | Required when `LLM_BACKEND=claude` | | `CLAUDE_MODEL` | `claude-opus-4-6` | Claude model name | -| `EMBEDDING_MODEL` | `sentence-transformers/all-MiniLM-L6-v2` | Sentence-transformers model for embedding | +| `EMBEDDING_MODEL` | `nomic-ai/nomic-embed-text-v1.5` | Sentence-transformers model for embedding (768-dim) | + +**Performance tuning** + +Scale these to the machine you are running on; the defaults are conservative so the stack starts anywhere. + +| Variable | Default | Description | +| ----------------------- | ------- | --------------------------------------------------------------------------------------------------------------------------------------- | +| `OLLAMA_NUM_CTX` | `8192` | Context window requested from Ollama. The entire RAG prompt must fit or Ollama truncates it from the start, dropping the system prompt. | +| `EMBEDDING_BATCH_SIZE` | `64` | Chunks embedded per batch in the embedding worker. Larger batches raise throughput at the cost of memory. | +| `POSTGRES_SHARED_BUFFERS` | `2GB` | PostgreSQL page cache. Worth raising only once the corpus outgrows it. | +| `POSTGRES_WORK_MEM` | `256MB` | Per-operation sort/hash memory for PostgreSQL. | +| `QUERY_API_CPUS` | `4.0` | CPU limit for the Query API container. Retrieval runs in a thread pool, so this caps concurrent query throughput. | +| `EMBEDDING_WORKER_CPUS` | `12.0` | CPU limit for the embedding worker container. | **Security & API** @@ -389,7 +404,12 @@ All configuration is driven by environment variables. See `.env.example` for the The Query API supports 3 LLM backends, selectable via the `LLM_BACKEND` environment variable: -**Ollama (default)** -- Runs locally inside the Docker Compose stack. No API key required. The model is pulled automatically on first start. Start with `make run` (or `docker compose --profile local-llm up`). Suitable for development and self-hosted deployments. Requires ~6 GB RAM for the `mistral:7b` model weights (~8 GB total including Docker overhead). +**Ollama (default)** -- Runs locally, with no API key required; the model named by `OLLAMA_MODEL` is pulled automatically on first start. Suitable for development and self-hosted deployments. There are two ways to run it: + +- *Containerised* (`make run`, or `docker compose --profile local-llm up`) — zero setup, fully portable. On macOS and Windows this runs inference **on CPU only**: Docker's Linux VM has no access to the host GPU, so generation is slow regardless of how much CPU or memory you give the container. +- *Host-native* — run `ollama serve` on the host so it uses the host GPU (Metal on macOS, CUDA/ROCm on Linux), then start the stack **without** the `local-llm` profile (`make run-remote`) and set `OLLAMA_URL=http://host.docker.internal:11434`. This is typically an order of magnitude faster than the containerised path. Some runtimes (OrbStack) forward `host.docker.internal` to the host loopback, so a server bound to `127.0.0.1` is reachable as-is; on Docker Desktop the native server must listen on all interfaces instead (`OLLAMA_HOST=0.0.0.0:11434`), which does expose it to your local network. + +Size `OLLAMA_NUM_CTX` to your model and machine — the whole RAG prompt must fit inside it, or Ollama silently truncates the prompt from the start, discarding the system instructions first. **OpenAI** -- Calls the OpenAI chat completions API. Set `LLM_BACKEND=openai` and provide a valid `OPENAI_API_KEY`. Uses `gpt-4o-mini` by default. Start with `make run-remote`. Useful for higher-quality answers and evaluation comparisons. diff --git a/docker-compose.yml b/docker-compose.yml index d58f424..34953e6 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -2,6 +2,14 @@ # FinDocDRAG — Full Local Stack # Usage: docker compose up # See: docs/technical-design-document.md Section 10.2 +# +# CPU and memory limits below are environment variables with conservative +# defaults so the stack starts on a modest machine. Raise them in .env to +# suit the host — see the "Performance tuning" table in README.md. +# +# Note: a limit larger than the container runtime's own memory allowance is +# silently ignored (Docker Desktop / OrbStack run a fixed-size Linux VM), so +# check `docker info` before setting large values here. # ============================================================ services: @@ -21,18 +29,20 @@ services: command: - postgres - -c - - shared_buffers=2GB + - shared_buffers=${POSTGRES_SHARED_BUFFERS:-2GB} - -c - - work_mem=256MB + - work_mem=${POSTGRES_WORK_MEM:-256MB} - -c - - maintenance_work_mem=1GB + - maintenance_work_mem=${POSTGRES_MAINTENANCE_WORK_MEM:-1GB} - -c - - effective_cache_size=4GB + - effective_cache_size=${POSTGRES_EFFECTIVE_CACHE_SIZE:-4GB} + - -c + - max_parallel_workers_per_gather=${POSTGRES_PARALLEL_WORKERS_PER_GATHER:-2} deploy: resources: limits: - cpus: "4.0" - memory: 8G + cpus: "${POSTGRES_CPUS:-4.0}" + memory: ${POSTGRES_MEMORY:-8G} healthcheck: test: ["CMD-SHELL", "pg_isready -U findocdrag -d findocdrag"] interval: 5s @@ -95,8 +105,8 @@ services: deploy: resources: limits: - cpus: "2.0" - memory: 4G + cpus: "${KAFKA_CPUS:-2.0}" + memory: ${KAFKA_MEMORY:-4G} healthcheck: test: ["CMD-SHELL", "/opt/kafka/bin/kafka-topics.sh --bootstrap-server localhost:9092 --list || exit 1"] interval: 10s @@ -132,6 +142,12 @@ services: # Only started when the "local-llm" profile is active: # docker compose --profile local-llm up -d # Omit the profile flag (or set LLM_BACKEND=claude/openai) to skip Ollama entirely. + # + # IMPORTANT: on macOS and Windows this container has no access to the host + # GPU — Docker runs it inside a Linux VM — so inference is CPU-only and slow + # no matter how much CPU/memory it is given. For GPU-accelerated local + # inference, run Ollama natively on the host, start the stack WITHOUT this + # profile, and set OLLAMA_URL=http://host.docker.internal:11434. ollama: profiles: [local-llm] image: ollama/ollama:latest @@ -139,10 +155,16 @@ services: deploy: resources: limits: - cpus: "8.0" - memory: 16G + cpus: "${OLLAMA_CPUS:-8.0}" + # Must stay within the container runtime's VM allowance or Docker + # silently drops the limit. Size it to the model plus its KV cache. + memory: ${OLLAMA_MEMORY:-8G} + # Not published by default: binding the host's 11434 shadows a natively + # installed Ollama on that port (the container binds all interfaces, + # including IPv6, so "localhost:11434" resolves to the container). Set + # OLLAMA_HOST_PORT to expose it deliberately. ports: - - "11434:11434" + - "${OLLAMA_HOST_PORT:-127.0.0.1:11435}:11434" volumes: - ollama_models:/root/.ollama healthcheck: @@ -177,8 +199,8 @@ services: deploy: resources: limits: - cpus: "2.0" - memory: 2G + cpus: "${INGESTION_CPUS:-2.0}" + memory: ${INGESTION_MEMORY:-2G} depends_on: db-migrate: condition: service_completed_successfully @@ -213,8 +235,9 @@ services: deploy: resources: limits: - cpus: "12.0" - memory: 16G + cpus: "${EMBEDDING_WORKER_CPUS:-12.0}" + # Keep within the runtime's VM allowance; a larger value is ignored. + memory: ${EMBEDDING_WORKER_MEMORY:-6G} depends_on: db-migrate: condition: service_completed_successfully @@ -248,11 +271,17 @@ services: deploy: resources: limits: - cpus: "4.0" - memory: 4G + # Retrieval (embedding + pgvector) runs in a thread pool, so this + # caps how much concurrent query work the service can actually do. + cpus: "${QUERY_API_CPUS:-4.0}" + memory: ${QUERY_API_MEMORY:-4G} depends_on: db-migrate: condition: service_completed_successfully + # Lets OLLAMA_URL reach a host-native Ollama. Docker Desktop and OrbStack + # resolve host.docker.internal already; this adds it on Linux too. + extra_hosts: + - "host.docker.internal:host-gateway" environment: POSTGRES_HOST: postgres POSTGRES_PORT: "5432" @@ -260,8 +289,13 @@ services: POSTGRES_USER: findocdrag POSTGRES_PASSWORD: changeme LLM_BACKEND: "${LLM_BACKEND:-ollama}" - OLLAMA_URL: "http://ollama:11434" + # Defaults to the containerised Ollama on the compose network. Point this + # at http://host.docker.internal:11434 to use a host-native Ollama (the + # only way to get GPU acceleration on macOS/Windows); the native server + # must then listen on all interfaces, e.g. OLLAMA_HOST=0.0.0.0:11434. + OLLAMA_URL: "${OLLAMA_URL:-http://ollama:11434}" OLLAMA_MODEL: "${OLLAMA_MODEL:-mistral:7b}" + OLLAMA_NUM_CTX: "${OLLAMA_NUM_CTX:-8192}" OPENAI_API_KEY: "${OPENAI_API_KEY:-}" ANTHROPIC_API_KEY: "${ANTHROPIC_API_KEY:-}" CLAUDE_MODEL: "${CLAUDE_MODEL:-claude-opus-4-6}" diff --git a/docs/design-decisions.ja.md b/docs/design-decisions.ja.md index f630acd..4985ec4 100644 --- a/docs/design-decisions.ja.md +++ b/docs/design-decisions.ja.md @@ -128,19 +128,20 @@ Query API は、検索されたコンテキストから引用付きの回答を トリプルバックエンドアプローチは複数のデプロイメントシナリオを満たします。ローカル開発とポートフォリオデモンストレーションでは、システムは `docker compose --profile local-llm up` で起動し、API キーなしで動作しなければなりません -- Ollama はこれを提供します。評価と品質比較のために、OpenAI の GPT-4o-mini と Anthropic の Claude はどちらも、ローカルで実行される 7B モデルより測定可能なほど優れた回答を生成します。Claude は OpenAI キーを持たず Anthropic サブスクリプションを持つ開発者に特に有用であり、その逆も同様です。 -Ollama は Docker Compose プロファイル（`--profile local-llm`）を通じてオプトイン方式となっています。プロファイルなしで実行すると Ollama コンテナが完全にスキップされ、スタックの RAM フットプリントが約 6 GB 削減され、数 GB のモデルダウンロードを回避できます -- これによりリモートバックエンドのワークフローの起動が大幅に高速化されます。 +Ollama は Docker Compose プロファイル（`--profile local-llm`）を通じてオプトイン方式となっています。プロファイルなしで実行すると Ollama コンテナが完全にスキップされ、選択したモデルが占有していたリソースが解放され、数 GB のモデルダウンロードを回避できます -- これによりリモートバックエンドのワークフローの起動が大幅に高速化されます。また、このパスはスタックをホストネイティブの Ollama に向ける場合にも使用します。macOS および Windows のコンテナはホストの GPU にアクセスできないため、これがローカルバックエンドで GPU アクセラレーションを得る唯一の方法です。 LangChain ベースの抽象化は、最終的には単一の `generate(prompt) -> text` 呼び出しに対して大きなトランザクティブな依存関係ツリーをもたらすため却下されました。カスタムの `LLMBackend` Protocol は 10 行のコードで、フレームワーク結合なしに必要なことを正確に実行します。 -vLLM と TGI は、GPU バックアップされた高スループットサービング向けに設計された本番推論サーバーです。CPU のみの Docker Compose デモ環境には適していません。 +vLLM と TGI は、GPU バックアップされた高スループットサービング向けに設計された本番推論サーバーです。ポータブルな Docker Compose 環境に対しては本システムの必要以上に重く、そもそもコンテナ化されたランタイムには GPU アクセスがありません。 ### 結果 - 利点: システムはそのままの状態でオフラインで完全に動作し（Ollama プロファイル）、Ollama が現実的でない場合はリモート API でも動作する。 - 利点: 評価ハーネスが Ollama、OpenAI、Claude の品質を並べて比較できる。 - 利点: 最小限の抽象化レイヤー（Protocol + 3 つの実装）は理解とテストが容易。 -- 利点: Ollama は Docker Compose プロファイルを通じてオプトイン; リモートバックエンドのセットアップでは約 6 GB のモデルダウンロードをスキップできる。 -- 欠点: Ollama と Mistral 7B は約 6 GB の RAM を必要とし、メモリ制約の環境では重大（TinyLlama は 637 MB でドキュメント化されたフォールバックとして機能する）。 +- 利点: Ollama は Docker Compose プロファイルを通じてオプトイン; リモートバックエンドのセットアップではモデルダウンロードを完全にスキップできる。 +- 利点: 同じプロファイル切り替えでホストネイティブの Ollama を利用できる。これは macOS および Windows で GPU アクセラレーションされたローカル推論を得る唯一の方法である。 +- 欠点: ローカルモデルのサイズとメモリ使用量は `OLLAMA_MODEL` で選択したモデルに応じて変わる。制約のあるマシンではより小さなモデルがドキュメント化されたフォールバックとなる。 - 欠点: 3 つのコードパスにより、3 つのバックエンドすべてにテストカバレッジが必要になる。 --- diff --git a/docs/design-decisions.md b/docs/design-decisions.md index b39757d..3b50944 100644 --- a/docs/design-decisions.md +++ b/docs/design-decisions.md @@ -128,19 +128,20 @@ The Query API needs an LLM to generate cited answers from retrieved context. The The triple-backend approach satisfies multiple deployment scenarios. For local development and portfolio demonstration, the system must start with `docker compose --profile local-llm up` and work without any API key -- Ollama provides this. For evaluation and quality comparison, both OpenAI's GPT-4o-mini and Anthropic's Claude produce measurably better answers than a locally-run 7B model. Claude is particularly useful for developers who have an Anthropic subscription but not an OpenAI key, and vice versa. -Ollama is made opt-in via Docker Compose profiles (`--profile local-llm`). Running without the profile skips the Ollama container entirely, reducing the stack's RAM footprint by ~6 GB and avoiding a multi-GB model download -- this makes remote-backend workflows significantly faster to start. +Ollama is made opt-in via Docker Compose profiles (`--profile local-llm`). Running without the profile skips the Ollama container entirely, freeing whatever the chosen model would have occupied and avoiding a multi-GB model download -- this makes remote-backend workflows significantly faster to start. It is also the path used when pointing the stack at a host-native Ollama, which is how the local backend gets GPU acceleration: containers on macOS and Windows cannot reach the host GPU. A LangChain-based abstraction was rejected because it introduces a large transitive dependency tree for what is ultimately a single `generate(prompt) -> text` call. The custom `LLMBackend` Protocol is 10 lines of code and does exactly what we need without framework coupling. -vLLM and TGI are production inference servers designed for GPU-backed, high-throughput serving. They are not suitable for a CPU-only Docker Compose demo environment. +vLLM and TGI are production inference servers designed for GPU-backed, high-throughput serving. They are heavier than this system needs for a portable Docker Compose environment, where the containerised runtime has no GPU access anyway. ### Consequences - Positive: System works fully offline out of the box (Ollama profile) and with remote APIs when Ollama is impractical. - Positive: Evaluation harness can compare Ollama, OpenAI, and Claude quality side by side. - Positive: Minimal abstraction layer (Protocol + three implementations) is easy to understand and test. -- Positive: Ollama is opt-in via Docker Compose profiles; remote-backend setups skip the ~6 GB model download. -- Negative: Ollama with Mistral 7B requires approximately 6 GB of RAM, which is significant in memory-constrained environments (TinyLlama at 637 MB is a documented fallback). +- Positive: Ollama is opt-in via Docker Compose profiles; remote-backend setups skip the model download entirely. +- Positive: The same profile switch allows a host-native Ollama to be used instead, which is the only way to get GPU-accelerated local inference on macOS and Windows. +- Negative: Local model size and memory use scale with the model chosen via `OLLAMA_MODEL`; a smaller model is the documented fallback on constrained machines. - Negative: Three code paths means all three backends need test coverage. --- diff --git a/docs/evaluation-results.ja.md b/docs/evaluation-results.ja.md index e8100ee..67d9c7a 100644 --- a/docs/evaluation-results.ja.md +++ b/docs/evaluation-results.ja.md @@ -9,7 +9,7 @@ ```bash # 1. スタックを起動する — どちらか一方を選択: -make run # Ollama を含む（約 6 GB の RAM とモデルのダウンロードが必要） +make run # コンテナ化された Ollama を含む（初回起動時に OLLAMA_MODEL をダウンロード） make run-remote # Ollama なし; 事前に LLM_BACKEND=claude または openai を設定すること # 2. ragas ジャッジ用の認証情報をエクスポートする（いずれか一つ; 優先順位: Anthropic > OpenAI > Ollama） diff --git a/docs/evaluation-results.md b/docs/evaluation-results.md index bff7854..34bcc45 100644 --- a/docs/evaluation-results.md +++ b/docs/evaluation-results.md @@ -9,7 +9,7 @@ Each run appends a new section below. ```bash # 1. Start the stack — choose one: -make run # includes Ollama (needs ~6 GB RAM + model download) +make run # includes containerised Ollama (downloads OLLAMA_MODEL on first start) make run-remote # without Ollama; set LLM_BACKEND=claude or openai first # 2. Export credentials for the ragas judge (pick one; priority: Anthropic > OpenAI > Ollama) diff --git a/docs/technical-design-document.ja.md b/docs/technical-design-document.ja.md index f9ca89f..d29a7f1 100644 --- a/docs/technical-design-document.ja.md +++ b/docs/technical-design-document.ja.md @@ -198,11 +198,11 @@ Answer: | ID | 要件 | | ----- | ------------------------------------------------------------------------------------------------------------------------------------------------- | | NFR-1 | ベクトル類似性検索（LLM 生成を除く検索のみ）は、最大 500,000 チャンクのコーパスに対して p99 で 200ms 以内に完了しなければならない（SHALL）。 | -| NFR-2 | エンドツーエンドのクエリレイテンシ（検索 + LLM 生成）は、少なくとも 16 GB RAM を搭載したマシンで Ollama の `mistral:7b` を使用する場合、p95 で 10 秒以内に完了しなければならない（SHALL）。 | +| NFR-2 | エンドツーエンドのクエリレイテンシ（検索 + LLM 生成）は、LLM バックエンドが GPU アクセラレーション対応またはホスト型である場合、p95 で 10 秒以内に完了しなければならない（SHALL）。CPU のみのローカル推論はこの目標の対象外である。そのレイテンシは本システムではなく、選択したモデルとホストのハードウェアによって決まる。 | | NFR-3 | 埋め込みワーカーは、単一のワーカーインスタンスで少なくとも毎秒 50 チャンクを処理して保存しなければならない（SHALL）。 | | NFR-4 | 取り込みサービスは、SEC の 1 秒あたり 10 リクエストのレート制限を遵守しながら、SEC EDGAR から少なくとも毎分 10 件の書類を取得して公開しなければならない（SHALL）。 | -> **注記:** NFR-1 から NFR-4 は、想定されるデプロイメントプロファイル（RAM 16 GB 以上、CPU のみの埋め込み）に合わせたデザインターゲットである。ベンチマーク測定はまだ行われていない。実際の測定値に対する検証は、最初の完全な評価実行（評価ハードウェアが利用可能になる 2026 年 4 月頃に予定）で計画されている。実際のデータが収集された後、数値が修正される可能性がある。 +> **注記:** NFR-1 から NFR-4 は測定結果ではなくデザインターゲットであり、意図的に特定のハードウェアプロファイルを固定していない。CPU、メモリ、バッチサイズはデプロイごとに設定可能である（セクション 10.2 を参照）。実際の測定値に対する検証は、最初の完全な評価実行で計画されている。実際のデータが収集された後、数値が修正される可能性がある。 ### 4.2 キャパシティ diff --git a/docs/technical-design-document.md b/docs/technical-design-document.md index 94799fe..2b75caa 100644 --- a/docs/technical-design-document.md +++ b/docs/technical-design-document.md @@ -198,11 +198,11 @@ Answer: | ID | Requirement | | ----- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | NFR-1 | Vector similarity search (retrieval only, excluding LLM generation) SHALL complete within 200ms at p99 for a corpus of up to 500,000 chunks. | -| NFR-2 | End-to-end query latency (retrieval + LLM generation) SHALL complete within 10 seconds at p95 when using Ollama with `mistral:7b` on a machine with at least 16 GB RAM. | +| NFR-2 | End-to-end query latency (retrieval + LLM generation) SHALL complete within 10 seconds at p95 when the LLM backend is GPU-accelerated or hosted. CPU-only local inference is out of scope for this target: its latency is set by the chosen model and host hardware, not by this system. | | NFR-3 | The embedding worker SHALL process and store at least 50 chunks per second on a single worker instance. | | NFR-4 | The ingestion service SHALL fetch and publish at least 10 filings per minute from SEC EDGAR, respecting the SEC's rate limit of 10 requests per second. | -> **Note:** NFR-1 through NFR-4 are design targets sized for the intended deployment profile (≥16 GB RAM, CPU-only embedding). No benchmark measurements have been taken yet. Validation against actual measurements is planned for the first full evaluation run (expected April 2026 once evaluation hardware is available). Numbers may be revised once real data is collected. +> **Note:** NFR-1 through NFR-4 are design targets, not measured results, and they deliberately do not pin a hardware profile — CPU, memory, and batch sizes are configurable per deployment (see Section 10.2). Validation against actual measurements is planned for the first full evaluation run. Numbers may be revised once real data is collected. ### 4.2 Capacity diff --git a/examples/README.ja.md b/examples/README.ja.md index bac21e2..df9b4e2 100644 --- a/examples/README.ja.md +++ b/examples/README.ja.md @@ -33,7 +33,7 @@ curl -X POST http://localhost:8001/v1/ingest \ } ``` -埋め込みワーカーが Kafka メッセージを受け取り、各ファイリングを約 512 トークンのウィンドウにチャンク分割し、`sentence-transformers/all-MiniLM-L6-v2` で埋め込みを生成して、ベクトルを pgvector に書き込みます。CPU 上では数分かかります。 +埋め込みワーカーが Kafka メッセージを受け取り、各ファイリングを約 512 トークンのウィンドウにチャンク分割し、`nomic-ai/nomic-embed-text-v1.5` で埋め込みを生成して、ベクトルを pgvector に書き込みます。CPU 上では数分かかります。 --- diff --git a/examples/README.md b/examples/README.md index 6f80145..c434478 100644 --- a/examples/README.md +++ b/examples/README.md @@ -33,7 +33,7 @@ curl -X POST http://localhost:8001/v1/ingest \ } ``` -The embedding worker picks up the Kafka messages, chunks each filing into ~512-token windows, embeds them with `sentence-transformers/all-MiniLM-L6-v2`, and writes the vectors to pgvector. This takes a few minutes on CPU. +The embedding worker picks up the Kafka messages, chunks each filing into ~512-token windows, embeds them with `nomic-ai/nomic-embed-text-v1.5`, and writes the vectors to pgvector. This takes a few minutes on CPU. --- diff --git a/helm/findoc-rag/templates/configmaps.yaml b/helm/findoc-rag/templates/configmaps.yaml index 6aa3a25..8acbf36 100644 --- a/helm/findoc-rag/templates/configmaps.yaml +++ b/helm/findoc-rag/templates/configmaps.yaml @@ -14,6 +14,7 @@ data: LLM_BACKEND: {{ .Values.queryApi.llmBackend | quote }} OLLAMA_URL: http://{{ .Release.Name }}-ollama:11434 OLLAMA_MODEL: {{ .Values.ollama.model | quote }} + OLLAMA_NUM_CTX: {{ .Values.queryApi.ollamaNumCtx | quote }} EDGAR_USER_AGENT: {{ .Values.ingestion.edgarUserAgent | quote }} EDGAR_RATE_LIMIT_RPS: {{ .Values.ingestion.edgarRateLimitRps | quote }} LOG_LEVEL: {{ .Values.ingestion.logLevel | quote }} diff --git a/helm/findoc-rag/tests/configmaps_test.yaml b/helm/findoc-rag/tests/configmaps_test.yaml index d166554..7db7504 100644 --- a/helm/findoc-rag/tests/configmaps_test.yaml +++ b/helm/findoc-rag/tests/configmaps_test.yaml @@ -100,6 +100,24 @@ tests: path: data.OLLAMA_MODEL value: mistral:7b + # Without an explicit num_ctx, Ollama falls back to its own small default + # and silently truncates the RAG prompt, dropping the system instructions. + - it: app-config configmap sets an explicit OLLAMA_NUM_CTX + documentIndex: 0 + asserts: + - equal: + path: data.OLLAMA_NUM_CTX + value: "8192" + + - it: OLLAMA_NUM_CTX is overridable from values + documentIndex: 0 + set: + queryApi.ollamaNumCtx: 32768 + asserts: + - equal: + path: data.OLLAMA_NUM_CTX + value: "32768" + # ── db-migrations configmap (documentIndex 1) ───────────────────── - it: db-migrations configmap has correct kind diff --git a/helm/findoc-rag/values.yaml b/helm/findoc-rag/values.yaml index 0380580..442a2e5 100644 --- a/helm/findoc-rag/values.yaml +++ b/helm/findoc-rag/values.yaml @@ -96,6 +96,9 @@ queryApi: port: 8000 llmBackend: ollama ollamaModel: mistral:7b + # Context window requested from Ollama. The full RAG prompt must fit or + # Ollama truncates it silently, dropping the system prompt first. + ollamaNumCtx: 8192 embeddingModel: "nomic-ai/nomic-embed-text-v1.5" apiKeys: "dev-key-1,dev-key-2" openaiApiKey: "" diff --git a/services/embedding-worker/src/chunker.py b/services/embedding-worker/src/chunker.py index 138e63d..bbef69d 100644 --- a/services/embedding-worker/src/chunker.py +++ b/services/embedding-worker/src/chunker.py @@ -1,7 +1,10 @@ """Section-aware chunking for 10-K filings. Implements the chunking strategy from TDD Section 5.2.2: - 1. Section split — by 10-K item headers (Item 1, 1A, 7, etc.) + 1. Section split — by 10-K item headers (Item 1, 1A, 7, etc.), keeping + only the occurrence of each item that holds its actual body text + (every 10-K lists all of its items twice: once in the table of + contents, once as the real section header) 2. Paragraph split — by double newlines within each section 3. Paragraph packing — consecutive paragraphs are greedily packed into chunks of up to 512 tokens, so short paragraphs (headings, single @@ -38,13 +41,6 @@ re.MULTILINE, ) -# Canonical ordering so we can sort matches deterministically -_ITEM_ORDER = [ - "1", "1A", "1B", "2", "3", "4", "5", "6", - "7", "7A", "8", "9", "9A", "9B", "10", "11", - "12", "13", "14", "15", -] - DEFAULT_CHUNK_SIZE = 512 # tokens DEFAULT_OVERLAP = 64 # tokens TIKTOKEN_ENCODING = "cl100k_base" @@ -109,10 +105,44 @@ def make_chunk_id(accession_number: str, section_name: str, chunk_index: int) -> # ── Stage 1: Section split ────────────────────────────────────── +def _select_section_boundaries( + text: str, + matches: list[re.Match[str]], +) -> list[re.Match[str]]: + """Choose which Item matches are real section headers. + + Every 10-K names each of its items at least twice: once as a row in the + table of contents near the front, and once as the header of the actual + section. Treating all of them as boundaries produces ~20 junk sections + holding nothing but TOC rows, and — worse — mislabels real body text with + whichever TOC row happened to precede it. + + For each item number we therefore keep only the occurrence that owns the + most text, which is the body header rather than the one-line TOC row. + Ties go to the later occurrence, since the body always follows the TOC. + Matches that lose are not boundaries, so their text is absorbed into the + preceding section instead of being dropped — the TOC ends up in the + Preamble, where it belongs. + """ + best: dict[str, tuple[re.Match[str], int]] = {} + + for i, match in enumerate(matches): + end = matches[i + 1].start() if i + 1 < len(matches) else len(text) + span = end - match.end() + item_number = match.group(1) + current = best.get(item_number) + # >= so that, at equal length, the later (body) occurrence wins. + if current is None or span >= current[1]: + best[item_number] = (match, span) + + return sorted((match for match, _ in best.values()), key=lambda m: m.start()) + + def split_into_sections(text: str) -> list[tuple[str, str]]: """Split a 10-K filing into (section_name, section_text) pairs. - Uses regex to find Item headers. Text before the first item is + Uses regex to find Item headers, keeping one boundary per item number + (see _select_section_boundaries). Text before the first item is labelled "Preamble". If no items are found, the entire document is returned as a single "Full Document" section. """ @@ -121,19 +151,21 @@ def split_into_sections(text: str) -> list[tuple[str, str]]: if not matches: return [("Full Document", text)] + boundaries = _select_section_boundaries(text, matches) + sections: list[tuple[str, str]] = [] - # Text before the first match → Preamble - if matches[0].start() > 0: - preamble = text[: matches[0].start()].strip() + # Text before the first boundary → Preamble (cover page + table of contents) + if boundaries[0].start() > 0: + preamble = text[: boundaries[0].start()].strip() if preamble: sections.append(("Preamble", preamble)) - for i, match in enumerate(matches): + for i, match in enumerate(boundaries): item_number = match.group(1) section_name = f"Item {item_number}" start = match.end() - end = matches[i + 1].start() if i + 1 < len(matches) else len(text) + end = boundaries[i + 1].start() if i + 1 < len(boundaries) else len(text) section_text = text[start:end].strip() if section_text: sections.append((section_name, section_text)) diff --git a/services/embedding-worker/src/drift_detector.py b/services/embedding-worker/src/drift_detector.py index 424ffe6..137f0d4 100644 --- a/services/embedding-worker/src/drift_detector.py +++ b/services/embedding-worker/src/drift_detector.py @@ -115,9 +115,19 @@ def main() -> None: corpus_mean = row[0] if row else None with conn.cursor() as cur: + # NOW() is already an absolute timestamptz, so it is compared + # directly against created_at (also timestamptz). Stripping the + # zone with AT TIME ZONE 'UTC' would yield a naive timestamp that + # Postgres re-interprets in the server's local zone, sliding the + # lookback window by that offset on any non-UTC server. + # + # make_interval() takes the window as a real parameter. The + # interval cannot be written as INTERVAL '%s days': the + # placeholder would sit inside a string literal, which psycopg2 + # documents as unsupported. cur.execute( "SELECT avg(embedding) FROM document_chunks " - "WHERE created_at >= (NOW() AT TIME ZONE 'UTC') - INTERVAL '%s days'", + "WHERE created_at >= NOW() - make_interval(days => %s)", (DRIFT_LOOKBACK_DAYS,), ) row = cur.fetchone() diff --git a/services/embedding-worker/tests/test_chunker.py b/services/embedding-worker/tests/test_chunker.py index 63d7309..3b28ccc 100644 --- a/services/embedding-worker/tests/test_chunker.py +++ b/services/embedding-worker/tests/test_chunker.py @@ -87,6 +87,96 @@ def test_case_insensitive_item(self) -> None: assert "Item 1" in names +# ── Stage 1: table-of-contents handling ───────────────────────── + +def _filing_with_toc() -> str: + """A miniature 10-K: cover page, table of contents, then the real body. + + Every real filing names each item twice — once as a TOC row with a page + number, once as the header of the actual section. + """ + return ( + "UNITED STATES SECURITIES AND EXCHANGE COMMISSION\n" + "Annual Report on Form 10-K\n" + "\n" + "TABLE OF CONTENTS\n" + "Item 1. Business 3\n" + "Item 1A. Risk Factors 12\n" + "Item 7. Management's Discussion and Analysis 40\n" + "\n" + "Item 1. Business\n" + "The Company designs and sells consumer electronics worldwide. " * 12 + "\n" + "Item 1A. Risk Factors\n" + "The Company depends on a concentrated supply chain in Asia. " * 12 + "\n" + "Item 7. Management's Discussion and Analysis\n" + "Total net sales increased 2% year over year to $391.0 billion. " * 12 + ) + + +class TestTableOfContentsHandling: + def test_each_item_yields_exactly_one_section(self) -> None: + """TOC rows must not create a second section for the same item.""" + sections = split_into_sections(_filing_with_toc()) + names = [name for name, _ in sections] + for item in ("Item 1", "Item 1A", "Item 7"): + assert names.count(item) == 1, f"{item} appeared {names.count(item)} times" + + def test_sections_hold_body_text_not_toc_rows(self) -> None: + """The kept boundary must be the body header, not the TOC row.""" + sections = dict(split_into_sections(_filing_with_toc())) + assert "concentrated supply chain" in sections["Item 1A"] + assert "Total net sales increased" in sections["Item 7"] + + def test_toc_rows_land_in_the_preamble(self) -> None: + """Losing boundaries are absorbed, never dropped — no text is lost.""" + sections = dict(split_into_sections(_filing_with_toc())) + assert "TABLE OF CONTENTS" in sections["Preamble"] + # The page-numbered TOC row stays with the TOC, out of the body sections. + assert "Risk Factors 12" in sections["Preamble"] + assert "Risk Factors 12" not in sections["Item 1A"] + + def test_body_text_is_not_mislabelled_by_a_toc_row(self) -> None: + """Regression: body prose must carry its own item's section name. + + With every TOC row treated as a boundary, the last TOC row swallowed + the text that followed it, so real Item 1 prose was labelled with + whichever item the TOC happened to list last. + """ + sections = dict(split_into_sections(_filing_with_toc())) + assert "designs and sells consumer electronics" in sections["Item 1"] + + def test_no_text_is_lost(self) -> None: + sections = split_into_sections(_filing_with_toc()) + recombined = "".join(text for _, text in sections) + assert "designs and sells consumer electronics" in recombined + assert "concentrated supply chain" in recombined + assert "Total net sales increased" in recombined + + def test_duplicate_item_keeps_the_longer_span(self) -> None: + """A passing mention loses to the section that actually holds the body.""" + text = ( + "Cover.\n" + "Item 3. Legal Proceedings 5\n" + "Item 3. Legal Proceedings\n" + + "The Company is party to various legal proceedings. " * 20 + ) + sections = dict(split_into_sections(text)) + assert "The Company is party" in sections["Item 3"] + # The TOC row (title plus page number) stayed behind in the Preamble. + assert "Legal Proceedings 5" not in sections["Item 3"] + assert "Legal Proceedings 5" in sections["Preamble"] + + def test_single_occurrence_is_unaffected(self) -> None: + """Documents without a TOC keep their original one-section-per-item split.""" + text = ( + "Cover page.\n" + "Item 1. Business\nWe sell things.\n" + "Item 7. MD&A\nRevenue grew." + ) + names = [name for name, _ in split_into_sections(text)] + assert names == ["Preamble", "Item 1", "Item 7"] + + # ── Stage 2: Paragraph split ──────────────────────────────────── class TestSplitIntoParagraphs: diff --git a/services/embedding-worker/tests/test_drift_detector.py b/services/embedding-worker/tests/test_drift_detector.py index 2d99e79..9a4cb52 100644 --- a/services/embedding-worker/tests/test_drift_detector.py +++ b/services/embedding-worker/tests/test_drift_detector.py @@ -117,6 +117,56 @@ def test_returns_early_when_no_recent_chunks( dd.main() +# ── main() — lookback window SQL ───────────────────────────────── + +class TestLookbackWindowQuery: + """The recent-window query must compare two timestamptz values. + + Subtracting the interval from a naive `NOW() AT TIME ZONE 'UTC'` makes + Postgres re-interpret the result in the server's local zone, shifting the + window by that offset whenever the server is not on UTC. + """ + + def _run_and_capture_sql(self) -> list[tuple[str, tuple[object, ...] | None]]: + executed: list[tuple[str, tuple[object, ...] | None]] = [] + results = iter([(np.array([1.0, 0.0]),), (np.array([1.0, 0.0]),)]) + + mock_cur = MagicMock() + mock_cur.__enter__ = MagicMock(return_value=mock_cur) + mock_cur.__exit__ = MagicMock(return_value=False) + mock_cur.fetchone.side_effect = lambda: next(results) + mock_cur.execute.side_effect = lambda sql, params=None: executed.append( + (sql, params) + ) + + mock_conn = MagicMock() + mock_conn.cursor.return_value = mock_cur + + with ( + patch("src.drift_detector.psycopg2.connect", return_value=mock_conn), + patch("src.drift_detector.register_vector"), + patch.object(dd, "PUSHGATEWAY_URL", None), + ): + dd.main() + return executed + + def test_recent_window_compares_against_timestamptz(self) -> None: + recent_sql = self._run_and_capture_sql()[1][0] + assert "NOW()" in recent_sql + assert "AT TIME ZONE" not in recent_sql + + def test_lookback_days_is_a_bound_parameter(self) -> None: + """The window must not be interpolated inside a quoted literal. + + `INTERVAL '%s days'` puts the placeholder inside a string, which + psycopg2 does not support. + """ + recent_sql, params = self._run_and_capture_sql()[1] + assert "make_interval" in recent_sql + assert "INTERVAL '" not in recent_sql + assert params == (dd.DRIFT_LOOKBACK_DAYS,) + + # ── main() — normal run (no alert) ─────────────────────────────── class TestMainNormalRun: diff --git a/services/ingestion/src/kafka_producer.py b/services/ingestion/src/kafka_producer.py index 1d2512b..5841a6a 100644 --- a/services/ingestion/src/kafka_producer.py +++ b/services/ingestion/src/kafka_producer.py @@ -11,10 +11,11 @@ import json from datetime import UTC, datetime +from enum import StrEnum from typing import TYPE_CHECKING, Any import structlog -from confluent_kafka import KafkaError, Producer +from confluent_kafka import KafkaError, KafkaException, Producer from src.config import settings from src.metrics import FILING_SIZE_BYTES, KAFKA_PUBLISH_TOTAL @@ -28,6 +29,18 @@ _KAFKA_MAX_MESSAGE_BYTES = 157_286_400 # 150 MB — must match broker KAFKA_MESSAGE_MAX_BYTES +# How long to wait for the broker to acknowledge one filing before treating it +# as undelivered. Generous because acks=all on a multi-MB message is slow. +_DELIVERY_TIMEOUT_SECONDS = 120.0 + + +class PublishOutcome(StrEnum): + """Result of attempting to publish one filing.""" + + DELIVERED = "delivered" # broker acknowledged the message + TOO_LARGE = "too_large" # payload over MAX_RAW_BYTES, never produced + FAILED = "failed" # produce errored, was rejected, or timed out + def _delivery_callback(err: KafkaError | None, msg: Any) -> None: """Callback invoked on message delivery (or failure).""" @@ -73,11 +86,20 @@ def __init__(self, bootstrap_servers: str | None = None) -> None: } ) - def publish_filing(self, filing: Filing) -> bool: - """Serialize a Filing and publish it to Kafka (FR-3). - - Returns True if the message was produced, False if it was skipped - because the raw payload exceeded MAX_RAW_BYTES. + def publish_filing( + self, + filing: Filing, + *, + timeout: float = _DELIVERY_TIMEOUT_SECONDS, + ) -> PublishOutcome: + """Serialize a Filing and publish it to Kafka, waiting for the ack (FR-3). + + ``produce()`` only enqueues into librdkafka's buffer; delivery happens + later and can still fail. The caller records the filing in + ``ingestion_log``, which permanently suppresses re-ingestion of that + accession number, so returning before the broker has acknowledged the + message risks losing a filing with no way to notice. This method + therefore blocks until delivery is confirmed and reports the outcome. """ message = { "accession_number": filing.accession_number, @@ -103,16 +125,57 @@ def publish_filing(self, filing: Filing) -> bool: raw_mb=round(raw_bytes / 1024 / 1024, 1), limit_mb=round(self.MAX_RAW_BYTES / 1024 / 1024, 1), ) - return False + return PublishOutcome.TOO_LARGE + + # Captured by the per-message callback below; non-empty means the + # broker rejected this specific message. + delivery_errors: list[str] = [] + + def _on_delivery(err: KafkaError | None, msg: Any) -> None: + _delivery_callback(err, msg) + if err is not None: + delivery_errors.append(str(err)) + + try: + self._producer.produce( + topic=TOPIC_FILINGS_RAW, + key=filing.accession_number, + value=payload, + callback=_on_delivery, + ) + except (BufferError, KafkaException) as exc: + KAFKA_PUBLISH_TOTAL.labels(topic=TOPIC_FILINGS_RAW, status="error").inc() + logger.error( + "kafka_produce_failed", + ticker=filing.ticker, + accession=filing.accession_number, + error=str(exc), + ) + return PublishOutcome.FAILED + + # Block until this message is acknowledged. flush() returns the number + # of messages still queued, so a non-zero result means we timed out. + remaining = self._producer.flush(timeout) + if remaining: + logger.error( + "kafka_delivery_timeout", + ticker=filing.ticker, + accession=filing.accession_number, + timeout_seconds=timeout, + still_queued=remaining, + ) + return PublishOutcome.FAILED - self._producer.produce( - topic=TOPIC_FILINGS_RAW, - key=filing.accession_number, - value=payload, - callback=_delivery_callback, - ) - self._producer.poll(0) # Trigger any pending delivery callbacks - return True + if delivery_errors: + logger.error( + "kafka_delivery_rejected", + ticker=filing.ticker, + accession=filing.accession_number, + error=delivery_errors[0], + ) + return PublishOutcome.FAILED + + return PublishOutcome.DELIVERED def flush(self, timeout: float = 10.0) -> int: """Wait for all in-flight messages to be delivered. diff --git a/services/ingestion/src/main.py b/services/ingestion/src/main.py index b27ecae..dfac03f 100644 --- a/services/ingestion/src/main.py +++ b/services/ingestion/src/main.py @@ -35,7 +35,7 @@ from src.db import IngestionDB from src.edgar_client import EdgarClient from src.facts import extract_annual_facts -from src.kafka_producer import FilingProducer +from src.kafka_producer import FilingProducer, PublishOutcome from src.metrics import FILINGS_FETCHED_TOTAL # ── Structured logging (TDD: Section 8.3) ─────────────────────── @@ -90,6 +90,9 @@ class IngestResponse(BaseModel): tickers_processed: list[str] filings_published: int filings_skipped: int + # Filings Kafka never acknowledged. Deliberately not recorded as ingested, + # so a subsequent /v1/ingest call retries them. + filings_failed: int facts_stored: int errors: list[str] @@ -193,7 +196,16 @@ def _publish_single_filing( db: Any, producer: Any, ) -> str: - """Attempt to publish one filing. Returns 'published', 'skipped', or raises.""" + """Attempt to publish one filing. + + Returns 'published', 'skipped' (already ingested, or too large to publish), + or 'failed' (Kafka did not acknowledge the message). + + ``ingestion_log`` is only written after the broker confirms delivery. + Writing it earlier would permanently mark the accession as ingested — + ``is_already_ingested`` would skip it on every future run — even though the + filing never reached the topic and no chunks will ever be produced. + """ if db.is_already_ingested(filing.accession_number): logger.info( "filing_skipped_duplicate", @@ -209,8 +221,22 @@ def _publish_single_filing( filing_date=filing.filing_date, ) t_publish = time.perf_counter() - if not producer.publish_filing(filing): + outcome = producer.publish_filing(filing) + + if outcome is PublishOutcome.TOO_LARGE: return "skipped" + + if outcome is not PublishOutcome.DELIVERED: + # Left unrecorded on purpose: the next ingest run retries this filing. + logger.error( + "filing_publish_failed", + ticker=symbol, + accession=filing.accession_number, + outcome=str(outcome), + elapsed_ms=round((time.perf_counter() - t_publish) * 1000, 1), + ) + return "failed" + db.record_ingestion(filing) logger.info( "filing_published", @@ -263,6 +289,7 @@ async def ingest(body: IngestRequest | None = None) -> IngestResponse: filings_published = 0 filings_skipped = 0 + filings_failed = 0 facts_stored = 0 errors: list[str] = [] tickers_processed: list[str] = [] @@ -289,6 +316,13 @@ async def ingest(body: IngestRequest | None = None) -> IngestResponse: if result == "published": filings_published += 1 FILINGS_FETCHED_TOTAL.labels(ticker=symbol, status="success").inc() + elif result == "failed": + filings_failed += 1 + errors.append( + f"Kafka delivery failed for {symbol} " + f"{filing.accession_number}; will retry on next run" + ) + FILINGS_FETCHED_TOTAL.labels(ticker=symbol, status="error").inc() else: filings_skipped += 1 FILINGS_FETCHED_TOTAL.labels(ticker=symbol, status="skipped").inc() @@ -304,6 +338,7 @@ async def ingest(body: IngestRequest | None = None) -> IngestResponse: ticker=symbol, filings_published=filings_published, filings_skipped=filings_skipped, + filings_failed=filings_failed, elapsed_ms=round((time.perf_counter() - t_ticker) * 1000, 1), ) @@ -325,6 +360,7 @@ async def ingest(body: IngestRequest | None = None) -> IngestResponse: tickers=tickers_processed, filings_published=filings_published, filings_skipped=filings_skipped, + filings_failed=filings_failed, facts_stored=facts_stored, errors=len(errors), elapsed_ms=round((time.perf_counter() - t_request) * 1000, 1), @@ -335,6 +371,7 @@ async def ingest(body: IngestRequest | None = None) -> IngestResponse: tickers_processed=tickers_processed, filings_published=filings_published, filings_skipped=filings_skipped, + filings_failed=filings_failed, facts_stored=facts_stored, errors=errors, ) diff --git a/services/ingestion/tests/test_edgar_client.py b/services/ingestion/tests/test_edgar_client.py index 7f714d3..1bc973d 100644 --- a/services/ingestion/tests/test_edgar_client.py +++ b/services/ingestion/tests/test_edgar_client.py @@ -17,6 +17,7 @@ from src.config import Settings, load_tickers from src.edgar_client import EdgarClient, Filing +from src.kafka_producer import PublishOutcome if TYPE_CHECKING: from collections.abc import AsyncGenerator @@ -609,6 +610,49 @@ def test_load_tickers_empty_file(self, tmp_path: Any) -> None: # ── Kafka Producer ─────────────────────────────────────────────── +def _delivery_test_filing() -> Filing: + return Filing( + accession_number="0001-24-000001", + ticker="AAPL", + company_name="Apple Inc.", + filing_date="2024-11-01", + filing_type="10-K", + source_url="https://sec.gov/...", + raw_text="Item 1. Business...", + ) + + +def _make_acking_producer_mock(delivery_error: str | None = None) -> MagicMock: + """Mock Producer whose flush() fires the pending delivery callback. + + Mirrors librdkafka: produce() only enqueues, and the callback registered + with it runs later — during flush() — carrying the delivery result. + """ + instance = MagicMock() + pending: list[Any] = [] + + def _produce(**kwargs: Any) -> None: + pending.append(kwargs["callback"]) + + def _flush(_timeout: float = 0.0) -> int: + msg = MagicMock() + msg.topic.return_value = "filings.raw" + msg.partition.return_value = 0 + msg.offset.return_value = 1 + err = None + if delivery_error is not None: + err = MagicMock() + err.__str__ = lambda _self: delivery_error # type: ignore[assignment] + for callback in pending: + callback(err, msg) + pending.clear() + return 0 # nothing left queued + + instance.produce.side_effect = _produce + instance.flush.side_effect = _flush + return instance + + class TestKafkaProducer: """Tests for kafka_producer.py (serialisation logic).""" @@ -647,6 +691,78 @@ def test_publish_filing_calls_produce(self) -> None: assert msg["source_url"] == "https://sec.gov/..." assert msg["accession_number"] == "0001-24-000001" + def test_publish_filing_returns_delivered_once_broker_acks(self) -> None: + """A clean flush with no callback error means the broker took the message.""" + from src.kafka_producer import FilingProducer + + with patch("src.kafka_producer.Producer") as mock_producer: + mock_instance = _make_acking_producer_mock() + mock_producer.return_value = mock_instance + + producer = FilingProducer(bootstrap_servers="localhost:9092") + outcome = producer.publish_filing(_delivery_test_filing()) + + assert outcome is PublishOutcome.DELIVERED + # The ack must be waited for, not left in the background buffer. + mock_instance.flush.assert_called_once() + + def test_publish_filing_returns_failed_when_broker_rejects(self) -> None: + """A delivery callback carrying an error must not be reported as success.""" + from src.kafka_producer import FilingProducer + + with patch("src.kafka_producer.Producer") as mock_producer: + mock_instance = _make_acking_producer_mock(delivery_error="Broker unavailable") + mock_producer.return_value = mock_instance + + producer = FilingProducer(bootstrap_servers="localhost:9092") + outcome = producer.publish_filing(_delivery_test_filing()) + + assert outcome is PublishOutcome.FAILED + + def test_publish_filing_returns_failed_when_flush_times_out(self) -> None: + """Messages still queued after flush() are undelivered, not delivered.""" + from src.kafka_producer import FilingProducer + + with patch("src.kafka_producer.Producer") as mock_producer: + mock_instance = MagicMock() + mock_instance.flush.return_value = 1 # one message still queued + mock_producer.return_value = mock_instance + + producer = FilingProducer(bootstrap_servers="localhost:9092") + outcome = producer.publish_filing(_delivery_test_filing(), timeout=0.1) + + assert outcome is PublishOutcome.FAILED + + def test_publish_filing_returns_failed_when_produce_raises(self) -> None: + """A full local queue surfaces as a failure rather than an exception.""" + from src.kafka_producer import FilingProducer + + with patch("src.kafka_producer.Producer") as mock_producer: + mock_instance = MagicMock() + mock_instance.produce.side_effect = BufferError("Local queue full") + mock_producer.return_value = mock_instance + + producer = FilingProducer(bootstrap_servers="localhost:9092") + outcome = producer.publish_filing(_delivery_test_filing()) + + assert outcome is PublishOutcome.FAILED + + def test_publish_filing_returns_too_large_without_producing(self) -> None: + """Oversize filings are rejected before hitting the broker.""" + from src.kafka_producer import FilingProducer + + with patch("src.kafka_producer.Producer") as mock_producer: + mock_instance = _make_acking_producer_mock() + mock_producer.return_value = mock_instance + + producer = FilingProducer(bootstrap_servers="localhost:9092") + producer.MAX_RAW_BYTES = 10 # smaller than any serialised filing + + outcome = producer.publish_filing(_delivery_test_filing()) + + assert outcome is PublishOutcome.TOO_LARGE + mock_instance.produce.assert_not_called() + def test_flush_delegates_to_producer(self) -> None: from src.kafka_producer import FilingProducer @@ -849,7 +965,7 @@ def test_ingest_success_with_tickers(self) -> None: mock_kafka = MagicMock() mock_db = MagicMock() - mock_kafka.publish_filing = MagicMock() + mock_kafka.publish_filing = MagicMock(return_value=PublishOutcome.DELIVERED) mock_kafka.flush = MagicMock() mock_db.is_already_ingested.return_value = False @@ -869,6 +985,7 @@ def test_ingest_success_with_tickers(self) -> None: assert data["tickers_processed"] == ["AAPL"] assert data["filings_published"] == 1 assert data["filings_skipped"] == 0 + assert data["filings_failed"] == 0 assert data["facts_stored"] == 0 assert data["errors"] == [] mock_kafka.publish_filing.assert_called_once_with(filing) @@ -1048,6 +1165,95 @@ def test_ingest_skips_already_ingested_filing(self) -> None: main_module._kafka_producer = original_kafka main_module._db = original_db + def test_ingest_does_not_record_filing_kafka_never_acked(self) -> None: + """A filing Kafka never acknowledged must stay eligible for re-ingestion. + + Recording it in ingestion_log would make is_already_ingested() skip it + on every future run, losing the filing permanently even though no + chunks were ever produced for it. + """ + from fastapi.testclient import TestClient + + import src.main as main_module + + filing = _delivery_test_filing() + + mock_edgar = _make_edgar_mock(filings=[filing]) + mock_kafka = MagicMock() + mock_kafka.publish_filing = MagicMock(return_value=PublishOutcome.FAILED) + mock_kafka.flush = MagicMock() + mock_db = MagicMock() + mock_db.is_already_ingested.return_value = False + + original_edgar = main_module._edgar_client + original_kafka = main_module._kafka_producer + original_db = main_module._db + main_module._edgar_client = mock_edgar + main_module._kafka_producer = mock_kafka + main_module._db = mock_db + + try: + client = TestClient(app=main_module.app, raise_server_exceptions=False) + response = client.post("/v1/ingest", json={"tickers": ["AAPL"]}) + assert response.status_code == 200 + data = response.json() + + # The filing counts as failed, never as published or merely skipped. + assert data["filings_failed"] == 1 + assert data["filings_published"] == 0 + assert data["filings_skipped"] == 0 + # The critical assertion: nothing was written to ingestion_log. + mock_db.record_ingestion.assert_not_called() + # And the failure is visible to the caller, not swallowed. + assert len(data["errors"]) == 1 + assert filing.accession_number in data["errors"][0] + finally: + main_module._edgar_client = original_edgar + main_module._kafka_producer = original_kafka + main_module._db = original_db + + def test_ingest_records_filing_only_after_delivery_ack(self) -> None: + """ingestion_log is written only once the broker has confirmed delivery.""" + from fastapi.testclient import TestClient + + import src.main as main_module + + filing = _delivery_test_filing() + call_order: list[str] = [] + + mock_edgar = _make_edgar_mock(filings=[filing]) + mock_kafka = MagicMock() + mock_db = MagicMock() + mock_db.is_already_ingested.return_value = False + + def _publish(_filing: Filing) -> PublishOutcome: + call_order.append("publish") + return PublishOutcome.DELIVERED + + mock_kafka.publish_filing = MagicMock(side_effect=_publish) + mock_kafka.flush = MagicMock() + mock_db.record_ingestion = MagicMock( + side_effect=lambda _f: call_order.append("record") + ) + + original_edgar = main_module._edgar_client + original_kafka = main_module._kafka_producer + original_db = main_module._db + main_module._edgar_client = mock_edgar + main_module._kafka_producer = mock_kafka + main_module._db = mock_db + + try: + client = TestClient(app=main_module.app, raise_server_exceptions=False) + response = client.post("/v1/ingest", json={"tickers": ["AAPL"]}) + assert response.status_code == 200 + assert response.json()["filings_published"] == 1 + assert call_order == ["publish", "record"] + finally: + main_module._edgar_client = original_edgar + main_module._kafka_producer = original_kafka + main_module._db = original_db + def test_ingest_mixed_new_and_duplicate_filings(self) -> None: """POST /v1/ingest publishes new filings and skips duplicates in the same run.""" from fastapi.testclient import TestClient @@ -1075,6 +1281,7 @@ def test_ingest_mixed_new_and_duplicate_filings(self) -> None: mock_edgar = _make_edgar_mock(filings=[new_filing, dup_filing]) mock_kafka = MagicMock() + mock_kafka.publish_filing = MagicMock(return_value=PublishOutcome.DELIVERED) mock_kafka.flush = MagicMock() mock_db = MagicMock() # new_filing is new, dup_filing is already ingested diff --git a/services/query-api/src/llm/ollama_backend.py b/services/query-api/src/llm/ollama_backend.py index 381491d..5af7e18 100644 --- a/services/query-api/src/llm/ollama_backend.py +++ b/services/query-api/src/llm/ollama_backend.py @@ -17,6 +17,19 @@ DEFAULT_OLLAMA_URL = "http://ollama:11434" DEFAULT_MODEL = "mistral:7b" +# Ollama defaults num_ctx to 4096 and silently discards whatever does not fit — +# from the *start* of the prompt, which is exactly where the system prompt and +# the citation instruction live. A top_k=5 RAG prompt is already ~3 K tokens +# (5 × 512-token chunks plus headers), and injected XBRL facts push it further, +# so the default is not enough headroom. 8192 covers top_k up to ~10 while +# keeping the KV cache small enough for the 8 GB local-LLM budget in the README; +# raise it (and the RAM allowance) if you routinely query with a larger top_k. +DEFAULT_NUM_CTX = 8192 + +# Rough chars-per-token ratio for English prose, used only to warn when a +# prompt is about to overflow the context window. +_CHARS_PER_TOKEN = 4 + class OllamaBackend: """Calls the local Ollama server's /api/generate endpoint.""" @@ -26,21 +39,47 @@ def __init__( base_url: str = DEFAULT_OLLAMA_URL, model: str = DEFAULT_MODEL, timeout_seconds: float = 300.0, + num_ctx: int = DEFAULT_NUM_CTX, ) -> None: self._base_url = base_url.rstrip("/") self._model = model self._timeout = aiohttp.ClientTimeout(total=timeout_seconds) + self._num_ctx = num_ctx @property def model_name(self) -> str: return self._model + @property + def num_ctx(self) -> int: + return self._num_ctx + + def _warn_if_context_exceeded(self, prompt: str, max_tokens: int) -> None: + """Log when prompt + answer budget cannot fit in the context window. + + Ollama truncates silently, so without this the only symptom is an + answer that ignores its instructions or cites nothing. + """ + estimated = len(prompt) // _CHARS_PER_TOKEN + if estimated + max_tokens <= self._num_ctx: + return + logger.warning( + "ollama_prompt_may_exceed_context", + estimated_prompt_tokens=estimated, + max_tokens=max_tokens, + num_ctx=self._num_ctx, + advice="Raise OLLAMA_NUM_CTX or lower top_k; Ollama truncates the " + "start of the prompt, dropping the system instructions.", + ) + async def generate(self, prompt: str, max_tokens: int = 1024) -> LLMResponse: """Call Ollama /api/generate and return the response. Raises aiohttp.ClientError or asyncio.TimeoutError on failure (caller handles graceful degradation). """ + self._warn_if_context_exceeded(prompt, max_tokens) + payload = { "model": self._model, "prompt": prompt, @@ -52,6 +91,9 @@ async def generate(self, prompt: str, max_tokens: int = 1024) -> LLMResponse: "think": False, "options": { "num_predict": max_tokens, + # Without an explicit num_ctx, Ollama falls back to its own + # small default and truncates the RAG context (see above). + "num_ctx": self._num_ctx, }, } diff --git a/services/query-api/src/main.py b/services/query-api/src/main.py index 57dbc05..2acd5f0 100644 --- a/services/query-api/src/main.py +++ b/services/query-api/src/main.py @@ -14,6 +14,7 @@ - TDD: Section 9.3 (rate limiting) """ +import asyncio import logging import os import uuid @@ -34,7 +35,7 @@ from src.auth import verify_api_key from src.llm.anthropic_backend import AnthropicBackend -from src.llm.ollama_backend import OllamaBackend +from src.llm.ollama_backend import DEFAULT_NUM_CTX, OllamaBackend from src.llm.openai_backend import OpenAIBackend from src.metrics import DEGRADED_RESPONSES_TOTAL, REQUESTS_TOTAL from src.models import ( @@ -59,6 +60,9 @@ LLM_BACKEND = os.getenv("LLM_BACKEND", "ollama") OLLAMA_URL = os.getenv("OLLAMA_URL", "http://ollama:11434") OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "mistral:7b") +# Context window requested from Ollama. Must be large enough to hold the whole +# RAG prompt or Ollama truncates it silently (see src/llm/ollama_backend.py). +OLLAMA_NUM_CTX = int(os.getenv("OLLAMA_NUM_CTX", str(DEFAULT_NUM_CTX))) OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o-mini") CLAUDE_MODEL = os.getenv("CLAUDE_MODEL", "claude-opus-4-6") EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "nomic-ai/nomic-embed-text-v1.5") @@ -146,7 +150,9 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: elif LLM_BACKEND == "claude": llm = AnthropicBackend(model=CLAUDE_MODEL) # type: ignore[assignment] else: - llm = OllamaBackend(base_url=OLLAMA_URL, model=OLLAMA_MODEL) # type: ignore[assignment] + llm = OllamaBackend( # type: ignore[assignment] + base_url=OLLAMA_URL, model=OLLAMA_MODEL, num_ctx=OLLAMA_NUM_CTX + ) _generator = RAGGenerator(retriever=_retriever, llm=llm, facts=_facts_repo) @@ -212,13 +218,16 @@ async def _rate_limit_handler(request: Request, exc: RateLimitExceeded) -> Respo @app.get("/health") async def health() -> dict[str, str]: - """Liveness probe (FR-21) — healthy when DB is reachable.""" + """Liveness probe (FR-21) — healthy when DB is reachable. + + The probe query is synchronous psycopg2, so it runs in a worker thread: + a slow or hung database must not stall the event loop for every other + in-flight request. + """ if _retriever is None: raise HTTPException(status_code=503, detail="unhealthy: not initialized") try: - cur = _retriever._get_conn().cursor() - cur.execute("SELECT 1") - cur.close() + await asyncio.to_thread(_retriever.ping) except Exception as exc: raise HTTPException(status_code=503, detail="unhealthy: db_unreachable") from exc return {"status": "healthy"} @@ -294,7 +303,9 @@ async def list_documents( REQUESTS_TOTAL.labels(endpoint="/v1/documents", status="success").inc() - rows, total = _retriever.list_documents(ticker=ticker, limit=limit, offset=offset) + rows, total = await asyncio.to_thread( + _retriever.list_documents, ticker=ticker, limit=limit, offset=offset + ) documents = [ DocumentInfo( diff --git a/services/query-api/src/rag/generator.py b/services/query-api/src/rag/generator.py index a4c0864..85cb329 100644 --- a/services/query-api/src/rag/generator.py +++ b/services/query-api/src/rag/generator.py @@ -8,6 +8,7 @@ from __future__ import annotations +import asyncio import time from typing import TYPE_CHECKING @@ -33,7 +34,14 @@ class RAGGenerator: - """Orchestrates the full RAG pipeline: retrieve → facts → prompt → generate.""" + """Orchestrates the full RAG pipeline: retrieve → facts → prompt → generate. + + Retrieval and fact lookup are synchronous, CPU- and IO-bound work (a + sentence-transformers forward pass plus psycopg2 queries). Calling them + directly from this coroutine would block the event loop for the whole + duration, serialising every concurrent request behind one query, so they + are dispatched to the default thread pool with ``asyncio.to_thread``. + """ def __init__( self, @@ -81,7 +89,8 @@ async def answer( # embedding_ms — query vector generation (returned by retriever) # retrieval_ms — pgvector SQL query (total minus embedding) t0_retrieve = time.perf_counter() - chunks, _, embedding_ms = self._retriever.retrieve( + chunks, _, embedding_ms = await asyncio.to_thread( + self._retriever.retrieve, question=question, top_k=top_k, ticker_filter=ticker_filter, @@ -103,7 +112,9 @@ async def answer( # authoritative XBRL facts so figures come from structured data, # not prose retrieval. if self._facts is not None: - fact_chunks = self._facts.facts_for_question(question, ticker_filter) + fact_chunks = await asyncio.to_thread( + self._facts.facts_for_question, question, ticker_filter + ) if fact_chunks: chunks = fact_chunks + chunks diff --git a/services/query-api/src/rag/retriever.py b/services/query-api/src/rag/retriever.py index 553bc75..23c880f 100644 --- a/services/query-api/src/rag/retriever.py +++ b/services/query-api/src/rag/retriever.py @@ -25,6 +25,7 @@ from __future__ import annotations import contextlib +import threading import time from typing import TYPE_CHECKING, Any @@ -58,6 +59,19 @@ # so one leg cannot dominate the fusion. _RRF_K = 60 +# pgvector explores `hnsw.ef_search` candidates per HNSW search and defaults to +# 40. Asking for more rows than that (candidate_k reaches 100 here) silently +# returns fewer and lower-quality neighbours, starving RRF and MMR of material. +# Keeping ef_search at or above the requested LIMIT is pgvector's documented +# guidance for recall. +_MIN_EF_SEARCH = 40 + +# With a WHERE clause, HNSW filters *after* the index scan, so a selective +# ticker/date filter can leave almost nothing behind. pgvector 0.8's iterative +# scan re-searches until enough rows survive the filter. "relaxed_order" keeps +# recall high at a small ordering cost, which RRF re-ranks away anyway. +_ITERATIVE_SCAN_MODE = "relaxed_order" + # Columns shared by both retrieval legs (score is appended per leg). _CHUNK_COLUMNS = "chunk_id, ticker, filing_date, section_name, chunk_text, embedding" @@ -166,6 +180,27 @@ def __init__( self._dsn = dsn self._conn: psycopg2.extensions.connection | None = None self._query_prefix = query_prefix + # Retrieval runs in a worker thread (see RAGGenerator), so the lazy + # reconnect below must not race: two threads both observing a closed + # connection would otherwise open two and leak one. + self._conn_lock = threading.Lock() + # SentenceTransformer.encode() is NOT thread-safe for this model: + # nomic-bert keeps its rotary-embedding tables (_cos_cached/_sin_cached) + # as mutable state on the shared module and resizes them to each input's + # sequence length. Two concurrent encodes with different lengths corrupt + # each other mid-forward, raising a tensor-size mismatch. Serialising + # the forward pass costs little — torch already parallelises internally + # across cores, so concurrent encodes mostly contend for the same CPU — + # and the event loop stays free either way, which is the point. + self._encode_lock = threading.Lock() + # hnsw.ef_search is *session* state on a connection shared by all worker + # threads, so the SET and the SELECT that depends on it must be applied + # as a unit; otherwise one thread's tuning silently applies to another + # thread's query. + self._search_lock = threading.Lock() + # Set once, the first time pgvector rejects the iterative-scan GUC + # (pgvector < 0.8), so the warning is logged once and not per query. + self._iterative_scan_unsupported = False logger.info("retriever_loading_model", model=model_name) # trust_remote_code: nomic ships a custom BERT (nomic-bert-2048); the # flag is ignored by standard sentence-transformers models. @@ -185,11 +220,22 @@ def close(self) -> None: self._conn = None def _get_conn(self) -> psycopg2.extensions.connection: - if self._conn is None or self._conn.closed: - self.connect() - if self._conn is None: - raise RuntimeError("Failed to establish database connection") - return self._conn + # psycopg2 connections are safe to share between threads (the driver + # serialises libpq access); only the reconnect needs guarding. + with self._conn_lock: + if self._conn is None or self._conn.closed: + self.connect() + if self._conn is None: + raise RuntimeError("Failed to establish database connection") + return self._conn + + def ping(self) -> None: + """Execute a trivial query to confirm the database is reachable. + + Raises whatever psycopg2 raises when the connection is unusable. + """ + with self._cursor() as cur: + cur.execute("SELECT 1") @contextlib.contextmanager def _cursor(self) -> Generator[psycopg2.extensions.cursor, None, None]: @@ -206,11 +252,12 @@ def embed_query(self, question: str) -> list[float]: The ``search_query:`` task prefix mirrors the ``search_document:`` prefix applied to chunks at ingestion time (asymmetric retrieval). """ - embedding = self._model.encode( - f"{self._query_prefix}{question}", - normalize_embeddings=True, - show_progress_bar=False, - ) + with self._encode_lock: + embedding = self._model.encode( + f"{self._query_prefix}{question}", + normalize_embeddings=True, + show_progress_bar=False, + ) return embedding.tolist() def verify_embedding_model_consistency(self) -> None: @@ -265,6 +312,33 @@ def _build_filters( params.append(filing_date_to) return clauses, params + def _apply_hnsw_tuning( + self, + cur: psycopg2.extensions.cursor, + candidate_k: int, + *, + filtered: bool, + ) -> None: + """Raise HNSW search effort to match how many candidates we ask for. + + Without this, pgvector's default ef_search of 40 caps the vector leg + well below candidate_k and quietly degrades recall. The connection + runs in autocommit mode, so SET LOCAL would be reverted immediately — + these are session-level SETs, which is what we want on a long-lived + connection. + """ + cur.execute("SET hnsw.ef_search = %s", (max(candidate_k, _MIN_EF_SEARCH),)) + + if not filtered or self._iterative_scan_unsupported: + return + try: + cur.execute("SET hnsw.iterative_scan = %s", (_ITERATIVE_SCAN_MODE,)) + except psycopg2.Error as exc: + # pgvector < 0.8 has no such GUC. Filtered recall is worse without + # it, but the query itself is still correct, so carry on. + self._iterative_scan_unsupported = True + logger.warning("hnsw_iterative_scan_unsupported", error=str(exc)) + def _vector_search( self, query_embedding: list[float], @@ -283,7 +357,10 @@ def _vector_search( LIMIT %s """ params = [query_embedding, *filter_params, query_embedding, candidate_k] - with self._cursor() as cur: + # Lock spans the tuning SET and the query it applies to — see + # _search_lock in __init__. + with self._search_lock, self._cursor() as cur: + self._apply_hnsw_tuning(cur, candidate_k, filtered=bool(filter_clauses)) cur.execute(sql, params) rows: list[tuple[Any, ...]] = cur.fetchall() return rows diff --git a/services/query-api/tests/test_generator.py b/services/query-api/tests/test_generator.py index 301d58b..7f9ab60 100644 --- a/services/query-api/tests/test_generator.py +++ b/services/query-api/tests/test_generator.py @@ -5,7 +5,10 @@ from __future__ import annotations +import asyncio import os +import threading +import time from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -128,6 +131,121 @@ async def test_no_chunks_returns_no_answer(self) -> None: mock_llm.generate.assert_not_called() +# ── Event loop must stay free during retrieval ─────────────────── + +class TestEventLoopIsNotBlocked: + """Retrieval is synchronous CPU + IO work (a sentence-transformers forward + pass and psycopg2 queries). Awaiting it inline would pin the event loop + for the whole query, so a single-worker uvicorn would serve exactly one + request at a time regardless of how many arrive. + """ + + @staticmethod + def _retriever_recording_thread(record: dict[str, int]) -> MagicMock: + def _retrieve(**_kwargs: object) -> tuple: + record["retrieve"] = threading.get_ident() + return ([], [0.1] * 768, 1.0) + + mock_retriever = MagicMock() + mock_retriever.retrieve.side_effect = _retrieve + return mock_retriever + + @pytest.mark.asyncio + async def test_retrieval_runs_off_the_event_loop_thread(self) -> None: + record: dict[str, int] = {} + mock_llm = AsyncMock() + mock_llm.model_name = "test-model" + + gen = RAGGenerator( + retriever=self._retriever_recording_thread(record), llm=mock_llm + ) + await gen.answer("What are Apple's risks?") + + assert record["retrieve"] != threading.get_ident() + + @pytest.mark.asyncio + async def test_fact_lookup_runs_off_the_event_loop_thread(self) -> None: + """The XBRL lookup is another blocking psycopg2 round trip.""" + record: dict[str, int] = {} + + def _facts_for_question(*_args: object) -> list: + record["facts"] = threading.get_ident() + return [] + + mock_facts = MagicMock() + mock_facts.facts_for_question.side_effect = _facts_for_question + + mock_llm = AsyncMock() + mock_llm.model_name = "test-model" + + gen = RAGGenerator( + retriever=self._retriever_recording_thread({}), + llm=mock_llm, + facts=mock_facts, + ) + await gen.answer("What was revenue in 2024?", ticker_filter="AAPL") + + assert record["facts"] != threading.get_ident() + + @pytest.mark.asyncio + async def test_concurrent_queries_overlap(self) -> None: + """Three slow retrievals must run concurrently, not back to back.""" + retrieve_seconds = 0.2 + concurrency = 3 + + def _slow_retrieve(**_kwargs: object) -> tuple: + time.sleep(retrieve_seconds) # blocking on purpose + return ([], [0.1] * 768, 1.0) + + mock_retriever = MagicMock() + mock_retriever.retrieve.side_effect = _slow_retrieve + mock_llm = AsyncMock() + mock_llm.model_name = "test-model" + + gen = RAGGenerator(retriever=mock_retriever, llm=mock_llm) + + started = time.perf_counter() + await asyncio.gather(*(gen.answer("question") for _ in range(concurrency))) + elapsed = time.perf_counter() - started + + serial = retrieve_seconds * concurrency + # Generous margin: the point is "clearly not serial", not a precise time. + assert elapsed < serial * 0.7, ( + f"{concurrency} queries took {elapsed:.2f}s; " + f"serial execution would be ~{serial:.2f}s" + ) + + @pytest.mark.asyncio + async def test_other_coroutines_progress_during_retrieval(self) -> None: + """A blocked loop would starve every unrelated task on it.""" + ticks = 0 + + async def _ticker() -> None: + nonlocal ticks + while True: + await asyncio.sleep(0.01) + ticks += 1 + + def _slow_retrieve(**_kwargs: object) -> tuple: + time.sleep(0.2) + return ([], [0.1] * 768, 1.0) + + mock_retriever = MagicMock() + mock_retriever.retrieve.side_effect = _slow_retrieve + mock_llm = AsyncMock() + mock_llm.model_name = "test-model" + + gen = RAGGenerator(retriever=mock_retriever, llm=mock_llm) + + background = asyncio.create_task(_ticker()) + try: + await gen.answer("question") + finally: + background.cancel() + + assert ticks > 0, "event loop made no progress while retrieval ran" + + # ── Auth ───────────────────────────────────────────────────────── class TestAuth: diff --git a/services/query-api/tests/test_llm_backends.py b/services/query-api/tests/test_llm_backends.py index 0245091..dd80f2b 100644 --- a/services/query-api/tests/test_llm_backends.py +++ b/services/query-api/tests/test_llm_backends.py @@ -135,6 +135,91 @@ async def test_generate_disables_thinking(self) -> None: assert payload["options"]["num_predict"] == 64 +class TestOllamaContextWindow: + """Ollama defaults num_ctx to 4096 and silently drops the overflow from the + START of the prompt — taking the system prompt and citation instruction + with it. A top_k=5 RAG prompt already approaches that, and injected XBRL + facts push it over, so num_ctx must be requested explicitly. + """ + + @staticmethod + def _mock_session() -> MagicMock: + mock_resp = AsyncMock() + mock_resp.raise_for_status = MagicMock() + mock_resp.json = AsyncMock(return_value={"response": "ok"}) + + mock_post_ctx = MagicMock() + mock_post_ctx.__aenter__ = AsyncMock(return_value=mock_resp) + mock_post_ctx.__aexit__ = AsyncMock(return_value=False) + + mock_session = MagicMock() + mock_session.post = MagicMock(return_value=mock_post_ctx) + return mock_session + + @staticmethod + def _session_ctx(mock_session: MagicMock) -> MagicMock: + ctx = MagicMock() + ctx.__aenter__ = AsyncMock(return_value=mock_session) + ctx.__aexit__ = AsyncMock(return_value=False) + return ctx + + @pytest.mark.asyncio + async def test_payload_requests_an_explicit_context_window(self) -> None: + from src.llm.ollama_backend import DEFAULT_NUM_CTX, OllamaBackend + + mock_session = self._mock_session() + with patch( + "src.llm.ollama_backend.aiohttp.ClientSession", + return_value=self._session_ctx(mock_session), + ): + backend = OllamaBackend(base_url="http://localhost:11434", model="mistral:7b") + await backend.generate("prompt") + + options = mock_session.post.call_args.kwargs["json"]["options"] + assert options["num_ctx"] == DEFAULT_NUM_CTX + + @pytest.mark.asyncio + async def test_context_window_is_configurable(self) -> None: + from src.llm.ollama_backend import OllamaBackend + + mock_session = self._mock_session() + with patch( + "src.llm.ollama_backend.aiohttp.ClientSession", + return_value=self._session_ctx(mock_session), + ): + backend = OllamaBackend( + base_url="http://localhost:11434", model="mistral:7b", num_ctx=32768 + ) + await backend.generate("prompt") + + assert mock_session.post.call_args.kwargs["json"]["options"]["num_ctx"] == 32768 + + def test_default_holds_a_top_k_5_rag_prompt(self) -> None: + """Guards the default against being lowered below a realistic prompt.""" + from src.llm.ollama_backend import DEFAULT_NUM_CTX + + chunk_tokens = 512 * 5 # top_k=5 chunks at the chunker's budget + answer_tokens = 1024 # generator's max_tokens + assert chunk_tokens + answer_tokens <= DEFAULT_NUM_CTX + + def test_oversized_prompt_is_flagged_not_silently_truncated(self) -> None: + from src.llm.ollama_backend import OllamaBackend + + backend = OllamaBackend(model="mistral:7b", num_ctx=1024) + with patch("src.llm.ollama_backend.logger") as mock_logger: + backend._warn_if_context_exceeded("x" * 40_000, max_tokens=1024) + mock_logger.warning.assert_called_once() + assert mock_logger.warning.call_args.args[0] == "ollama_prompt_may_exceed_context" + + def test_prompt_within_budget_is_not_flagged(self) -> None: + from src.llm.ollama_backend import OllamaBackend + + backend = OllamaBackend(model="mistral:7b", num_ctx=8192) + with patch("src.llm.ollama_backend.logger") as mock_logger: + backend._warn_if_context_exceeded("short prompt", max_tokens=1024) + mock_logger.warning.assert_not_called() + + # ── OpenAIBackend ──────────────────────────────────────────────── diff --git a/services/query-api/tests/test_main.py b/services/query-api/tests/test_main.py index f6d0632..343c7c0 100644 --- a/services/query-api/tests/test_main.py +++ b/services/query-api/tests/test_main.py @@ -98,7 +98,7 @@ def test_health_returns_503_on_db_error(self) -> None: import src.main as main_mod mock_retriever = MagicMock() - mock_retriever._get_conn.side_effect = Exception("db unreachable") + mock_retriever.ping.side_effect = Exception("db unreachable") original_retriever = main_mod._retriever original_generator = main_mod._generator @@ -113,6 +113,36 @@ def test_health_returns_503_on_db_error(self) -> None: main_mod._retriever = original_retriever main_mod._generator = original_generator + def test_health_probe_runs_off_the_event_loop_thread(self) -> None: + """A hung database must not stall the loop for every in-flight request.""" + import threading + + from fastapi.testclient import TestClient + + import src.main as main_mod + + probe_thread: dict[str, int] = {} + mock_retriever = MagicMock() + mock_retriever.ping.side_effect = lambda: probe_thread.setdefault( + "id", threading.get_ident() + ) + + original_retriever = main_mod._retriever + original_generator = main_mod._generator + main_mod._retriever = mock_retriever + main_mod._generator = MagicMock() + + try: + client = TestClient(app=main_mod.app, raise_server_exceptions=False) + response = client.get("/health") + assert response.status_code == 200 + # TestClient drives the loop from its own thread; the probe must + # not have run on whichever thread the handler coroutine used. + assert probe_thread["id"] != threading.get_ident() + finally: + main_mod._retriever = original_retriever + main_mod._generator = original_generator + def test_ready_returns_ready(self) -> None: from fastapi.testclient import TestClient diff --git a/services/query-api/tests/test_retriever.py b/services/query-api/tests/test_retriever.py index 0db428f..812db2b 100644 --- a/services/query-api/tests/test_retriever.py +++ b/services/query-api/tests/test_retriever.py @@ -5,6 +5,7 @@ from __future__ import annotations +import time from unittest.mock import MagicMock, patch import numpy as np @@ -25,6 +26,15 @@ def _row(chunk_id: str, embedding: np.ndarray, score: float = 0.9) -> tuple: return (chunk_id, "AAPL", "2024-11-01", "Item 1A", f"Text of {chunk_id}", embedding, score) +def _search_sql(mock_cur: MagicMock) -> list[str]: + """Executed statements, excluding the HNSW tuning SETs that precede them.""" + return [ + call.args[0] + for call in mock_cur.execute.call_args_list + if not call.args[0].lstrip().upper().startswith("SET ") + ] + + def _make_retriever(query_embedding: np.ndarray) -> tuple: """Build a Retriever with mocked model and DB connection.""" from src.rag.retriever import Retriever @@ -79,7 +89,7 @@ def test_retrieve_with_ticker_filter(self) -> None: assert chunks[0].relevance_score == 1.0 assert len(embedding) == _DIM # Both legs (vector + lexical) ran, each with the ticker filter. - executed = [call[0][0] for call in mock_cur.execute.call_args_list] + executed = _search_sql(mock_cur) assert len(executed) == 2 for sql in executed: assert "ticker = %s" in sql @@ -91,8 +101,8 @@ def test_retrieve_without_ticker_filter(self) -> None: chunks, _, _ = retriever.retrieve("General question", top_k=3) assert chunks == [] - for call in mock_cur.execute.call_args_list: - assert "ticker = %s" not in call[0][0] + for sql in _search_sql(mock_cur): + assert "ticker = %s" not in sql def test_retrieve_with_date_filters(self) -> None: retriever, mock_cur = _make_retriever(_unit(0)) @@ -105,8 +115,7 @@ def test_retrieve_with_date_filters(self) -> None: filing_date_to="2024-12-31", ) - for call in mock_cur.execute.call_args_list: - sql = call[0][0] + for sql in _search_sql(mock_cur): assert "filing_date >= %s" in sql assert "filing_date <= %s" in sql @@ -128,6 +137,237 @@ def execute(sql: str, params: object = None) -> None: assert chunks[0].chunk_id == "chunk1" +class TestHnswTuning: + """pgvector caps HNSW candidates at hnsw.ef_search (default 40). + + Requesting more rows than that silently returns fewer, lower-quality + neighbours, which starves RRF and MMR of material to work with. + """ + + @staticmethod + def _executed_sql(mock_cur: MagicMock) -> list[str]: + return [call.args[0] for call in mock_cur.execute.call_args_list] + + @staticmethod + def _set_params(mock_cur: MagicMock, guc: str) -> object | None: + for call in mock_cur.execute.call_args_list: + if guc in call.args[0] and len(call.args) > 1: + return call.args[1][0] + return None + + def test_ef_search_covers_the_requested_candidate_count(self) -> None: + from src.rag.retriever import _CANDIDATE_MULTIPLIER + + retriever, mock_cur = _make_retriever(_unit(0)) + mock_cur.fetchall.return_value = [_row("chunk1", _unit(0))] + + top_k = 20 + retriever.retrieve("revenue", top_k=top_k) + + ef_search = self._set_params(mock_cur, "hnsw.ef_search") + assert ef_search is not None, "ef_search was never set" + assert ef_search >= top_k * _CANDIDATE_MULTIPLIER + + def test_ef_search_never_drops_below_the_pgvector_default(self) -> None: + from src.rag.retriever import _MIN_EF_SEARCH + + retriever, mock_cur = _make_retriever(_unit(0)) + mock_cur.fetchall.return_value = [_row("chunk1", _unit(0))] + + retriever.retrieve("revenue", top_k=1) + + assert self._set_params(mock_cur, "hnsw.ef_search") == _MIN_EF_SEARCH + + def test_ef_search_is_applied_before_the_vector_query(self) -> None: + """A SET issued after the search would not affect it.""" + retriever, mock_cur = _make_retriever(_unit(0)) + mock_cur.fetchall.return_value = [_row("chunk1", _unit(0))] + + retriever.retrieve("revenue", top_k=5) + + statements = self._executed_sql(mock_cur) + set_idx = next(i for i, s in enumerate(statements) if "hnsw.ef_search" in s) + scan_idx = next(i for i, s in enumerate(statements) if "<=>" in s) + assert set_idx < scan_idx + + def test_filtered_search_enables_iterative_scan(self) -> None: + """HNSW filters after the scan, so a selective filter needs re-searching.""" + from src.rag.retriever import _ITERATIVE_SCAN_MODE + + retriever, mock_cur = _make_retriever(_unit(0)) + mock_cur.fetchall.return_value = [_row("chunk1", _unit(0))] + + retriever.retrieve("revenue", top_k=5, ticker_filter="AAPL") + + assert self._set_params(mock_cur, "hnsw.iterative_scan") == _ITERATIVE_SCAN_MODE + + def test_unfiltered_search_leaves_iterative_scan_off(self) -> None: + """Without a filter every neighbour survives, so plain HNSW is faster.""" + retriever, mock_cur = _make_retriever(_unit(0)) + mock_cur.fetchall.return_value = [_row("chunk1", _unit(0))] + + retriever.retrieve("revenue", top_k=5) + + assert all( + "hnsw.iterative_scan" not in s for s in self._executed_sql(mock_cur) + ) + + def test_retrieval_survives_pgvector_without_iterative_scan(self) -> None: + """pgvector < 0.8 has no such GUC; recall suffers but queries still run.""" + import psycopg2 + + retriever, mock_cur = _make_retriever(_unit(0)) + mock_cur.fetchall.return_value = [_row("chunk1", _unit(0))] + + def _execute(sql: str, params: object = None) -> None: + if "hnsw.iterative_scan" in sql: + raise psycopg2.Error("unrecognized configuration parameter") + + mock_cur.execute.side_effect = _execute + + chunks, _, _ = retriever.retrieve("revenue", top_k=5, ticker_filter="AAPL") + + assert len(chunks) == 1 + assert retriever._iterative_scan_unsupported is True + + def test_unsupported_iterative_scan_is_not_retried(self) -> None: + """The probe must not repeat on every single query.""" + import psycopg2 + + retriever, mock_cur = _make_retriever(_unit(0)) + mock_cur.fetchall.return_value = [_row("chunk1", _unit(0))] + mock_cur.execute.side_effect = lambda sql, params=None: ( + _ for _ in () + ).throw(psycopg2.Error("nope")) if "hnsw.iterative_scan" in sql else None + + retriever.retrieve("revenue", top_k=5, ticker_filter="AAPL") + first_attempts = sum( + 1 for c in mock_cur.execute.call_args_list if "hnsw.iterative_scan" in c.args[0] + ) + retriever.retrieve("revenue", top_k=5, ticker_filter="AAPL") + total_attempts = sum( + 1 for c in mock_cur.execute.call_args_list if "hnsw.iterative_scan" in c.args[0] + ) + + assert first_attempts == 1 + assert total_attempts == 1 + + +class TestThreadSafety: + """Retrieval runs in a thread pool, so the shared model and connection are + touched concurrently. + + SentenceTransformer.encode() is not thread-safe for nomic-bert: the model + keeps its rotary-embedding tables as mutable state and resizes them to each + input's sequence length, so two overlapping encodes raise a tensor-size + mismatch mid-forward. hnsw.ef_search has the same class of problem at the + database level — it is session state on a shared connection. + """ + + @staticmethod + def _overlap_detector() -> tuple[MagicMock, list[int]]: + """Return a side_effect that records the max concurrent entry count.""" + import threading + + state = {"active": 0, "peak": 0} + guard = threading.Lock() + peak: list[int] = [] + + def _tracked(*_args: object, **_kwargs: object) -> np.ndarray: + with guard: + state["active"] += 1 + state["peak"] = max(state["peak"], state["active"]) + time.sleep(0.02) # widen the window for a race to show up + with guard: + state["active"] -= 1 + peak.append(state["peak"]) + return _unit(0) + + return MagicMock(side_effect=_tracked), peak + + def test_concurrent_embed_query_never_overlaps(self) -> None: + from concurrent.futures import ThreadPoolExecutor + + retriever, _ = _make_retriever(_unit(0)) + tracked, peak = self._overlap_detector() + retriever._model.encode = tracked + + with ThreadPoolExecutor(max_workers=8) as pool: + list(pool.map(retriever.embed_query, [f"question {i}" for i in range(8)])) + + assert max(peak) == 1, ( + f"{max(peak)} concurrent encodes observed — the model forward pass " + "must be serialised or nomic-bert's cached rotary tables corrupt" + ) + + def test_concurrent_embed_query_returns_correct_results(self) -> None: + from concurrent.futures import ThreadPoolExecutor + + retriever, _ = _make_retriever(_unit(0)) + with ThreadPoolExecutor(max_workers=8) as pool: + results = list( + pool.map(retriever.embed_query, [f"question {i}" for i in range(8)]) + ) + assert all(len(r) == _DIM for r in results) + + def test_concurrent_vector_search_never_overlaps(self) -> None: + """The ef_search SET and the query it tunes must apply as a unit.""" + from concurrent.futures import ThreadPoolExecutor + + retriever, mock_cur = _make_retriever(_unit(0)) + tracked, peak = self._overlap_detector() + mock_cur.execute = tracked + mock_cur.fetchall.return_value = [_row("chunk1", _unit(0))] + + def _search(_i: int) -> object: + return retriever._vector_search([0.1] * _DIM, 20, [], []) + + with ThreadPoolExecutor(max_workers=6) as pool: + list(pool.map(_search, range(6))) + + assert max(peak) == 1 + + def test_reconnect_is_not_raced(self) -> None: + """Two threads seeing a closed connection must not both open one.""" + from concurrent.futures import ThreadPoolExecutor + + retriever, _ = _make_retriever(_unit(0)) + retriever._conn = None + + connects: list[int] = [] + + def _connect() -> None: + time.sleep(0.01) + conn = MagicMock() + conn.closed = False + retriever._conn = conn + connects.append(1) + + with ( + patch.object(retriever, "connect", side_effect=_connect), + ThreadPoolExecutor(max_workers=8) as pool, + ): + list(pool.map(lambda _i: retriever._get_conn(), range(8))) + + assert sum(connects) == 1 + + +class TestPing: + def test_ping_executes_a_trivial_query(self) -> None: + retriever, mock_cur = _make_retriever(_unit(0)) + retriever.ping() + assert mock_cur.execute.call_args.args[0] == "SELECT 1" + + def test_ping_propagates_connection_errors(self) -> None: + import psycopg2 + import pytest + + retriever, mock_cur = _make_retriever(_unit(0)) + mock_cur.execute.side_effect = psycopg2.OperationalError("server closed") + with pytest.raises(psycopg2.OperationalError): + retriever.ping() + + class TestFuseRRF: """Tests for Reciprocal Rank Fusion of the two retrieval legs."""