diff --git a/.env.example b/.env.example index 33ac18d..d0db7f6 100644 --- a/.env.example +++ b/.env.example @@ -11,3 +11,6 @@ LLM_MODEL=claude-haiku-5-5 LLM_TIMEOUT_S=30 LLM_MAX_RETRIES=0 LLM_MAX_TOKENS=4096 + +EMBEDDING_MODEL=sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2 +RAG_TOP_K=5 diff --git a/Dockerfile b/Dockerfile index c0f070e..3e4fedb 100644 --- a/Dockerfile +++ b/Dockerfile @@ -11,6 +11,14 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ build-essential \ && rm -rf /var/lib/apt/lists/* +ARG EMBEDDING_MODEL=sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2 +ENV EMBEDDING_MODEL=${EMBEDDING_MODEL} \ + EMBEDDING_CACHE_DIR=/opt/fastembed \ + EMBEDDING_LOCAL_FILES_ONLY=true + +RUN pip install "fastembed>=0.9.1" \ + && python -c "import os; from fastembed import TextEmbedding; TextEmbedding(model_name=os.environ['EMBEDDING_MODEL'], cache_dir=os.environ['EMBEDDING_CACHE_DIR'])" + COPY pyproject.toml ./ COPY src ./src COPY main.py ./ diff --git a/Makefile b/Makefile index 07450c9..f86c00d 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: help install run up down stop sh test unit-test integration-test e2e-test lint format build migrate migration migrate-test +.PHONY: help install run up down stop sh test unit-test integration-test e2e-test lint format build migrate migration migrate-test reindex DC = docker compose SERVICE = api @@ -51,5 +51,8 @@ migrate: ## Apply migrations to the dev database migrate-test: ## Apply migrations to the test database $(DC) run --rm $(SERVICE) sh -lc 'DATABASE_URL=$$TEST_DATABASE_URL alembic upgrade head' +reindex: ## Index catalog products for the assistant (incremental) + $(DC) run --rm $(SERVICE) python -m contexts.assistant.ui.cli.reindex + migration: ## Create a new migration. Usage: make migration m="add products table" $(DC) run --rm $(SERVICE) alembic revision --autogenerate -m "$(m)" diff --git a/README.md b/README.md index b3d697d..8b6a94a 100644 --- a/README.md +++ b/README.md @@ -4,13 +4,14 @@ **Generative AI on FastAPI, done the clean way.** -A reference **Generative AI backend** in Python: an LLM-powered assistant over a product catalog, built with **FastAPI**, **Domain-Driven Design (DDD)** and the **Hexagonal (Ports & Adapters)** pattern. The model is an adapter behind a port — swap providers, fake it in tests, and keep it out of your domain. +A reference **Generative AI backend** in Python: a **RAG** assistant that answers from a product catalog, built with **FastAPI**, **Domain-Driven Design (DDD)** and the **Hexagonal (Ports & Adapters)** pattern. The LLM, the embeddings and the vector store are adapters behind ports — swap providers, fake them in tests, and keep them out of your domain. [![CI](https://github.com/jcmoro/python-fastapi/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/jcmoro/python-fastapi/actions/workflows/ci.yml) -![Tests](https://img.shields.io/badge/tests-99%20passing-2e7d32?style=flat) +![Tests](https://img.shields.io/badge/tests-137%20passing-2e7d32?style=flat) ![Architecture](https://img.shields.io/badge/architecture-DDD%20%2B%20Hexagonal-2e7d32?style=flat) ![mypy](https://img.shields.io/badge/mypy-strict-1565c0?style=flat) -![GenAI](https://img.shields.io/badge/GenAI-LLM%20assistant-6a1b9a?style=flat) +![GenAI](https://img.shields.io/badge/GenAI-RAG%20assistant-6a1b9a?style=flat) +![pgvector](https://img.shields.io/badge/vectors-pgvector-336791?style=flat&logo=postgresql&logoColor=white) ![Claude](https://img.shields.io/badge/LLM-Claude-D97757?style=flat&logo=anthropic&logoColor=white) ![Python](https://img.shields.io/badge/Python-3.12-3776AB?style=flat&logo=python&logoColor=white) ![FastAPI](https://img.shields.io/badge/FastAPI-0.162-009688?style=flat&logo=fastapi&logoColor=white) @@ -36,8 +37,10 @@ A reference **Generative AI backend** in Python: an LLM-powered assistant over a Most GenAI demos wire the SDK straight into the route handler. This one treats the LLM like any other infrastructure dependency: - **The LLM is a port.** The domain defines `LLMClient`; `AnthropicLLMClient` is just one adapter. The domain and application layers never import `anthropic`. -- **Tests never call a real model.** Unit and e2e tests use `FakeLLMClient` at the port boundary. The real adapter is tested against a mocked HTTP transport — no network, no cost. -- **Failures have a contract.** Transient LLM failures (connection, timeout, 429, 5xx) return **503**. Configuration errors (bad key, unknown model) are not disguised as outages. +- **Retrieval is a port too.** `EmbeddingClient` and `DocumentSearch` are ports; embeddings run locally with a multilingual model (no API key, no cost) and vectors live in the existing Postgres with `pgvector`. +- **Indexing is explicit and incremental.** `make reindex` re-embeds only products whose content or embedding model changed (a `content_hash` per document), in batches, and deletes documents only after confirming their product no longer exists. `catalog` doesn't know the assistant exists. +- **Tests never call a real model.** Unit and e2e tests use fakes at the port boundary (`FakeLLMClient`, `FakeEmbeddingClient`, `FakeDocumentSearch`). The real LLM adapter is tested against a mocked HTTP transport, and the vector adapters against a real pgvector database. +- **Failures have a contract.** If the LLM is unavailable (connection, timeout, 429, 5xx), the answer still comes back — **200** with the retrieved products and `fallback_used: true`. Configuration errors (bad key, unknown model) are not disguised as outages. - **Latency is bounded.** `LLM_TIMEOUT_S` applies per attempt and `LLM_MAX_RETRIES` defaults to 0, so the worst case is `timeout * (retries + 1)`. - **Every answer reports its cost.** Responses include the model used, input/output tokens and latency. @@ -45,9 +48,11 @@ Most GenAI demos wire the SDK straight into the route handler. This one treats t ## What works today -`POST /assistant/ask` — ask a natural-language question and get the model's answer plus usage. +Index the catalog, then ask. `POST /assistant/ask` retrieves the most similar products, answers from them and returns them as `sources`. ```bash +make reindex # → indexed=3 unchanged=0 deleted=0 (incremental on later runs) + curl -s -X POST http://localhost:8000/assistant/ask \ -H 'Content-Type: application/json' \ -d '{"question": "Which widgets cost less than 20 EUR?"}' @@ -58,6 +63,14 @@ Example response (values are illustrative): ```json { "text": "…", + "sources": [ + { + "product_id": "8f1c…", + "content": "SKU: PROD-1\nName: Blue widget\nPrice: 19.99 EUR", + "score": 0.61 + } + ], + "fallback_used": false, "usage": { "model": "claude-haiku-5-5", "input_tokens": 52, @@ -69,12 +82,11 @@ Example response (values are illustrative): | Status | When | |--------|------| -| `200` | The model answered. | +| `200` | The model answered from the retrieved products. Also when the LLM is unavailable (connection, timeout, 429, 5xx): then `fallback_used` is `true`, `usage` is `null` and the text lists the retrieved products. | | `422` | Empty or blank question, more than 1,000 characters, or unknown fields. | -| `503` | The LLM is unavailable: connection error, timeout, rate limit (429) or server error (5xx). | | `500` | Misconfiguration: missing or invalid API key, unknown model, invalid request. | -The assistant does not read the catalog yet — that is what RAG (next on the roadmap) adds. +The index reflects the catalog as of the last `make reindex`. There is deliberately no indexing endpoint while the API has no authentication. --- @@ -85,7 +97,7 @@ Delivered as small, shippable PRs. Full plan: [`docs/plan-assistant-genai.md`](d | Step | Scope | Status | |------|-------|--------| | 1 | `LLMClient` port, Anthropic adapter, `POST /assistant/ask`, failure contract | ✅ Done | -| 2 | **RAG** with `pgvector` in the existing Postgres: embeddings behind a port, top-k retrieval, `sources` in the answer, deterministic fallback when the LLM fails | ⏳ Planned | +| 2 | **RAG** with `pgvector` in the existing Postgres: local multilingual embeddings behind a port, incremental `make reindex`, top-k retrieval, `sources` in the answer, deterministic fallback when the LLM fails | ✅ Done | | 3 | **Agent** with LangGraph: catalog tools, step limits, token usage summed across steps, offline evals | ⏳ Planned | | 4 | **Deploy**: Kubernetes manifests, local `kind` cluster, optional cloud target | ⏳ Planned | @@ -98,6 +110,8 @@ Delivered as small, shippable PRs. Full plan: [`docs/plan-assistant-genai.md`](d | Language | Python ≥ 3.10 (type hints required) | | Web framework | FastAPI + Uvicorn | | LLM | Anthropic Claude (`anthropic` SDK), behind `LLMClient` | +| Embeddings | `fastembed`, multilingual MiniLM (384 dims), local | +| Vector store | `pgvector` (HNSW, cosine) in PostgreSQL 16 | | Validation | Pydantic v2 | | Settings | `pydantic-settings` (`BaseSettings`) | | Serialization | FastAPI native (Pydantic core) | @@ -121,7 +135,7 @@ DDD + Hexagonal. Each **bounded context** is a vertical slice with four layers: **Direction of dependencies points inward.** Outer layers depend on inner ones, never the reverse. -Cross-context contracts travel as **DTOs only** — no shared domain models between bounded contexts. `assistant` never imports `catalog`. +Cross-context contracts travel as **DTOs only** — no shared domain models between bounded contexts. `assistant` reaches `catalog` in exactly one adapter, `ListProductsCatalogReader`, through `catalog`'s application API (`ListProducts` + `ProductDTO`), plus its own dependency wiring. It never imports `catalog.domain` or `catalog.infrastructure`, and `catalog` never imports `assistant`. ### Project Layout @@ -157,25 +171,42 @@ src/ domain/ # No anthropic, no FastAPI. question.py # Value object (stripped, non-empty, ≤ 1,000 chars) completion.py # Completion: text, model, tokens, latency + catalog_product.py # assistant's own view of a product + product_document.py # ProductDocument (content + content_hash), EmbeddedDocument + retrieved_document.py # Search hit with cosine similarity exceptions.py # InvalidQuestionError, LLMUnavailableError ports/ llm_client.py # LLMClient ABC — documents which failures it raises + embedding_client.py # One embedding per text, same order + document_search.py # Top-k, most similar first + document_index.py # Hashes, upsert, delete + catalog_reader.py # Products from the catalog, no duplicates application/ - dtos.py # AskCommand, AnswerDTO, UsageDTO + dtos.py # AskCommand, AnswerDTO, SourceDTO, UsageDTO, IndexReportDTO use_cases/ - ask_assistant.py - infrastructure/ # The only place the anthropic SDK lives. + ask_assistant.py # Embed → retrieve → generate, fallback without generation + index_catalog.py # Incremental reindex + orphan deletion + infrastructure/ # The only place anthropic, fastembed and pgvector live. llm/ anthropic_llm_client.py # Timeout, retries, error classification + embeddings/ + fastembed_embedding_client.py + persistence/ + models.py # ProductDocumentModel (vector(384), HNSW cosine index) + sqlalchemy_document_index.py + sqlalchemy_document_search.py + catalog/ + list_products_catalog_reader.py # The only adapter that talks to catalog ui/ - schemas.py # AskRequest, AnswerResponse + schemas.py # AskRequest, AnswerResponse, SourceResponse routers/assistant.py # POST /assistant/ask - dependencies.py # Per-context DI wiring + cli/reindex.py # make reindex → python -m contexts.assistant.ui.cli.reindex + dependencies.py # Per-context DI wiring (shared by API and CLI) shared/ # Cross-cutting only. Keep it small. domain/ errors.py # DomainError + NotFound / Conflict / InvalidInput / ServiceUnavailable infrastructure/ - settings.py # Pydantic BaseSettings (DB + LLM) + settings.py # Pydantic BaseSettings (DB + LLM + embeddings) database.py # async engine + session factory + get_session dep db_base.py # SQLAlchemy declarative Base exception_handlers.py # Domain errors → 404 / 409 / 422 / 503 @@ -186,12 +217,12 @@ alembic/ alembic.ini tests/ conftest.py # make_product, db_engine, db_session, client (sync) - fakes/ # Port-boundary fakes shared by suites (FakeLLMClient) + fakes/ # Port-boundary fakes shared by suites (LLM, embeddings, documents) unit/ # Pure logic. No I/O. - integration/ # Real adapters: postgres, Anthropic client on a mocked transport. + integration/ # Real adapters: postgres + pgvector, Anthropic client on a mocked transport. e2e/ # AsyncClient + ASGITransport against the full app. main.py # create_app() — composes routers + handlers -docker-compose.yml # api + db + db-test (db-test on tmpfs) +docker-compose.yml # api + db + db-test (pgvector images; db-test on tmpfs) Dockerfile Makefile ``` @@ -210,6 +241,7 @@ cp .env.example .env # then set ANTHROPIC_API_KEY to use the assistant make install # Build the docker image (installs deps inside) make migrate # Apply migrations to the dev database make migrate-test # Apply migrations to the test database (required before make test) +make reindex # Index catalog products for the assistant (incremental) make run # Start dev server in the foreground (uvicorn --reload) make up # Same, in the background make down # Stop and remove containers @@ -225,7 +257,7 @@ make sh # Shell into a fresh api container make migration m="add foo" # Generate a new Alembic migration via autogenerate ``` -LLM settings (all in `.env.example`): `ANTHROPIC_API_KEY`, `LLM_MODEL` (default `claude-haiku-5-5`), `LLM_TIMEOUT_S`, `LLM_MAX_RETRIES`, `LLM_MAX_TOKENS`. +LLM settings (all in `.env.example`): `ANTHROPIC_API_KEY`, `LLM_MODEL` (default `claude-haiku-5-5`), `LLM_TIMEOUT_S`, `LLM_MAX_RETRIES`, `LLM_MAX_TOKENS`. Retrieval: `EMBEDDING_MODEL` (default `sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2`) and `RAG_TOP_K` (default 5, max 40). The embedding model (~0.24 GB) is **baked into the Docker image at build time**; at runtime it is only loaded from disk (`EMBEDDING_LOCAL_FILES_ONLY=true`, never downloaded) and warmed up when the app starts, so computing embeddings needs no network (the LLM call still does). API is served on `http://localhost:8000` with auto-generated docs: @@ -239,8 +271,8 @@ API is served on `http://localhost:8000` with auto-generated docs: | Suite | Location | What it covers | |--------------|------------------------|-------------------------------------------------------------| -| Unit | `tests/unit/` | Pure domain and application logic. Mocks at the port boundary via `FakeProductRepository` and `FakeLLMClient`. **No I/O.** | -| Integration | `tests/integration/` | Real `SqlAlchemyProductRepository` against `db-test` (transaction-rollback per test). Real `AnthropicLLMClient` on an `httpx2.MockTransport`: error classification and retry bound. | -| E2E | `tests/e2e/` | Full HTTP via `httpx.AsyncClient` + `ASGITransport`. The DB session and the `LLMClient` are overridden; each domain error is checked against its status code. | +| Unit | `tests/unit/` | Pure domain and application logic. Mocks at the port boundary via `FakeProductRepository`, `FakeLLMClient`, `FakeEmbeddingClient`, `FakeDocumentIndex` and `FakeDocumentSearch`. Covers incremental indexing and orphan deletion. **No I/O.** | +| Integration | `tests/integration/` | Real `SqlAlchemyProductRepository` against `db-test` (transaction-rollback per test). Real `AnthropicLLMClient` on an `httpx2.MockTransport`: error classification and retry bound. Real pgvector adapters: upsert, hashes, delete, nearest-first search. Catalog reader across pages. | +| E2E | `tests/e2e/` | Full HTTP via `httpx.AsyncClient` + `ASGITransport`. The DB session, `LLMClient`, `EmbeddingClient` and `DocumentSearch` are overridden; covers sources, fallback and that no indexing endpoint exists. Each domain error is checked against its status code. | --- diff --git a/alembic/env.py b/alembic/env.py index 92261d9..98b6b8b 100644 --- a/alembic/env.py +++ b/alembic/env.py @@ -10,6 +10,7 @@ # Import the project's declarative Base and every module that registers models # against it. Adding a new context with its own ORM models requires importing # them here so autogenerate sees them. +from contexts.assistant.infrastructure.persistence import models as _assistant_models # noqa: F401 from contexts.catalog.infrastructure.persistence import models as _catalog_models # noqa: F401 from shared.infrastructure.db_base import Base from shared.infrastructure.settings import get_settings diff --git a/alembic/versions/4aa1b2bf2668_create_product_documents_table.py b/alembic/versions/4aa1b2bf2668_create_product_documents_table.py new file mode 100644 index 0000000..f1eb1c4 --- /dev/null +++ b/alembic/versions/4aa1b2bf2668_create_product_documents_table.py @@ -0,0 +1,44 @@ +"""create product_documents table + +Revision ID: 4aa1b2bf2668 +Revises: 85318a34438d +Create Date: 2026-10-27 18:45:10 + +""" +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa +from pgvector.sqlalchemy import Vector + + +revision: str = '4aa1b2bf2668' +down_revision: str | None = '85318a34438d' +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.execute('CREATE EXTENSION IF NOT EXISTS vector') + op.create_table('product_documents', + sa.Column('product_id', sa.Uuid(), nullable=False), + sa.Column('content', sa.Text(), nullable=False), + sa.Column('content_hash', sa.String(length=64), nullable=False), + sa.Column('embedding', Vector(384), nullable=False), + sa.Column('updated_at', sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint('product_id') + ) + op.create_index( + 'ix_product_documents_embedding', + 'product_documents', + ['embedding'], + unique=False, + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 64}, + postgresql_ops={'embedding': 'vector_cosine_ops'}, + ) + + +def downgrade() -> None: + op.drop_index('ix_product_documents_embedding', table_name='product_documents') + op.drop_table('product_documents') diff --git a/docker-compose.yml b/docker-compose.yml index 307d391..306ee8d 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -25,7 +25,7 @@ services: command: uvicorn main:app --host 0.0.0.0 --port 8000 --reload db: - image: postgres:16-alpine + image: pgvector/pgvector:pg16 container_name: catalog-db environment: POSTGRES_USER: catalog @@ -42,7 +42,7 @@ services: retries: 10 db-test: - image: postgres:16-alpine + image: pgvector/pgvector:pg16 container_name: catalog-db-test environment: POSTGRES_USER: catalog diff --git a/docs/plan-assistant-genai.md b/docs/plan-assistant-genai.md index 23e2408..dcc0379 100644 --- a/docs/plan-assistant-genai.md +++ b/docs/plan-assistant-genai.md @@ -13,7 +13,9 @@ Delivered in four PRs. Each one is shippable on its own and keeps CI green. settings; default to the cheapest current Claude model for dev. - **Vector store**: `pgvector` inside the existing Postgres. No new infrastructure. - **Embeddings**: behind an `EmbeddingClient` port. Dev/CI adapter: `fastembed` (local ONNX model, - no API key, no cost). A hosted adapter (e.g. Voyage AI) can be added later without touching the domain. + no API key, no cost) with `sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2` + (384 dimensions, ~50 languages: Spanish and English questions both work). A hosted adapter + (e.g. Voyage AI) can be added later without touching the domain. - **Agent framework**: LangGraph, confined to `infrastructure/agent/`. - **Cross-context rule**: `assistant` never imports `catalog.domain`. It keeps its own read model (`product_documents`) fed through DTOs. @@ -21,12 +23,16 @@ Delivered in four PRs. Each one is shippable on its own and keeps CI green. boundary, same pattern as `FakeProductRepository`. - **When the LLM fails**: PR 1 answers 503 (there is nothing to retrieve yet). From PR 2 on, a deterministic fallback answer (top-k search results without generation) with `fallback_used=true`. - -## Still open (need confirmation) - -- **Indexing trigger**: synchronous re-index inside `CreateProduct` / `UpdateProduct` (simple) vs. - a domain event + background worker (closer to production). _Recommendation_: synchronous now, - event later. +- **Indexing trigger: reindex on demand.** `assistant` owns an `IndexCatalog` use case that reads + products through its own `CatalogReader` port (DTOs only) and updates `product_documents`. + `catalog` does not change and does not know `assistant` exists. It runs only via `make reindex`; + there is no public endpoint while the API has no authentication. Trade-off accepted: the index + is stale until the next reindex. + - *Rejected — synchronous inside `CreateProduct` / `UpdateProduct`*: makes `catalog` depend on + `assistant` and puts embedding latency on every write. + - *Future evolution (not in PR 2)*: `catalog` publishes `ProductCreated` / `ProductUpdated` / + `ProductDeleted` domain events through a transactional **outbox**; `assistant` consumes them + and reindexes only what changed. Needs an event/outbox mechanism the repo does not have yet. --- @@ -89,21 +95,53 @@ Delivered in four PRs. Each one is shippable on its own and keeps CI green. - `docker-compose.yml` — `db` and `db-test` move to `pgvector/pgvector:pg16`. - Alembic migration — `CREATE EXTENSION IF NOT EXISTS vector`; table `product_documents` - (`product_id UUID PK`, `content TEXT`, `embedding vector()`, `updated_at`), HNSW index - on `embedding` (cosine). -- Domain ports — `EmbeddingClient.embed(texts) -> list[list[float]]`; - `DocumentSearch.search(embedding, k) -> list[RetrievedDocument]`; `DocumentIndex.upsert(...)`. -- Infrastructure — `FastEmbedEmbeddingClient`, `SqlAlchemyDocumentSearch` / `SqlAlchemyDocumentIndex` - (`pgvector.sqlalchemy.Vector`). -- Application — `IndexProduct` use case (takes a product DTO); `AskAssistant` now retrieves top-k, - puts them in the prompt and returns them as `sources`. `AnswerDTO` gains `sources` and - `fallback_used`. + (`product_id UUID PK`, `content TEXT`, `content_hash TEXT`, `embedding vector(384)`, + `updated_at`), HNSW index on `embedding` (cosine). +- Domain ports — + - `EmbeddingClient.embed(texts) -> list[list[float]]`. + - `DocumentSearch.search(embedding, k) -> list[RetrievedDocument]`. + - `DocumentIndex`: read the stored `product_id → content_hash`, upsert documents, delete + documents by `product_id`. + - `CatalogReader.read_all() -> list[CatalogProduct]` — `assistant`'s own view of a product, + not `catalog`'s DTO. `CatalogReader.missing(ids)` confirms which ids really no longer exist + (paging is not snapshot-consistent, so absence from `read_all` is not proof of deletion). + - `EmbeddingClient.model_name` — part of the `content_hash`, so changing the model re-embeds. +- Application — + - `IndexCatalog` use case, **incremental**: builds each product's document content and its + `content_hash`; embeds and upserts only new or changed products (Product has no + `updated_at`, so the hash — over model + content — is the change detector), in batches of + 256; **deletes orphaned documents** only after `CatalogReader.missing` confirms the product no + longer exists. Returns counts (indexed, unchanged, deleted). + - `AskAssistant` now embeds the question, retrieves top-k, puts them in the prompt and returns + them as `sources`. `AnswerDTO` gains `sources` and `fallback_used`. - Fallback — `AskAssistant` catches `LLMUnavailableError` and returns the top-k results without generation, `fallback_used=true`, 200. The 503 from PR 1 no longer reaches the client. +- Infrastructure — + - `FastEmbedEmbeddingClient` (model, cache dir and `local_files_only` from settings). The + `Dockerfile` downloads the model at build time; at runtime it is only loaded from disk and + warmed up in the app lifespan. + - `SqlAlchemyDocumentSearch` (own short-lived session, so no connection is held during the LLM + call) / `SqlAlchemyDocumentIndex` (`pgvector.sqlalchemy.Vector`). `RAG_TOP_K` ≤ 40 + (HNSW default `ef_search`). + - `CatalogReader` adapter: the **only** place `assistant` touches `catalog`. It calls + `catalog`'s public application API (`ListProducts`, paginated) and maps `ProductDTO` to + `CatalogProduct`. Never `catalog.domain` or `catalog.infrastructure`. +- UI (input adapters) — + - `ui/cli/reindex.py` — entry point `python -m contexts.assistant.ui.cli.reindex`. Builds + `IndexCatalog` with the **same dependency providers** as the API (`ui/dependencies.py`, + `shared/infrastructure/database.py`), runs it and prints the counts. An input adapter like + the router; reusable as a Kubernetes `Job` in PR 4. + - `Makefile` — `make reindex` runs that module in the `api` container. + - No HTTP endpoint for indexing. - `pyproject.toml` — add `pgvector`, `fastembed`. -- Tests — unit with fakes (including fallback when the fake LLM raises); e2e asserts the fallback - flag instead of 503; integration for the search adapter against real `db-test` - (seed three documents, assert the nearest one comes back first). +- Tests — + - Unit with fakes: incremental indexing (unchanged products are not re-embedded), **orphan + deletion**, retrieval in `AskAssistant`, fallback when the fake LLM raises. + - Integration against real `db-test`: search adapter (seed three documents, assert the nearest + one comes back first) and index adapter (upsert, hashes, delete). + - E2E asserts the fallback flag instead of 503. +- README — mark roadmap step 2 as done and replace "`assistant` never imports `catalog`" with the + real rule (only the `CatalogReader` adapter, through `catalog`'s application API). --- @@ -129,7 +167,8 @@ Delivered in four PRs. Each one is shippable on its own and keeps CI green. ## PR 4 — Deploy - `k8s/` — `Deployment` (liveness on `/healthz`), `Service`, `ConfigMap`, `Secret` (template only), - `Job` running `alembic upgrade head`, Postgres with pgvector for local use. + `Job` running `alembic upgrade head`, `Job` running `python -m contexts.assistant.ui.cli.reindex`, + Postgres with pgvector for local use. - `make k8s-up` / `make k8s-down` — runs on `kind` locally. - CI — optional job that creates a kind cluster, applies the manifests and hits `/healthz`. - Optional — the same image on AWS (App Runner or ECS Fargate + RDS with pgvector). diff --git a/docs/step-16-assistant-rag.md b/docs/step-16-assistant-rag.md new file mode 100644 index 0000000..c594a0e --- /dev/null +++ b/docs/step-16-assistant-rag.md @@ -0,0 +1,210 @@ +# Paso 16 — Contexto `assistant`: RAG con pgvector (PR 2) + +Rama `feature/assistant-rag`. Plan: `docs/plan-assistant-genai.md` → «PR 2 — RAG with pgvector». + +## Objetivo + +Que una pregunta a `POST /assistant/ask` sobre productos del catálogo devuelva una respuesta +basada en esos productos, con los productos usados como `sources`, y que siga respondiendo, con +`fallback_used=true`, aunque el LLM no esté disponible. + +## Qué + +El asistente responde a partir del catálogo real: + +1. **Indexado bajo demanda.** `make reindex` ejecuta el caso de uso `IndexCatalog`, que: + - lee los productos a través del puerto `CatalogReader`; + - construye el texto de cada producto y su `content_hash`; + - calcula embeddings y hace upsert solo de los productos nuevos o cambiados; + - borra los documentos huérfanos, es decir, los de productos que ya no existen. +2. **Recuperación.** `POST /assistant/ask` calcula el embedding de la pregunta, recupera los k + documentos más cercanos (`pgvector`, distancia coseno), los incluye en el prompt y los + devuelve como `sources`. +3. **Fallback.** Si el LLM no está disponible, la respuesta es 200 con los resultados recuperados, + sin generación, y `fallback_used=true`. Sustituye al 503 del PR 1. + +## Por qué + +- El PR 1 dejó un LLM que no conoce el catálogo. Este PR es el que convierte el endpoint en un + asistente sobre los productos. +- Decisiones tomadas con el usuario (detalle en el plan): + - **Indexado bajo demanda** y no síncrono en `catalog`: `catalog` no cambia y no sabe que + `assistant` existe. + - **Solo `make reindex`**, sin endpoint público, porque la API no tiene autenticación. + - **Incremental con `content_hash`**: `Product` no tiene `updated_at`, así que el hash es lo que + detecta los cambios y evita recalcular embeddings de lo que no ha cambiado. + - **El script es un adaptador de entrada.** `python -m contexts.assistant.ui.cli.reindex` usa los + mismos proveedores de dependencias que la API, y servirá como Job de Kubernetes en el PR 4. + - **Modelo multilingüe**: `sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2`, de 384 + dimensiones. Verificado con `fastembed 0.9.1`. + - **Eventos de dominio + outbox** quedan como evolución futura, fuera de este PR. + +## Fuera de alcance + +- **Endpoint HTTP de indexado**: no se crea ninguno mientras la API no tenga autenticación. +- **Indexado automático** al crear, editar o borrar productos, por eventos de dominio y outbox: + queda como evolución futura en el plan. El índice puede quedar desactualizado hasta el + siguiente `make reindex`. +- **Cambios en `catalog`**: ni una línea. +- **Proveedor de embeddings alojado** (Voyage AI u otro): el puerto lo permite, pero no se añade. +- **Agente / LangGraph y despliegue**: son los PR 3 y 4. +- **Evaluar la calidad de las respuestas**: las evals son del PR 3. Aquí se comprueba que la + recuperación ordena bien, no la calidad del texto generado. + +## Detalles de diseño fijados aquí + +| Tema | Decisión | Porqué | +|---|---|---| +| Texto del documento | `sku`, `name`, `price_cents` y `currency` en un formato estable | Es lo que hay en `ProductDTO`. Si el formato es estable, el hash solo cambia cuando cambia el producto. | +| `content_hash` | SHA-256 del texto del documento | Determinista y sin colisiones prácticas. | +| `k` | Setting `rag_top_k`, 5 por defecto | Configurable, igual que el resto de settings del LLM. | +| Caché del modelo | ~~Volumen Docker con nombre~~ → **modelo dentro de la imagen** (R7) | El volumen evitaba descargas repetidas, pero cada contenedor nuevo seguía necesitando red. Ver R7. | +| Lectura del catálogo | Adaptador `CatalogReader` sobre `ListProducts` (API de aplicación de `catalog`), paginado hasta agotar | Es el único punto de contacto con `catalog`, y solo a través de su capa de aplicación y sus DTOs. | +| Índice vacío | `sources=[]`. El LLM responde con lo que sabe, como en el PR 1. En fallback, mensaje fijo. | No inventar resultados. | +| Salida del CLI | Una línea con `indexed=N unchanged=N deleted=N` | Es la evidencia observable del indexado incremental. | + +## Criterios de aceptación + +| # | Criterio | Cómo se demuestra | +|---|---|---| +| C1 | Con el índice vacío, `make reindex` indexa todos los productos | Ejecución real: `indexed` = número de productos | +| C2 | Una segunda ejecución sin cambios **no recalcula ningún embedding** | Test unitario: el `FakeEmbeddingClient` no recibe llamadas. Ejecución real: `indexed=0 unchanged=N` | +| C3 | Si cambia un producto, solo ese se recalcula | Test unitario | +| C4 | Si un producto desaparece del catálogo, su documento se borra | Test unitario **e** integración contra `db-test` | +| C5 | `catalog` no cambia | `git diff main -- src/contexts/catalog` vacío | +| C6 | `assistant` solo toca `catalog` en el adaptador `CatalogReader` y en su cableado de dependencias | `grep contexts.catalog src/contexts/assistant` | +| C7 | `/assistant/ask` devuelve `sources` y los documentos van en el prompt | Tests unitario y e2e | +| C8 | Si el LLM falla: 200, `fallback_used=true` y `sources` presentes | Tests unitario y e2e | +| C9 | No hay endpoint de indexado | Rutas de OpenAPI: solo `/assistant/ask` en el contexto | +| C10 | La búsqueda devuelve primero el documento más cercano | Integración contra `db-test` con pgvector | +| C11 | Los quality gates y el CI pasan | `make lint`, `make test` y CI de la PR | +| C12 | Plan y README coinciden con el código | Plan actualizado; README con el paso 2 ✅ y la regla real sobre `catalog` | + +## Decisiones tomadas durante la implementación + +| Tema | Decisión | Porqué | +|---|---|---| +| Contrato de `CatalogReader` | `read_all` garantiza «sin duplicados», y `missing(ids)` confirma las bajas (R1) | `catalog` solo expone paginación por `offset` ordenada por `created_at`, así que un producto puede faltar en `read_all`. Que falte **no** demuestra que se haya borrado: antes de borrar se confirma con `GetProduct`. Si un producto queda fuera de una ejecución, su documento sigue en el índice y se actualiza en la siguiente. | +| Distancia en SQL | Operador `<=>` de pgvector vía `.op("<=>", return_type=Float)` | Tipado para mypy strict sin silenciar nada. `cosine_distance` no es visible para mypy sobre `Mapped[list[float]]`. | +| Lectura de filas | `result.all()` en lugar de `result.tuples()` | `Result.tuples()` está deprecado en SQLAlchemy 2.1 (salía un `SADeprecationWarning`). Además, `dict(result)` falla porque `Result` tiene `keys()`. | +| Carga del modelo | Perezosa y protegida con un lock, ejecutada en hilo con `asyncio.to_thread` | ONNX es trabajo de CPU y no debe bloquear el event loop. Si nada cambia, el reindexado ni siquiera carga el modelo. | +| CLI | `asynccontextmanager(get_session)` + los proveedores de `ui/dependencies.py` | Misma transacción (commit/rollback) y mismo cableado que la API, sin una configuración paralela. | +| `usage` | Pasa a ser nullable | En el fallback no hay ninguna llamada al LLM que reportar. | +| Router | Se quita el 503 de `responses` | El fallback sustituye al 503. Los errores de configuración siguen siendo 500. | +| BD de desarrollo | Al pasar de `postgres:16-alpine` a `pgvector/pgvector:pg16` se ejecutó `REINDEX TABLE products` | Es la misma versión mayor (datos compatibles), pero cambia la libc (musl → glibc) y la ordenación de texto podría diferir. El `REINDEX` reconstruye índices sin tocar datos. Se conservaron los 3 productos. | + +## Correcciones tras la code review + +La revisión (`/code-review medium`) dio 10 hallazgos. Se comprobó cada uno contra el código: 6 se +corrigen aquí (R1–R6), 2 repiten decisiones ya documentadas y 2 son falsos positivos (detalle al +final de esta sección). El usuario eligió corregir los 6 y añadió R7. Todos empiezan por un test +que falla con el código anterior. + +| # | Qué | Por qué | Criterio de aceptación | +|---|---|---|---| +| R1 | Antes de borrar un documento «huérfano», confirmar con `catalog` (`GetProduct`, por su capa de aplicación) que el producto ya no existe | La paginación por `offset` sobre un `created_at` no único puede saltarse un producto en el borde de una página, y `IndexCatalog` lo borraba del índice aunque el producto siga existiendo | Un producto existente que `read_all` no devuelve **no** se borra. Uno que no existe sí. | +| R2 | El `content_hash` incluye el nombre del modelo de embeddings | Cambiar `EMBEDDING_MODEL` no recalculaba nada y dejaba vectores de dos modelos mezclados | Si cambia el modelo, el siguiente reindex vuelve a calcular todos los embeddings. | +| R3 | Calcular embeddings y hacer upsert en lotes de tamaño fijo | Un único `INSERT` con más de ~13.000 productos supera el límite de 65.535 parámetros de Postgres, y todo el reindex hacía rollback | Con un lote de 2 y 5 productos: llamadas de 2, 2 y 1. | +| R4 | `RAG_TOP_K` limitado a 40 | Con `hnsw.ef_search=40` (el valor por defecto), un k mayor devuelve ~40 resultados sin avisar | `rag_top_k=41` es un error de configuración. | +| R5 | La búsqueda abre y cierra su propia sesión | La sesión de la petición mantenía una conexión del pool durante toda la llamada al LLM (hasta 30 s). Con unas 15 preguntas concurrentes el pool se agota y bloquea `catalog` | Durante la llamada al LLM no queda ninguna conexión ocupada. | +| R6 | Cargar el modelo de embeddings al arrancar la app (lifespan) | La primera pregunta tras cada arranque cargaba el modelo y bloqueaba también a las concurrentes | Al arrancar la app se ejecuta un embedding de calentamiento. En tests se puede desactivar para no descargar el modelo. | +| R7 | Descargar el modelo **al construir la imagen** (Dockerfile) y, en ejecución, solo cargarlo de disco (`local_files_only=True`). Petición del usuario. | Si el modelo no viene en la imagen, cada contenedor nuevo lo descarga al arrancar: es más lento y depende de la red (y de un volumen de caché) | La imagen contiene el modelo. Un contenedor **sin red** (`--network none`) calcula embeddings. Con `local_files_only` y la caché vacía, falla en vez de descargar. Ya no hace falta el volumen `fastembed-cache`. | + +**No aplicados, con su porqué:** + +- **Acoplamiento `assistant` → `catalog`.** Repite el diseño documentado: solo la capa de + aplicación de `catalog`, sus DTOs y su cableado. +- **Docstrings «narrativos» en los puertos.** Falso positivo: el `CLAUDE.md` del repo los permite + en los puertos. +- **503 sin documentar.** Falso positivo: en esta ruta nada lanza un `ServiceUnavailableError` + distinto de `LLMUnavailableError`, y el fallback lo captura. +- **500 si fallan los embeddings o la búsqueda.** Coherente con la especificación: el fallback + cubre la caída del LLM. Sin recuperación no hay nada que devolver. + +### Evidencia de las correcciones + +**Rojo** (código anterior). Los 4 primeros fallan por el bug; el resto, porque la API que +necesitan aún no existe: + +``` +R1 test_r1_keeps_a_product_that_exists_but_was_not_listed → deleted=1 (borraba un producto existente) +R2 test_r2_changing_the_embedding_model_reembeds_everything → indexed=0 (no recalculaba nada) +R4 test_r4_rejects_more_than_the_hnsw_default_ef_search → DID NOT RAISE ValidationError +R5 test_r5_no_connection_is_held_while_the_llm_answers → assert [1] == [0] (1 conexión ocupada durante el LLM) +R1 test_r1_reports_only_ids_not_in_the_catalog (integración) → falta CatalogReader.missing +R2 test_r2_same_content_with_another_model_has_another_hash → falta el parámetro embedding_model +R3 test_r3_embeds_and_upserts_in_batches → falta el parámetro batch_size +R6 tests/e2e/test_lifespan.py → ImportError: no existe contexts.assistant.ui.lifespan +R7 test_r7_fails_instead_of_downloading_when_the_model_is_not_cached → falta el parámetro local_files_only +``` + +**Verde** (tras las correcciones): esos 9 tests y 3 de control (R1 «sigue borrando lo que ya no +existe», R4 «acepta 40», R6 «sin calentamiento por defecto»): `12 passed`. + +**R7 y R6 en ejecución real:** + +- La imagen contiene el modelo: `/opt/fastembed` ocupa 241 MB y se ha construido con + `EMBEDDING_LOCAL_FILES_ONLY=true`. +- La descarga es una capa propia del `Dockerfile`, **antes** de copiar el código, así que cambiar + código no vuelve a descargarlo. +- Contenedor con `--network none`: `embed(["¿Cuánto cuesta el widget?"])` devuelve un vector de + dimensión 384. +- Contenedor con `--network none`: `uvicorn main:app` llega a «Application startup complete» tras + cargar el modelo en el lifespan. +- Con `local_files_only` y la caché vacía, fastembed lanza `ValueError: Could not load model … + from any source` sin descargar nada (comprobado sin red y cubierto por el test R7). +- Ya no existe el volumen `fastembed-cache` en `docker-compose.yml`. + +**Consecuencia de R2 en un índice ya existente:** el formato del hash cambia, así que el primer +`make reindex` tras desplegar recalcula todo una vez. En desarrollo: `indexed=3 unchanged=0 +deleted=0` y después `indexed=0 unchanged=3 deleted=0`. + +## Evidencia + +| # | Resultado | Evidencia | +|---|---|---| +| C1 | ✅ | `make reindex` sobre la BD de desarrollo (3 productos), primera ejecución: `indexed=3 unchanged=0 deleted=0` (1 min 23 s, incluida la descarga del modelo). `product_documents`: 3 filas, 3 hashes distintos. Test `test_c1_indexes_every_product_into_an_empty_index`. | +| C2 | ✅ | Segunda ejecución: `indexed=0 unchanged=3 deleted=0` (2,9 s; el modelo no llega a cargarse). Test `test_c2_second_run_without_changes_embeds_nothing`: el `FakeEmbeddingClient` no recibe llamadas. | +| C3 | ✅ | Test `test_c3_only_the_changed_product_is_embedded_again`: solo se recalcula el embedding del producto renombrado. No se modificaron datos reales para probarlo. | +| C4 | ✅ | Test unitario `test_c4_deletes_documents_of_products_no_longer_in_the_catalog`. Integración contra pgvector: `test_c4_delete_removes_only_the_given_documents`. | +| C5 | ✅ | `git diff origin/main -- src/contexts/catalog` vacío. | +| C6 | ✅ | `grep contexts.catalog src/contexts/assistant` solo aparece en `infrastructure/catalog/list_products_catalog_reader.py` (capa de aplicación de `catalog`), `ui/dependencies.py` y `ui/cli/reindex.py` (cableado). `catalog` no importa `assistant`. | +| C7 | ✅ | Unitarios `test_c7_returns_retrieved_documents_as_sources` y `test_c7_puts_the_documents_and_the_question_in_the_prompt`. E2E `test_c7_returns_200_with_answer_sources_and_usage`. | +| C8 | ✅ | Unitarios `test_c8_llm_unavailable_*`. E2E `test_c8_returns_200_with_fallback_when_llm_unavailable`: 200, `fallback_used=true`, `usage=null`. | +| C9 | ✅ | E2E `test_c9_assistant_exposes_no_indexing_endpoint`: las rutas `/assistant*` de OpenAPI son exactamente `["/assistant/ask"]`. | +| C10 | ✅ | Integración `test_c10_returns_the_nearest_document_first`, contra pgvector real: orden `near`, `middle`, `far`, con similitud ≈ 0,995 para el primero. | +| C11 | ✅ | `make lint`: ruff «All checks passed!», pylint 10.00/10, mypy «no issues found in 72 source files». `make test`: **137 passed** (antes 99), sin warnings. **CI de la PR #8:** run `40588443477` sobre `1d1d048`, en verde. En ambos jobs pasan el login en Docker Hub, la construcción de la imagen (con el modelo dentro), las migraciones, el lint y los tests. | +| C12 | ✅ | Plan actualizado, también con R1–R7. README: paso 2 del roadmap ✅, regla real sobre `catalog`, layout, stack, Quick Start (modelo en la imagen, `RAG_TOP_K` máximo 40) y testing al día, badge de 137 tests. Anclas y badges comprobados. | + +### Recuperación real con el modelo multilingüe + +Pregunta en español contra el índice de desarrollo, con los adaptadores reales (fastembed + pgvector): + +``` +¿Cuánto cuesta el widget? +0.361 SKU: PROD-LIVE-1 | Name: Real Widget | Price: 25.00 EUR +0.279 SKU: GATE-2 | Name: x | Price: 0.01 EUR +0.216 SKU: DUP-LIVE | Name: x | Price: 0.01 EUR +``` + +La pregunta, en español, recupera primero el producto cuyo nombre está en inglés. + +### Tests nuevos o cambiados + +| Fichero | Tests | +|---|---| +| `tests/unit/contexts/assistant/domain/test_product_document.py` | 5 (nuevo) | +| `tests/unit/contexts/assistant/application/test_index_catalog.py` | 5 (nuevo) | +| `tests/unit/contexts/assistant/application/test_ask_assistant.py` | 9 (antes 5) | +| `tests/unit/contexts/assistant/application/test_index_catalog.py` (R1–R3) | +4 | +| `tests/unit/shared/infrastructure/test_settings.py` (R4) | 2 (nuevo) | +| `tests/integration/contexts/assistant/application/test_ask_assistant_connection.py` (R5) | 1 (nuevo) | +| `tests/integration/contexts/assistant/infrastructure/test_fastembed_embedding_client.py` (R7) | 1 (nuevo) | +| `tests/e2e/test_lifespan.py` (R6) | 2 (nuevo) | +| `tests/integration/contexts/assistant/infrastructure/test_sqlalchemy_documents.py` | 7 (nuevo) | +| `tests/integration/contexts/assistant/infrastructure/test_list_products_catalog_reader.py` | 3 (nuevo) | +| `tests/e2e/test_assistant.py` | 7 (antes 5; el 503 pasa a fallback) | + +### Fuera de lo verificado + +- **Pregunta completa con el LLM real**: `.env` no tiene `ANTHROPIC_API_KEY`. La generación está cubierta por los tests con `FakeLLMClient` y por el adaptador probado en el PR 1. La recuperación real está demostrada arriba. diff --git a/main.py b/main.py index dc05577..084b639 100644 --- a/main.py +++ b/main.py @@ -1,12 +1,22 @@ +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + from fastapi import FastAPI +from contexts.assistant.ui import lifespan as assistant_lifespan from contexts.assistant.ui.routers.assistant import router as assistant_router from contexts.catalog.ui.routers.products import router as products_router from shared.infrastructure.exception_handlers import register_exception_handlers from shared.infrastructure.settings import get_settings -def create_app() -> FastAPI: +@asynccontextmanager +async def _warm_up_models(_application: FastAPI) -> AsyncIterator[None]: + await assistant_lifespan.warm_up_embeddings() + yield + + +def create_app(*, warm_up_models: bool = False) -> FastAPI: settings = get_settings() application = FastAPI( @@ -16,6 +26,7 @@ def create_app() -> FastAPI: docs_url="/docs", redoc_url="/redoc", openapi_url="/openapi.json", + lifespan=_warm_up_models if warm_up_models else None, ) @application.get( @@ -34,4 +45,4 @@ async def healthz() -> dict[str, str]: return application -app = create_app() +app = create_app(warm_up_models=True) diff --git a/pyproject.toml b/pyproject.toml index ea6a309..985c896 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,8 @@ dependencies = [ "asyncpg>=0.30", "alembic>=1.13", "anthropic>=2.3", + "pgvector>=0.5.1", + "fastembed>=0.9.1", ] [project.optional-dependencies] diff --git a/src/contexts/assistant/application/dtos.py b/src/contexts/assistant/application/dtos.py index 6974cf9..15753d4 100644 --- a/src/contexts/assistant/application/dtos.py +++ b/src/contexts/assistant/application/dtos.py @@ -1,8 +1,10 @@ from __future__ import annotations from dataclasses import dataclass +from uuid import UUID from contexts.assistant.domain.completion import Completion +from contexts.assistant.domain.retrieved_document import RetrievedDocument @dataclass(frozen=True, slots=True) @@ -18,15 +20,30 @@ class UsageDTO: latency_ms: int +@dataclass(frozen=True, slots=True) +class SourceDTO: + product_id: UUID + content: str + score: float + + @classmethod + def from_retrieved(cls, document: RetrievedDocument) -> SourceDTO: + return cls(product_id=document.product_id, content=document.content, score=document.score) + + @dataclass(frozen=True, slots=True) class AnswerDTO: text: str - usage: UsageDTO + sources: tuple[SourceDTO, ...] + fallback_used: bool + usage: UsageDTO | None @classmethod - def from_completion(cls, completion: Completion) -> AnswerDTO: + def from_completion(cls, completion: Completion, sources: tuple[SourceDTO, ...]) -> AnswerDTO: return cls( text=completion.text, + sources=sources, + fallback_used=False, usage=UsageDTO( model=completion.model, input_tokens=completion.input_tokens, @@ -34,3 +51,14 @@ def from_completion(cls, completion: Completion) -> AnswerDTO: latency_ms=completion.latency_ms, ), ) + + @classmethod + def fallback(cls, text: str, sources: tuple[SourceDTO, ...]) -> AnswerDTO: + return cls(text=text, sources=sources, fallback_used=True, usage=None) + + +@dataclass(frozen=True, slots=True) +class IndexReportDTO: + indexed: int + unchanged: int + deleted: int diff --git a/src/contexts/assistant/application/use_cases/ask_assistant.py b/src/contexts/assistant/application/use_cases/ask_assistant.py index cf24f2c..3fca6ce 100644 --- a/src/contexts/assistant/application/use_cases/ask_assistant.py +++ b/src/contexts/assistant/application/use_cases/ask_assistant.py @@ -1,24 +1,70 @@ -from contexts.assistant.application.dtos import AnswerDTO, AskCommand +from contexts.assistant.application.dtos import AnswerDTO, AskCommand, SourceDTO +from contexts.assistant.domain.exceptions import LLMUnavailableError +from contexts.assistant.domain.ports.document_search import DocumentSearch +from contexts.assistant.domain.ports.embedding_client import EmbeddingClient from contexts.assistant.domain.ports.llm_client import LLMClient from contexts.assistant.domain.question import Question SYSTEM_PROMPT = ( "You are the assistant of a product catalog. " - "Answer the user's question concisely. " - "If you do not know the answer, say so instead of guessing." + "Answer the user's question concisely, using only the catalog products provided. " + "If they do not answer the question, say so instead of guessing." +) + +FALLBACK_WITH_SOURCES = ( + "The assistant is unavailable right now. " + "These catalog products best match your question:" +) +FALLBACK_WITHOUT_SOURCES = ( + "The assistant is unavailable right now and no catalog products match your question." ) class AskAssistant: - def __init__(self, llm_client: LLMClient, max_tokens: int) -> None: + def __init__( + self, + llm_client: LLMClient, + embedding_client: EmbeddingClient, + document_search: DocumentSearch, + max_tokens: int, + top_k: int, + ) -> None: self._llm_client = llm_client + self._embedding_client = embedding_client + self._document_search = document_search self._max_tokens = max_tokens + self._top_k = top_k async def execute(self, command: AskCommand) -> AnswerDTO: question = Question.create(command.question) - completion = await self._llm_client.complete( - system=SYSTEM_PROMPT, - prompt=question.text, - max_tokens=self._max_tokens, - ) - return AnswerDTO.from_completion(completion) + [embedding] = await self._embedding_client.embed([question.text]) + documents = await self._document_search.search(embedding, self._top_k) + sources = tuple(SourceDTO.from_retrieved(document) for document in documents) + + try: + completion = await self._llm_client.complete( + system=SYSTEM_PROMPT, + prompt=_build_prompt(question.text, sources), + max_tokens=self._max_tokens, + ) + except LLMUnavailableError: + return AnswerDTO.fallback(_fallback_text(sources), sources) + return AnswerDTO.from_completion(completion, sources) + + +def _build_prompt(question: str, sources: tuple[SourceDTO, ...]) -> str: + if not sources: + return question + products = "\n\n".join(source.content for source in sources) + return f"Catalog products:\n\n{products}\n\nQuestion: {question}" + + +def _fallback_text(sources: tuple[SourceDTO, ...]) -> str: + if not sources: + return FALLBACK_WITHOUT_SOURCES + lines = "\n".join(f"- {_single_line(source.content)}" for source in sources) + return f"{FALLBACK_WITH_SOURCES}\n{lines}" + + +def _single_line(content: str) -> str: + return content.replace("\n", " | ") diff --git a/src/contexts/assistant/application/use_cases/index_catalog.py b/src/contexts/assistant/application/use_cases/index_catalog.py new file mode 100644 index 0000000..582f4ef --- /dev/null +++ b/src/contexts/assistant/application/use_cases/index_catalog.py @@ -0,0 +1,56 @@ +from contexts.assistant.application.dtos import IndexReportDTO +from contexts.assistant.domain.ports.catalog_reader import CatalogReader +from contexts.assistant.domain.ports.document_index import DocumentIndex +from contexts.assistant.domain.ports.embedding_client import EmbeddingClient +from contexts.assistant.domain.product_document import EmbeddedDocument, ProductDocument + +DEFAULT_BATCH_SIZE = 256 + + +class IndexCatalog: + def __init__( + self, + catalog_reader: CatalogReader, + embedding_client: EmbeddingClient, + document_index: DocumentIndex, + batch_size: int = DEFAULT_BATCH_SIZE, + ) -> None: + self._catalog_reader = catalog_reader + self._embedding_client = embedding_client + self._document_index = document_index + self._batch_size = batch_size + + async def execute(self) -> IndexReportDTO: + products = await self._catalog_reader.read_all() + model = self._embedding_client.model_name + documents = [ + ProductDocument.from_catalog_product(product, embedding_model=model) + for product in products + ] + stored_hashes = await self._document_index.content_hashes() + + changed = [ + doc for doc in documents if stored_hashes.get(doc.product_id) != doc.content_hash + ] + for start in range(0, len(changed), self._batch_size): + await self._index_batch(changed[start : start + self._batch_size]) + + unlisted_ids = sorted(set(stored_hashes) - {doc.product_id for doc in documents}) + orphan_ids = await self._catalog_reader.missing(unlisted_ids) if unlisted_ids else [] + if orphan_ids: + await self._document_index.delete(orphan_ids) + + return IndexReportDTO( + indexed=len(changed), + unchanged=len(documents) - len(changed), + deleted=len(orphan_ids), + ) + + async def _index_batch(self, batch: list[ProductDocument]) -> None: + embeddings = await self._embedding_client.embed([doc.content for doc in batch]) + await self._document_index.upsert( + [ + EmbeddedDocument(document=doc, embedding=embedding) + for doc, embedding in zip(batch, embeddings, strict=True) + ] + ) diff --git a/src/contexts/assistant/domain/catalog_product.py b/src/contexts/assistant/domain/catalog_product.py new file mode 100644 index 0000000..619c9d6 --- /dev/null +++ b/src/contexts/assistant/domain/catalog_product.py @@ -0,0 +1,11 @@ +from dataclasses import dataclass +from uuid import UUID + + +@dataclass(frozen=True, slots=True) +class CatalogProduct: + product_id: UUID + sku: str + name: str + price_cents: int + currency: str diff --git a/src/contexts/assistant/domain/ports/catalog_reader.py b/src/contexts/assistant/domain/ports/catalog_reader.py new file mode 100644 index 0000000..57f3d5b --- /dev/null +++ b/src/contexts/assistant/domain/ports/catalog_reader.py @@ -0,0 +1,17 @@ +from abc import ABC, abstractmethod +from uuid import UUID + +from contexts.assistant.domain.catalog_product import CatalogProduct + + +class CatalogReader(ABC): + @abstractmethod + async def read_all(self) -> list[CatalogProduct]: + """Returns the products currently in the catalog, without duplicates. A product may be + missing from the result even though it exists (paging is not snapshot-consistent), so + absence from this list is not proof of deletion: use ``missing`` for that. + """ + + @abstractmethod + async def missing(self, product_ids: list[UUID]) -> list[UUID]: + """Returns the given ids that do not exist in the catalog, in the same order.""" diff --git a/src/contexts/assistant/domain/ports/document_index.py b/src/contexts/assistant/domain/ports/document_index.py new file mode 100644 index 0000000..b16a837 --- /dev/null +++ b/src/contexts/assistant/domain/ports/document_index.py @@ -0,0 +1,17 @@ +from abc import ABC, abstractmethod +from uuid import UUID + +from contexts.assistant.domain.product_document import EmbeddedDocument + + +class DocumentIndex(ABC): + @abstractmethod + async def content_hashes(self) -> dict[UUID, str]: ... + + @abstractmethod + async def upsert(self, documents: list[EmbeddedDocument]) -> None: + """Inserts new documents and replaces existing ones with the same ``product_id``.""" + + @abstractmethod + async def delete(self, product_ids: list[UUID]) -> None: + """Ids that are not indexed are ignored.""" diff --git a/src/contexts/assistant/domain/ports/document_search.py b/src/contexts/assistant/domain/ports/document_search.py new file mode 100644 index 0000000..6d93f5b --- /dev/null +++ b/src/contexts/assistant/domain/ports/document_search.py @@ -0,0 +1,11 @@ +from abc import ABC, abstractmethod + +from contexts.assistant.domain.retrieved_document import RetrievedDocument + + +class DocumentSearch(ABC): + @abstractmethod + async def search(self, embedding: list[float], k: int) -> list[RetrievedDocument]: + """Returns at most ``k`` documents, most similar first. ``score`` is the cosine + similarity (higher is closer). An empty index returns an empty list. + """ diff --git a/src/contexts/assistant/domain/ports/embedding_client.py b/src/contexts/assistant/domain/ports/embedding_client.py new file mode 100644 index 0000000..2249b1b --- /dev/null +++ b/src/contexts/assistant/domain/ports/embedding_client.py @@ -0,0 +1,14 @@ +from abc import ABC, abstractmethod + + +class EmbeddingClient(ABC): + @property + @abstractmethod + def model_name(self) -> str: + """Identifies the model. Vectors from different models are not comparable.""" + + @abstractmethod + async def embed(self, texts: list[str]) -> list[list[float]]: + """Returns exactly one embedding per text, in the same order as ``texts``. Every + embedding has the same dimension, fixed by the configured model. + """ diff --git a/src/contexts/assistant/domain/product_document.py b/src/contexts/assistant/domain/product_document.py new file mode 100644 index 0000000..074e62b --- /dev/null +++ b/src/contexts/assistant/domain/product_document.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +import hashlib +from dataclasses import dataclass +from uuid import UUID + +from contexts.assistant.domain.catalog_product import CatalogProduct + + +@dataclass(frozen=True, slots=True) +class ProductDocument: + product_id: UUID + content: str + content_hash: str + + @classmethod + def from_catalog_product( + cls, product: CatalogProduct, *, embedding_model: str + ) -> ProductDocument: + content = _render(product) + fingerprint = f"{embedding_model}\n{content}" + return cls( + product_id=product.product_id, + content=content, + content_hash=hashlib.sha256(fingerprint.encode("utf-8")).hexdigest(), + ) + + +@dataclass(frozen=True, slots=True) +class EmbeddedDocument: + document: ProductDocument + embedding: list[float] + + +def _render(product: CatalogProduct) -> str: + units, cents = divmod(product.price_cents, 100) + return ( + f"SKU: {product.sku}\n" + f"Name: {product.name}\n" + f"Price: {units}.{cents:02d} {product.currency}" + ) diff --git a/src/contexts/assistant/domain/retrieved_document.py b/src/contexts/assistant/domain/retrieved_document.py new file mode 100644 index 0000000..ffd28f2 --- /dev/null +++ b/src/contexts/assistant/domain/retrieved_document.py @@ -0,0 +1,9 @@ +from dataclasses import dataclass +from uuid import UUID + + +@dataclass(frozen=True, slots=True) +class RetrievedDocument: + product_id: UUID + content: str + score: float diff --git a/src/contexts/assistant/infrastructure/catalog/__init__.py b/src/contexts/assistant/infrastructure/catalog/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/contexts/assistant/infrastructure/catalog/list_products_catalog_reader.py b/src/contexts/assistant/infrastructure/catalog/list_products_catalog_reader.py new file mode 100644 index 0000000..3ee54a5 --- /dev/null +++ b/src/contexts/assistant/infrastructure/catalog/list_products_catalog_reader.py @@ -0,0 +1,55 @@ +from uuid import UUID + +from contexts.assistant.domain.catalog_product import CatalogProduct +from contexts.assistant.domain.ports.catalog_reader import CatalogReader +from contexts.catalog.application.dtos import ListProductsQuery, ProductDTO +from contexts.catalog.application.use_cases.get_product import GetProduct +from contexts.catalog.application.use_cases.list_products import ListProducts +from shared.domain.errors import NotFoundError + +DEFAULT_PAGE_SIZE = 200 + + +class ListProductsCatalogReader(CatalogReader): + def __init__( + self, + list_products: ListProducts, + get_product: GetProduct, + page_size: int = DEFAULT_PAGE_SIZE, + ) -> None: + self._list_products = list_products + self._get_product = get_product + self._page_size = page_size + + async def read_all(self) -> list[CatalogProduct]: + products: dict[UUID, CatalogProduct] = {} + offset = 0 + while True: + page = await self._list_products.execute( + ListProductsQuery(limit=self._page_size, offset=offset) + ) + for dto in page: + products[dto.id] = _to_catalog_product(dto) + if len(page) < self._page_size: + return list(products.values()) + offset += self._page_size + + async def missing(self, product_ids: list[UUID]) -> list[UUID]: + return [product_id for product_id in product_ids if not await self._exists(product_id)] + + async def _exists(self, product_id: UUID) -> bool: + try: + await self._get_product.execute(product_id) + except NotFoundError: + return False + return True + + +def _to_catalog_product(dto: ProductDTO) -> CatalogProduct: + return CatalogProduct( + product_id=dto.id, + sku=dto.sku, + name=dto.name, + price_cents=dto.price_cents, + currency=dto.currency, + ) diff --git a/src/contexts/assistant/infrastructure/embeddings/__init__.py b/src/contexts/assistant/infrastructure/embeddings/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/contexts/assistant/infrastructure/embeddings/fastembed_embedding_client.py b/src/contexts/assistant/infrastructure/embeddings/fastembed_embedding_client.py new file mode 100644 index 0000000..2fa01c0 --- /dev/null +++ b/src/contexts/assistant/infrastructure/embeddings/fastembed_embedding_client.py @@ -0,0 +1,36 @@ +import asyncio +import threading + +from fastembed import TextEmbedding + +from contexts.assistant.domain.ports.embedding_client import EmbeddingClient + + +class FastEmbedEmbeddingClient(EmbeddingClient): + def __init__(self, model_name: str, cache_dir: str | None, local_files_only: bool) -> None: + self._model_name = model_name + self._cache_dir = cache_dir + self._local_files_only = local_files_only + self._model: TextEmbedding | None = None + self._lock = threading.Lock() + + @property + def model_name(self) -> str: + return self._model_name + + async def embed(self, texts: list[str]) -> list[list[float]]: + return await asyncio.to_thread(self._embed_blocking, texts) + + def _embed_blocking(self, texts: list[str]) -> list[list[float]]: + model = self._loaded_model() + return [[float(value) for value in vector] for vector in model.embed(texts)] + + def _loaded_model(self) -> TextEmbedding: + with self._lock: + if self._model is None: + self._model = TextEmbedding( + model_name=self._model_name, + cache_dir=self._cache_dir, + local_files_only=self._local_files_only, + ) + return self._model diff --git a/src/contexts/assistant/infrastructure/persistence/__init__.py b/src/contexts/assistant/infrastructure/persistence/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/contexts/assistant/infrastructure/persistence/models.py b/src/contexts/assistant/infrastructure/persistence/models.py new file mode 100644 index 0000000..a8a112e --- /dev/null +++ b/src/contexts/assistant/infrastructure/persistence/models.py @@ -0,0 +1,29 @@ +from datetime import datetime +from uuid import UUID + +from pgvector.sqlalchemy import Vector +from sqlalchemy import DateTime, Index, String, Text +from sqlalchemy.orm import Mapped, mapped_column + +from shared.infrastructure.db_base import Base + +EMBEDDING_DIMENSIONS = 384 + + +class ProductDocumentModel(Base): + __tablename__ = "product_documents" + __table_args__ = ( + Index( + "ix_product_documents_embedding", + "embedding", + postgresql_using="hnsw", + postgresql_with={"m": 16, "ef_construction": 64}, + postgresql_ops={"embedding": "vector_cosine_ops"}, + ), + ) + + product_id: Mapped[UUID] = mapped_column(primary_key=True) + content: Mapped[str] = mapped_column(Text, nullable=False) + content_hash: Mapped[str] = mapped_column(String(64), nullable=False) + embedding: Mapped[list[float]] = mapped_column(Vector(EMBEDDING_DIMENSIONS), nullable=False) + updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False) diff --git a/src/contexts/assistant/infrastructure/persistence/sqlalchemy_document_index.py b/src/contexts/assistant/infrastructure/persistence/sqlalchemy_document_index.py new file mode 100644 index 0000000..3176fcd --- /dev/null +++ b/src/contexts/assistant/infrastructure/persistence/sqlalchemy_document_index.py @@ -0,0 +1,55 @@ +from datetime import datetime, timezone +from uuid import UUID + +from sqlalchemy import delete, select +from sqlalchemy.dialects.postgresql import insert +from sqlalchemy.ext.asyncio import AsyncSession + +from contexts.assistant.domain.ports.document_index import DocumentIndex +from contexts.assistant.domain.product_document import EmbeddedDocument +from contexts.assistant.infrastructure.persistence.models import ProductDocumentModel + + +class SqlAlchemyDocumentIndex(DocumentIndex): + def __init__(self, session: AsyncSession) -> None: + self._session = session + + async def content_hashes(self) -> dict[UUID, str]: + result = await self._session.execute( + select(ProductDocumentModel.product_id, ProductDocumentModel.content_hash) + ) + return dict(result.all()) + + async def upsert(self, documents: list[EmbeddedDocument]) -> None: + if not documents: + return + now = datetime.now(timezone.utc) + stmt = insert(ProductDocumentModel).values( + [ + { + "product_id": item.document.product_id, + "content": item.document.content, + "content_hash": item.document.content_hash, + "embedding": item.embedding, + "updated_at": now, + } + for item in documents + ] + ) + stmt = stmt.on_conflict_do_update( + index_elements=[ProductDocumentModel.product_id], + set_={ + "content": stmt.excluded.content, + "content_hash": stmt.excluded.content_hash, + "embedding": stmt.excluded.embedding, + "updated_at": stmt.excluded.updated_at, + }, + ) + await self._session.execute(stmt) + + async def delete(self, product_ids: list[UUID]) -> None: + if not product_ids: + return + await self._session.execute( + delete(ProductDocumentModel).where(ProductDocumentModel.product_id.in_(product_ids)) + ) diff --git a/src/contexts/assistant/infrastructure/persistence/sqlalchemy_document_search.py b/src/contexts/assistant/infrastructure/persistence/sqlalchemy_document_search.py new file mode 100644 index 0000000..5f29dd3 --- /dev/null +++ b/src/contexts/assistant/infrastructure/persistence/sqlalchemy_document_search.py @@ -0,0 +1,25 @@ +from sqlalchemy import Float, select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from contexts.assistant.domain.ports.document_search import DocumentSearch +from contexts.assistant.domain.retrieved_document import RetrievedDocument +from contexts.assistant.infrastructure.persistence.models import ProductDocumentModel + + +class SqlAlchemyDocumentSearch(DocumentSearch): + def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None: + self._session_factory = session_factory + + async def search(self, embedding: list[float], k: int) -> list[RetrievedDocument]: + distance = ProductDocumentModel.embedding.op("<=>", return_type=Float)(embedding) + stmt = ( + select(ProductDocumentModel.product_id, ProductDocumentModel.content, distance) + .order_by(distance) + .limit(k) + ) + async with self._session_factory() as session: + rows = (await session.execute(stmt)).all() + return [ + RetrievedDocument(product_id=product_id, content=content, score=1.0 - float(dist)) + for product_id, content, dist in rows + ] diff --git a/src/contexts/assistant/ui/cli/__init__.py b/src/contexts/assistant/ui/cli/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/contexts/assistant/ui/cli/reindex.py b/src/contexts/assistant/ui/cli/reindex.py new file mode 100644 index 0000000..b9a6ad4 --- /dev/null +++ b/src/contexts/assistant/ui/cli/reindex.py @@ -0,0 +1,41 @@ +import asyncio +from contextlib import asynccontextmanager + +from contexts.assistant.application.dtos import IndexReportDTO +from contexts.assistant.ui.dependencies import ( + get_catalog_reader, + get_document_index, + get_embedding_client, + get_index_catalog_use_case, +) +from contexts.catalog.ui.dependencies import ( + get_get_product_use_case, + get_list_products_use_case, + get_product_repository, +) +from shared.infrastructure.database import engine, get_session + + +async def reindex() -> IndexReportDTO: + async with asynccontextmanager(get_session)() as session: + use_case = get_index_catalog_use_case( + catalog_reader=get_catalog_reader( + get_list_products_use_case(get_product_repository(session)), + get_get_product_use_case(get_product_repository(session)), + ), + embedding_client=get_embedding_client(), + document_index=get_document_index(session), + ) + return await use_case.execute() + + +async def _main() -> None: + try: + report = await reindex() + finally: + await engine.dispose() + print(f"indexed={report.indexed} unchanged={report.unchanged} deleted={report.deleted}") + + +if __name__ == "__main__": + asyncio.run(_main()) diff --git a/src/contexts/assistant/ui/dependencies.py b/src/contexts/assistant/ui/dependencies.py index e692933..a15c6f8 100644 --- a/src/contexts/assistant/ui/dependencies.py +++ b/src/contexts/assistant/ui/dependencies.py @@ -3,10 +3,35 @@ from anthropic import AsyncAnthropic from fastapi import Depends +from sqlalchemy.ext.asyncio import AsyncSession from contexts.assistant.application.use_cases.ask_assistant import AskAssistant +from contexts.assistant.application.use_cases.index_catalog import IndexCatalog +from contexts.assistant.domain.ports.catalog_reader import CatalogReader +from contexts.assistant.domain.ports.document_index import DocumentIndex +from contexts.assistant.domain.ports.document_search import DocumentSearch +from contexts.assistant.domain.ports.embedding_client import EmbeddingClient from contexts.assistant.domain.ports.llm_client import LLMClient +from contexts.assistant.infrastructure.catalog.list_products_catalog_reader import ( + ListProductsCatalogReader, +) +from contexts.assistant.infrastructure.embeddings.fastembed_embedding_client import ( + FastEmbedEmbeddingClient, +) from contexts.assistant.infrastructure.llm.anthropic_llm_client import AnthropicLLMClient +from contexts.assistant.infrastructure.persistence.sqlalchemy_document_index import ( + SqlAlchemyDocumentIndex, +) +from contexts.assistant.infrastructure.persistence.sqlalchemy_document_search import ( + SqlAlchemyDocumentSearch, +) +from contexts.catalog.application.use_cases.get_product import GetProduct +from contexts.catalog.application.use_cases.list_products import ListProducts +from contexts.catalog.ui.dependencies import ( + get_get_product_use_case, + get_list_products_use_case, +) +from shared.infrastructure.database import get_session, session_factory from shared.infrastructure.settings import get_settings @@ -22,7 +47,51 @@ def get_llm_client() -> LLMClient: ) +@lru_cache +def get_embedding_client() -> EmbeddingClient: + settings = get_settings() + return FastEmbedEmbeddingClient( + model_name=settings.embedding_model, + cache_dir=settings.embedding_cache_dir, + local_files_only=settings.embedding_local_files_only, + ) + + +def get_document_search() -> DocumentSearch: + return SqlAlchemyDocumentSearch(session_factory) + + +def get_document_index( + session: Annotated[AsyncSession, Depends(get_session)], +) -> DocumentIndex: + return SqlAlchemyDocumentIndex(session) + + +def get_catalog_reader( + list_products: Annotated[ListProducts, Depends(get_list_products_use_case)], + get_product: Annotated[GetProduct, Depends(get_get_product_use_case)], +) -> CatalogReader: + return ListProductsCatalogReader(list_products, get_product) + + def get_ask_assistant_use_case( llm_client: Annotated[LLMClient, Depends(get_llm_client)], + embedding_client: Annotated[EmbeddingClient, Depends(get_embedding_client)], + document_search: Annotated[DocumentSearch, Depends(get_document_search)], ) -> AskAssistant: - return AskAssistant(llm_client, max_tokens=get_settings().llm_max_tokens) + settings = get_settings() + return AskAssistant( + llm_client, + embedding_client, + document_search, + max_tokens=settings.llm_max_tokens, + top_k=settings.rag_top_k, + ) + + +def get_index_catalog_use_case( + catalog_reader: Annotated[CatalogReader, Depends(get_catalog_reader)], + embedding_client: Annotated[EmbeddingClient, Depends(get_embedding_client)], + document_index: Annotated[DocumentIndex, Depends(get_document_index)], +) -> IndexCatalog: + return IndexCatalog(catalog_reader, embedding_client, document_index) diff --git a/src/contexts/assistant/ui/lifespan.py b/src/contexts/assistant/ui/lifespan.py new file mode 100644 index 0000000..0dff8a5 --- /dev/null +++ b/src/contexts/assistant/ui/lifespan.py @@ -0,0 +1,7 @@ +from contexts.assistant.ui.dependencies import get_embedding_client + +WARM_UP_TEXT = "warm-up" + + +async def warm_up_embeddings() -> None: + await get_embedding_client().embed([WARM_UP_TEXT]) diff --git a/src/contexts/assistant/ui/routers/assistant.py b/src/contexts/assistant/ui/routers/assistant.py index b0a0fea..bdd2186 100644 --- a/src/contexts/assistant/ui/routers/assistant.py +++ b/src/contexts/assistant/ui/routers/assistant.py @@ -15,11 +15,12 @@ response_model=AnswerResponse, status_code=status.HTTP_200_OK, summary="Ask the assistant", - description="Answers a natural-language question using the configured LLM.", - responses={ - 422: {"description": "Invalid question."}, - 503: {"description": "The language model is unavailable."}, - }, + description=( + "Retrieves the catalog products most similar to the question and answers with the " + "configured LLM, returning them as `sources`. If the LLM is unavailable, answers 200 " + "with the retrieved products and `fallback_used=true`." + ), + responses={422: {"description": "Invalid question."}}, ) async def ask_endpoint( body: AskRequest, diff --git a/src/contexts/assistant/ui/schemas.py b/src/contexts/assistant/ui/schemas.py index 85f5baf..d77f88f 100644 --- a/src/contexts/assistant/ui/schemas.py +++ b/src/contexts/assistant/ui/schemas.py @@ -1,8 +1,10 @@ from __future__ import annotations +from uuid import UUID + from pydantic import BaseModel, ConfigDict, Field -from contexts.assistant.application.dtos import AnswerDTO +from contexts.assistant.application.dtos import AnswerDTO, SourceDTO, UsageDTO class AskRequest(BaseModel): @@ -23,19 +25,45 @@ class UsageResponse(BaseModel): output_tokens: int = Field(..., description="Output tokens billed for the completion.") latency_ms: int = Field(..., description="Wall-clock latency of the LLM call in ms.") + @classmethod + def from_dto(cls, dto: UsageDTO) -> UsageResponse: + return cls( + model=dto.model, + input_tokens=dto.input_tokens, + output_tokens=dto.output_tokens, + latency_ms=dto.latency_ms, + ) + + +class SourceResponse(BaseModel): + product_id: UUID = Field(..., description="Catalog product the document was built from.") + content: str = Field(..., description="Indexed text of the product.") + score: float = Field(..., description="Cosine similarity to the question (higher is closer).") + + @classmethod + def from_dto(cls, dto: SourceDTO) -> SourceResponse: + return cls(product_id=dto.product_id, content=dto.content, score=dto.score) + class AnswerResponse(BaseModel): - text: str = Field(..., description="Answer generated by the assistant.") - usage: UsageResponse + text: str = Field(..., description="Answer generated by the assistant, or the fallback text.") + sources: list[SourceResponse] = Field( + ..., description="Catalog products retrieved for the question, most similar first." + ) + fallback_used: bool = Field( + ..., + description="True when the LLM was unavailable and the answer lists the retrieved " + "products without generation.", + ) + usage: UsageResponse | None = Field( + ..., description="LLM usage. Null when the fallback was used (no LLM call succeeded)." + ) @classmethod def from_dto(cls, dto: AnswerDTO) -> AnswerResponse: return cls( text=dto.text, - usage=UsageResponse( - model=dto.usage.model, - input_tokens=dto.usage.input_tokens, - output_tokens=dto.usage.output_tokens, - latency_ms=dto.usage.latency_ms, - ), + sources=[SourceResponse.from_dto(source) for source in dto.sources], + fallback_used=dto.fallback_used, + usage=UsageResponse.from_dto(dto.usage) if dto.usage is not None else None, ) diff --git a/src/shared/infrastructure/settings.py b/src/shared/infrastructure/settings.py index 5068a22..f401a63 100644 --- a/src/shared/infrastructure/settings.py +++ b/src/shared/infrastructure/settings.py @@ -50,6 +50,27 @@ class Settings(BaseSettings): description="Maximum output tokens per LLM completion.", ) + embedding_model: str = Field( + default="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2", + description="fastembed model used to embed catalog documents and questions.", + ) + embedding_cache_dir: str | None = Field( + default=None, + description="Directory where fastembed caches the model. Unset: fastembed's default.", + ) + embedding_local_files_only: bool = Field( + default=False, + description="Only load the model from the cache, never download it. " + "The Docker image bakes the model in and sets this to true.", + ) + rag_top_k: int = Field( + default=5, + gt=0, + le=40, + description="Number of catalog documents retrieved per question. At most 40: the " + "HNSW index returns up to hnsw.ef_search (default 40) candidates per query.", + ) + @lru_cache def get_settings() -> Settings: diff --git a/tests/e2e/test_assistant.py b/tests/e2e/test_assistant.py index 1fc425d..135ee48 100644 --- a/tests/e2e/test_assistant.py +++ b/tests/e2e/test_assistant.py @@ -1,23 +1,42 @@ from collections.abc import AsyncIterator, Callable +from uuid import uuid4 import pytest_asyncio from httpx import ASGITransport, AsyncClient from main import create_app +from tests.fakes.documents import FakeDocumentSearch +from tests.fakes.embedding_client import FakeEmbeddingClient from tests.fakes.llm_client import FakeLLMClient +from contexts.assistant.application.use_cases.ask_assistant import FALLBACK_WITH_SOURCES from contexts.assistant.domain.completion import Completion -from contexts.assistant.ui.dependencies import get_llm_client +from contexts.assistant.domain.retrieved_document import RetrievedDocument +from contexts.assistant.ui.dependencies import ( + get_document_search, + get_embedding_client, + get_llm_client, +) -ClientFactory = Callable[[FakeLLMClient], AsyncClient] +WIDGET = RetrievedDocument( + product_id=uuid4(), content="SKU: W-1\nName: Blue widget\nPrice: 19.99 EUR", score=0.9 +) + +ClientFactory = Callable[..., AsyncClient] @pytest_asyncio.fixture async def client_with() -> AsyncIterator[ClientFactory]: clients: list[AsyncClient] = [] - def _factory(llm_client: FakeLLMClient) -> AsyncClient: + def _factory( + llm_client: FakeLLMClient, search: FakeDocumentSearch | None = None + ) -> AsyncClient: application = create_app() + document_search = search or FakeDocumentSearch([WIDGET]) + embedding_client = FakeEmbeddingClient() application.dependency_overrides[get_llm_client] = lambda: llm_client + application.dependency_overrides[get_embedding_client] = lambda: embedding_client + application.dependency_overrides[get_document_search] = lambda: document_search client = AsyncClient(transport=ASGITransport(app=application), base_url="http://testserver") clients.append(client) return client @@ -27,30 +46,52 @@ def _factory(llm_client: FakeLLMClient) -> AsyncClient: await client.aclose() +def _source_json(document: RetrievedDocument) -> dict[str, object]: + return { + "product_id": str(document.product_id), + "content": document.content, + "score": document.score, + } + + class TestAsk: - async def test_returns_200_with_answer_and_usage(self, client_with: ClientFactory) -> None: + async def test_c7_returns_200_with_answer_sources_and_usage( + self, client_with: ClientFactory + ) -> None: completion = Completion( - text="Widgets are small devices.", + text="The blue widget costs 19.99 EUR.", model="claude-test", input_tokens=50, - output_tokens=6, + output_tokens=9, latency_ms=120, ) client = client_with(FakeLLMClient(completion)) - response = await client.post("/assistant/ask", json={"question": "What is a widget?"}) + response = await client.post("/assistant/ask", json={"question": "How much is a widget?"}) assert response.status_code == 200 assert response.json() == { - "text": "Widgets are small devices.", + "text": "The blue widget costs 19.99 EUR.", + "sources": [_source_json(WIDGET)], + "fallback_used": False, "usage": { "model": "claude-test", "input_tokens": 50, - "output_tokens": 6, + "output_tokens": 9, "latency_ms": 120, }, } + async def test_returns_empty_sources_when_index_is_empty( + self, client_with: ClientFactory + ) -> None: + client = client_with(FakeLLMClient(), FakeDocumentSearch()) + + response = await client.post("/assistant/ask", json={"question": "What is a widget?"}) + + assert response.status_code == 200 + assert response.json()["sources"] == [] + async def test_returns_422_on_empty_question(self, client_with: ClientFactory) -> None: client = client_with(FakeLLMClient()) @@ -77,10 +118,23 @@ async def test_returns_422_on_unknown_field(self, client_with: ClientFactory) -> assert response.status_code == 422 - async def test_returns_503_when_llm_unavailable(self, client_with: ClientFactory) -> None: + async def test_c8_returns_200_with_fallback_when_llm_unavailable( + self, client_with: ClientFactory + ) -> None: client = client_with(FakeLLMClient(fail=True)) response = await client.post("/assistant/ask", json={"question": "What is a widget?"}) - assert response.status_code == 503 - assert response.json() == {"detail": "the language model is unavailable"} + assert response.status_code == 200 + body = response.json() + assert body["fallback_used"] is True + assert body["usage"] is None + assert body["sources"] == [_source_json(WIDGET)] + assert body["text"].startswith(FALLBACK_WITH_SOURCES) + + +class TestOpenApi: + def test_c9_assistant_exposes_no_indexing_endpoint(self) -> None: + paths = create_app().openapi()["paths"] + + assert [path for path in paths if path.startswith("/assistant")] == ["/assistant/ask"] diff --git a/tests/e2e/test_lifespan.py b/tests/e2e/test_lifespan.py new file mode 100644 index 0000000..02d2427 --- /dev/null +++ b/tests/e2e/test_lifespan.py @@ -0,0 +1,23 @@ +import pytest +from fastapi.testclient import TestClient +from main import create_app +from tests.fakes.embedding_client import FakeEmbeddingClient + +from contexts.assistant.ui import lifespan + + +@pytest.fixture +def embeddings(monkeypatch: pytest.MonkeyPatch) -> FakeEmbeddingClient: + fake = FakeEmbeddingClient() + monkeypatch.setattr(lifespan, "get_embedding_client", lambda: fake) + return fake + + +def test_r6_app_loads_the_embedding_model_on_startup(embeddings: FakeEmbeddingClient) -> None: + with TestClient(create_app(warm_up_models=True)): + assert embeddings.calls == [[lifespan.WARM_UP_TEXT]] + + +def test_r6_warm_up_is_off_unless_requested(embeddings: FakeEmbeddingClient) -> None: + with TestClient(create_app()): + assert embeddings.calls == [] diff --git a/tests/fakes/documents.py b/tests/fakes/documents.py new file mode 100644 index 0000000..5e76443 --- /dev/null +++ b/tests/fakes/documents.py @@ -0,0 +1,53 @@ +from uuid import UUID + +from contexts.assistant.domain.catalog_product import CatalogProduct +from contexts.assistant.domain.ports.catalog_reader import CatalogReader +from contexts.assistant.domain.ports.document_index import DocumentIndex +from contexts.assistant.domain.ports.document_search import DocumentSearch +from contexts.assistant.domain.product_document import EmbeddedDocument +from contexts.assistant.domain.retrieved_document import RetrievedDocument + + +class FakeCatalogReader(CatalogReader): + def __init__( + self, + products: list[CatalogProduct] | None = None, + unlisted: list[CatalogProduct] | None = None, + ) -> None: + self.products = list(products or []) + self.unlisted = list(unlisted or []) + + async def read_all(self) -> list[CatalogProduct]: + return list(self.products) + + async def missing(self, product_ids: list[UUID]) -> list[UUID]: + existing = {p.product_id for p in self.products + self.unlisted} + return [pid for pid in product_ids if pid not in existing] + + +class FakeDocumentIndex(DocumentIndex): + def __init__(self) -> None: + self.documents: dict[UUID, EmbeddedDocument] = {} + self.upsert_sizes: list[int] = [] + + async def content_hashes(self) -> dict[UUID, str]: + return {pid: item.document.content_hash for pid, item in self.documents.items()} + + async def upsert(self, documents: list[EmbeddedDocument]) -> None: + self.upsert_sizes.append(len(documents)) + for item in documents: + self.documents[item.document.product_id] = item + + async def delete(self, product_ids: list[UUID]) -> None: + for product_id in product_ids: + self.documents.pop(product_id, None) + + +class FakeDocumentSearch(DocumentSearch): + def __init__(self, results: list[RetrievedDocument] | None = None) -> None: + self._results = list(results or []) + self.calls: list[tuple[list[float], int]] = [] + + async def search(self, embedding: list[float], k: int) -> list[RetrievedDocument]: + self.calls.append((embedding, k)) + return self._results[:k] diff --git a/tests/fakes/embedding_client.py b/tests/fakes/embedding_client.py new file mode 100644 index 0000000..fff7bdf --- /dev/null +++ b/tests/fakes/embedding_client.py @@ -0,0 +1,20 @@ +from contexts.assistant.domain.ports.embedding_client import EmbeddingClient + +DEFAULT_VECTOR = [1.0, 0.0, 0.0] + + +class FakeEmbeddingClient(EmbeddingClient): + def __init__( + self, vectors: dict[str, list[float]] | None = None, model_name: str = "fake-model" + ) -> None: + self._vectors = vectors or {} + self.current_model_name = model_name + self.calls: list[list[str]] = [] + + @property + def model_name(self) -> str: + return self.current_model_name + + async def embed(self, texts: list[str]) -> list[list[float]]: + self.calls.append(list(texts)) + return [self._vectors.get(text, DEFAULT_VECTOR) for text in texts] diff --git a/tests/integration/contexts/assistant/application/test_ask_assistant_connection.py b/tests/integration/contexts/assistant/application/test_ask_assistant_connection.py new file mode 100644 index 0000000..4f58263 --- /dev/null +++ b/tests/integration/contexts/assistant/application/test_ask_assistant_connection.py @@ -0,0 +1,42 @@ +from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker +from tests.fakes.embedding_client import FakeEmbeddingClient +from tests.fakes.llm_client import FakeLLMClient + +from contexts.assistant.application.dtos import AskCommand +from contexts.assistant.application.use_cases.ask_assistant import AskAssistant +from contexts.assistant.domain.completion import Completion +from contexts.assistant.infrastructure.persistence.models import EMBEDDING_DIMENSIONS +from contexts.assistant.infrastructure.persistence.sqlalchemy_document_search import ( + SqlAlchemyDocumentSearch, +) + +QUESTION = "What is a widget?" + + +class PoolProbingLLMClient(FakeLLMClient): + def __init__(self, engine: AsyncEngine) -> None: + super().__init__() + self._engine = engine + self.checked_out_during_call: list[int] = [] + + async def complete(self, *, system: str, prompt: str, max_tokens: int) -> Completion: + self.checked_out_during_call.append(self._engine.pool.checkedout()) + return await super().complete(system=system, prompt=prompt, max_tokens=max_tokens) + + +async def test_r5_no_connection_is_held_while_the_llm_answers(db_engine: AsyncEngine) -> None: + llm_client = PoolProbingLLMClient(db_engine) + embeddings = FakeEmbeddingClient({QUESTION: [1.0] + [0.0] * (EMBEDDING_DIMENSIONS - 1)}) + session_factory = async_sessionmaker(db_engine, expire_on_commit=False) + + use_case = AskAssistant( + llm_client, + embeddings, + SqlAlchemyDocumentSearch(session_factory), + max_tokens=16, + top_k=5, + ) + + await use_case.execute(AskCommand(question=QUESTION)) + + assert llm_client.checked_out_during_call == [0] diff --git a/tests/integration/contexts/assistant/infrastructure/test_fastembed_embedding_client.py b/tests/integration/contexts/assistant/infrastructure/test_fastembed_embedding_client.py new file mode 100644 index 0000000..90c0c19 --- /dev/null +++ b/tests/integration/contexts/assistant/infrastructure/test_fastembed_embedding_client.py @@ -0,0 +1,23 @@ +from pathlib import Path + +import pytest + +from contexts.assistant.infrastructure.embeddings.fastembed_embedding_client import ( + FastEmbedEmbeddingClient, +) + +MODEL = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2" + + +class TestLocalFilesOnly: + async def test_r7_fails_instead_of_downloading_when_the_model_is_not_cached( + self, tmp_path: Path + ) -> None: + client = FastEmbedEmbeddingClient( + model_name=MODEL, cache_dir=str(tmp_path), local_files_only=True + ) + + with pytest.raises(ValueError, match="Could not load model"): + await client.embed(["hola"]) + + assert not any(tmp_path.rglob("*.onnx")) diff --git a/tests/integration/contexts/assistant/infrastructure/test_list_products_catalog_reader.py b/tests/integration/contexts/assistant/infrastructure/test_list_products_catalog_reader.py new file mode 100644 index 0000000..cc57b35 --- /dev/null +++ b/tests/integration/contexts/assistant/infrastructure/test_list_products_catalog_reader.py @@ -0,0 +1,75 @@ +from collections.abc import Callable +from uuid import uuid4 + +from contexts.assistant.domain.catalog_product import CatalogProduct +from contexts.assistant.infrastructure.catalog.list_products_catalog_reader import ( + ListProductsCatalogReader, +) +from contexts.catalog.application.use_cases.get_product import GetProduct +from contexts.catalog.application.use_cases.list_products import ListProducts +from contexts.catalog.domain.ports.product_repository import ProductRepository +from contexts.catalog.domain.product import Product + + +class TestReadAll: + async def test_reads_every_product_across_pages( + self, + product_repository: ProductRepository, + make_product: Callable[..., Product], + ) -> None: + products = [make_product() for _ in range(5)] + for product in products: + await product_repository.add(product) + reader = ListProductsCatalogReader( + ListProducts(product_repository), GetProduct(product_repository), page_size=2 + ) + + result = await reader.read_all() + + assert sorted(p.product_id for p in result) == sorted(p.id for p in products) + + async def test_maps_product_fields( + self, + product_repository: ProductRepository, + make_product: Callable[..., Product], + ) -> None: + product = make_product(sku="SKU-MAP", name="Mapped", price_cents=1234, currency="USD") + await product_repository.add(product) + reader = ListProductsCatalogReader( + ListProducts(product_repository), GetProduct(product_repository) + ) + + [result] = await reader.read_all() + + assert result == CatalogProduct( + product_id=product.id, + sku="SKU-MAP", + name="Mapped", + price_cents=1234, + currency="USD", + ) + + async def test_empty_catalog_returns_no_products( + self, product_repository: ProductRepository + ) -> None: + reader = ListProductsCatalogReader( + ListProducts(product_repository), GetProduct(product_repository) + ) + + assert await reader.read_all() == [] + + +class TestMissing: + async def test_r1_reports_only_ids_not_in_the_catalog( + self, + product_repository: ProductRepository, + make_product: Callable[..., Product], + ) -> None: + existing = make_product() + await product_repository.add(existing) + gone = uuid4() + reader = ListProductsCatalogReader( + ListProducts(product_repository), GetProduct(product_repository) + ) + + assert await reader.missing([existing.id, gone]) == [gone] diff --git a/tests/integration/contexts/assistant/infrastructure/test_sqlalchemy_documents.py b/tests/integration/contexts/assistant/infrastructure/test_sqlalchemy_documents.py new file mode 100644 index 0000000..1679222 --- /dev/null +++ b/tests/integration/contexts/assistant/infrastructure/test_sqlalchemy_documents.py @@ -0,0 +1,121 @@ +from uuid import UUID, uuid4 + +import pytest +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from contexts.assistant.domain.product_document import EmbeddedDocument, ProductDocument +from contexts.assistant.infrastructure.persistence.models import EMBEDDING_DIMENSIONS +from contexts.assistant.infrastructure.persistence.sqlalchemy_document_index import ( + SqlAlchemyDocumentIndex, +) +from contexts.assistant.infrastructure.persistence.sqlalchemy_document_search import ( + SqlAlchemyDocumentSearch, +) + + +def _vector(*weights: tuple[int, float]) -> list[float]: + vector = [0.0] * EMBEDDING_DIMENSIONS + for position, value in weights: + vector[position] = value + return vector + + +def _embedded( + content: str, vector: list[float], product_id: UUID | None = None +) -> EmbeddedDocument: + return EmbeddedDocument( + document=ProductDocument( + product_id=product_id or uuid4(), content=content, content_hash=f"hash-{content}" + ), + embedding=vector, + ) + + +@pytest.fixture +def index(db_session: AsyncSession) -> SqlAlchemyDocumentIndex: + return SqlAlchemyDocumentIndex(db_session) + + +@pytest.fixture +def search(db_session: AsyncSession) -> SqlAlchemyDocumentSearch: + factory = async_sessionmaker( + bind=db_session.bind, + expire_on_commit=False, + join_transaction_mode="create_savepoint", + ) + return SqlAlchemyDocumentSearch(factory) + + +class TestDocumentIndex: + async def test_upsert_inserts_and_reports_hashes(self, index: SqlAlchemyDocumentIndex) -> None: + first = _embedded("alpha", _vector((0, 1.0))) + second = _embedded("beta", _vector((1, 1.0))) + + await index.upsert([first, second]) + + assert await index.content_hashes() == { + first.document.product_id: "hash-alpha", + second.document.product_id: "hash-beta", + } + + async def test_upsert_replaces_an_existing_document( + self, index: SqlAlchemyDocumentIndex, search: SqlAlchemyDocumentSearch + ) -> None: + product_id = uuid4() + await index.upsert([_embedded("old", _vector((0, 1.0)), product_id)]) + + await index.upsert([_embedded("new", _vector((1, 1.0)), product_id)]) + + assert await index.content_hashes() == {product_id: "hash-new"} + [hit] = await search.search(_vector((1, 1.0)), k=1) + assert hit.product_id == product_id + assert hit.content == "new" + + async def test_c4_delete_removes_only_the_given_documents( + self, index: SqlAlchemyDocumentIndex + ) -> None: + kept = _embedded("kept", _vector((0, 1.0))) + orphan = _embedded("orphan", _vector((1, 1.0))) + await index.upsert([kept, orphan]) + + await index.delete([orphan.document.product_id]) + + assert await index.content_hashes() == {kept.document.product_id: "hash-kept"} + + async def test_delete_ignores_ids_that_are_not_indexed( + self, index: SqlAlchemyDocumentIndex + ) -> None: + kept = _embedded("kept", _vector((0, 1.0))) + await index.upsert([kept]) + + await index.delete([uuid4()]) + + assert await index.content_hashes() == {kept.document.product_id: "hash-kept"} + + +class TestDocumentSearch: + async def test_c10_returns_the_nearest_document_first( + self, index: SqlAlchemyDocumentIndex, search: SqlAlchemyDocumentSearch + ) -> None: + near = _embedded("near", _vector((0, 1.0), (1, 0.1))) + middle = _embedded("middle", _vector((0, 1.0), (1, 1.0))) + far = _embedded("far", _vector((1, 1.0))) + await index.upsert([far, near, middle]) + + results = await search.search(_vector((0, 1.0)), k=3) + + assert [r.content for r in results] == ["near", "middle", "far"] + assert results[0].score == pytest.approx(0.995, abs=1e-3) + assert results[0].score > results[1].score > results[2].score + + async def test_limits_results_to_k( + self, index: SqlAlchemyDocumentIndex, search: SqlAlchemyDocumentSearch + ) -> None: + await index.upsert( + [_embedded(name, _vector((i, 1.0))) for i, name in enumerate(["a", "b", "c"])] + ) + + assert len(await search.search(_vector((0, 1.0)), k=2)) == 2 + + async def test_empty_index_returns_no_results(self, search: SqlAlchemyDocumentSearch) -> None: + assert await search.search(_vector((0, 1.0)), k=5) == [] diff --git a/tests/unit/contexts/assistant/application/test_ask_assistant.py b/tests/unit/contexts/assistant/application/test_ask_assistant.py index 2333ef4..9d3d7a5 100644 --- a/tests/unit/contexts/assistant/application/test_ask_assistant.py +++ b/tests/unit/contexts/assistant/application/test_ask_assistant.py @@ -1,25 +1,108 @@ +from uuid import uuid4 + import pytest +from tests.fakes.documents import FakeDocumentSearch +from tests.fakes.embedding_client import FakeEmbeddingClient from tests.fakes.llm_client import FakeLLMClient -from contexts.assistant.application.dtos import AnswerDTO, AskCommand, UsageDTO -from contexts.assistant.application.use_cases.ask_assistant import SYSTEM_PROMPT, AskAssistant +from contexts.assistant.application.dtos import AnswerDTO, AskCommand, SourceDTO, UsageDTO +from contexts.assistant.application.use_cases.ask_assistant import ( + FALLBACK_WITH_SOURCES, + FALLBACK_WITHOUT_SOURCES, + SYSTEM_PROMPT, + AskAssistant, +) from contexts.assistant.domain.completion import Completion -from contexts.assistant.domain.exceptions import InvalidQuestionError, LLMUnavailableError +from contexts.assistant.domain.exceptions import InvalidQuestionError +from contexts.assistant.domain.retrieved_document import RetrievedDocument + +QUESTION_VECTOR = [0.0, 1.0, 0.0] +WIDGET = RetrievedDocument( + product_id=uuid4(), content="SKU: W-1\nName: Blue widget\nPrice: 19.99 EUR", score=0.91 +) +GADGET = RetrievedDocument( + product_id=uuid4(), content="SKU: G-1\nName: Gadget\nPrice: 5.00 EUR", score=0.42 +) + + +@pytest.fixture +def embeddings() -> FakeEmbeddingClient: + return FakeEmbeddingClient({"What is a widget?": QUESTION_VECTOR}) + + +@pytest.fixture +def search() -> FakeDocumentSearch: + return FakeDocumentSearch([WIDGET, GADGET]) + + +def _use_case( + llm_client: FakeLLMClient, embeddings: FakeEmbeddingClient, search: FakeDocumentSearch +) -> AskAssistant: + return AskAssistant(llm_client, embeddings, search, max_tokens=256, top_k=5) @pytest.fixture -def use_case(llm_client: FakeLLMClient) -> AskAssistant: - return AskAssistant(llm_client, max_tokens=256) +def use_case( + llm_client: FakeLLMClient, embeddings: FakeEmbeddingClient, search: FakeDocumentSearch +) -> AskAssistant: + return _use_case(llm_client, embeddings, search) + + +class TestRetrieval: + async def test_embeds_the_stripped_question_and_searches_top_k( + self, + use_case: AskAssistant, + embeddings: FakeEmbeddingClient, + search: FakeDocumentSearch, + ) -> None: + await use_case.execute(AskCommand(question=" What is a widget? ")) + assert embeddings.calls == [["What is a widget?"]] + assert search.calls == [(QUESTION_VECTOR, 5)] -class TestExecute: - async def test_returns_completion_text(self, use_case: AskAssistant) -> None: + async def test_c7_returns_retrieved_documents_as_sources(self, use_case: AskAssistant) -> None: + result = await use_case.execute(AskCommand(question="What is a widget?")) + + assert result.sources == ( + SourceDTO.from_retrieved(WIDGET), + SourceDTO.from_retrieved(GADGET), + ) + + async def test_c7_puts_the_documents_and_the_question_in_the_prompt( + self, use_case: AskAssistant, llm_client: FakeLLMClient + ) -> None: + await use_case.execute(AskCommand(question="What is a widget?")) + + request = llm_client.requests[0] + assert request.prompt == ( + f"Catalog products:\n\n{WIDGET.content}\n\n{GADGET.content}\n\n" + "Question: What is a widget?" + ) + assert request.system == SYSTEM_PROMPT + assert request.max_tokens == 256 + + async def test_empty_index_sends_only_the_question( + self, llm_client: FakeLLMClient, embeddings: FakeEmbeddingClient + ) -> None: + use_case = _use_case(llm_client, embeddings, FakeDocumentSearch()) + + result = await use_case.execute(AskCommand(question="What is a widget?")) + + assert llm_client.requests[0].prompt == "What is a widget?" + assert result.sources == () + + +class TestAnswer: + async def test_returns_completion_text_without_fallback(self, use_case: AskAssistant) -> None: result = await use_case.execute(AskCommand(question="What is a widget?")) assert isinstance(result, AnswerDTO) assert result.text == "An answer." + assert result.fallback_used is False - async def test_propagates_usage_to_dto(self) -> None: + async def test_propagates_usage_to_dto( + self, embeddings: FakeEmbeddingClient, search: FakeDocumentSearch + ) -> None: completion = Completion( text="Answer.", model="claude-test", @@ -27,7 +110,7 @@ async def test_propagates_usage_to_dto(self) -> None: output_tokens=20, latency_ms=350, ) - use_case = AskAssistant(FakeLLMClient(completion), max_tokens=256) + use_case = _use_case(FakeLLMClient(completion), embeddings, search) result = await use_case.execute(AskCommand(question="What is a widget?")) @@ -35,27 +118,46 @@ async def test_propagates_usage_to_dto(self) -> None: model="claude-test", input_tokens=100, output_tokens=20, latency_ms=350 ) - async def test_sends_stripped_question_with_system_prompt_and_max_tokens( - self, use_case: AskAssistant, llm_client: FakeLLMClient - ) -> None: - await use_case.execute(AskCommand(question=" What is a widget? ")) - - assert len(llm_client.requests) == 1 - request = llm_client.requests[0] - assert request.prompt == "What is a widget?" - assert request.system == SYSTEM_PROMPT - assert request.max_tokens == 256 - - async def test_rejects_invalid_question_without_calling_llm( - self, use_case: AskAssistant, llm_client: FakeLLMClient + async def test_rejects_invalid_question_without_embedding_or_calling_llm( + self, + use_case: AskAssistant, + llm_client: FakeLLMClient, + embeddings: FakeEmbeddingClient, ) -> None: with pytest.raises(InvalidQuestionError): await use_case.execute(AskCommand(question=" ")) + assert embeddings.calls == [] assert llm_client.requests == [] - async def test_propagates_llm_unavailable(self) -> None: - use_case = AskAssistant(FakeLLMClient(fail=True), max_tokens=256) - with pytest.raises(LLMUnavailableError): - await use_case.execute(AskCommand(question="What is a widget?")) +class TestFallback: + async def test_c8_llm_unavailable_returns_sources_without_generation( + self, embeddings: FakeEmbeddingClient, search: FakeDocumentSearch + ) -> None: + use_case = _use_case(FakeLLMClient(fail=True), embeddings, search) + + result = await use_case.execute(AskCommand(question="What is a widget?")) + + assert result.fallback_used is True + assert result.usage is None + assert result.sources == ( + SourceDTO.from_retrieved(WIDGET), + SourceDTO.from_retrieved(GADGET), + ) + assert result.text == ( + f"{FALLBACK_WITH_SOURCES}\n" + "- SKU: W-1 | Name: Blue widget | Price: 19.99 EUR\n" + "- SKU: G-1 | Name: Gadget | Price: 5.00 EUR" + ) + + async def test_c8_llm_unavailable_with_empty_index_returns_fixed_text( + self, embeddings: FakeEmbeddingClient + ) -> None: + use_case = _use_case(FakeLLMClient(fail=True), embeddings, FakeDocumentSearch()) + + result = await use_case.execute(AskCommand(question="What is a widget?")) + + assert result.fallback_used is True + assert result.sources == () + assert result.text == FALLBACK_WITHOUT_SOURCES diff --git a/tests/unit/contexts/assistant/application/test_index_catalog.py b/tests/unit/contexts/assistant/application/test_index_catalog.py new file mode 100644 index 0000000..d55e506 --- /dev/null +++ b/tests/unit/contexts/assistant/application/test_index_catalog.py @@ -0,0 +1,175 @@ +from dataclasses import replace +from uuid import uuid4 + +import pytest +from tests.fakes.documents import FakeCatalogReader, FakeDocumentIndex +from tests.fakes.embedding_client import FakeEmbeddingClient + +from contexts.assistant.application.dtos import IndexReportDTO +from contexts.assistant.application.use_cases.index_catalog import IndexCatalog +from contexts.assistant.domain.catalog_product import CatalogProduct +from contexts.assistant.domain.product_document import ProductDocument + +MODEL = "fake-model" + + +def _product(name: str) -> CatalogProduct: + return CatalogProduct( + product_id=uuid4(), + sku=f"SKU-{name.upper()}", + name=name, + price_cents=1000, + currency="EUR", + ) + + +@pytest.fixture +def catalog() -> FakeCatalogReader: + return FakeCatalogReader([_product("alpha"), _product("beta"), _product("gamma")]) + + +@pytest.fixture +def embeddings() -> FakeEmbeddingClient: + return FakeEmbeddingClient() + + +@pytest.fixture +def index() -> FakeDocumentIndex: + return FakeDocumentIndex() + + +@pytest.fixture +def use_case( + catalog: FakeCatalogReader, embeddings: FakeEmbeddingClient, index: FakeDocumentIndex +) -> IndexCatalog: + return IndexCatalog(catalog, embeddings, index) + + +class TestExecute: + async def test_c1_indexes_every_product_into_an_empty_index( + self, + use_case: IndexCatalog, + catalog: FakeCatalogReader, + embeddings: FakeEmbeddingClient, + index: FakeDocumentIndex, + ) -> None: + report = await use_case.execute() + + assert report == IndexReportDTO(indexed=3, unchanged=0, deleted=0) + assert set(index.documents) == {p.product_id for p in catalog.products} + expected = [ + ProductDocument.from_catalog_product(p, embedding_model=MODEL).content + for p in catalog.products + ] + assert embeddings.calls == [expected] + + async def test_c2_second_run_without_changes_embeds_nothing( + self, use_case: IndexCatalog, embeddings: FakeEmbeddingClient + ) -> None: + await use_case.execute() + embeddings.calls.clear() + + report = await use_case.execute() + + assert report == IndexReportDTO(indexed=0, unchanged=3, deleted=0) + assert embeddings.calls == [] + + async def test_c3_only_the_changed_product_is_embedded_again( + self, + use_case: IndexCatalog, + catalog: FakeCatalogReader, + embeddings: FakeEmbeddingClient, + index: FakeDocumentIndex, + ) -> None: + await use_case.execute() + embeddings.calls.clear() + renamed = replace(catalog.products[1], name="beta v2") + catalog.products[1] = renamed + + report = await use_case.execute() + + assert report == IndexReportDTO(indexed=1, unchanged=2, deleted=0) + new_content = ProductDocument.from_catalog_product(renamed, embedding_model=MODEL).content + assert embeddings.calls == [[new_content]] + assert index.documents[renamed.product_id].document.content == new_content + + async def test_c4_deletes_documents_of_products_no_longer_in_the_catalog( + self, + use_case: IndexCatalog, + catalog: FakeCatalogReader, + embeddings: FakeEmbeddingClient, + index: FakeDocumentIndex, + ) -> None: + await use_case.execute() + embeddings.calls.clear() + removed = catalog.products.pop(0) + + report = await use_case.execute() + + assert report == IndexReportDTO(indexed=0, unchanged=2, deleted=1) + assert removed.product_id not in index.documents + assert set(index.documents) == {p.product_id for p in catalog.products} + assert embeddings.calls == [] + + async def test_empty_catalog_and_empty_index_do_nothing( + self, embeddings: FakeEmbeddingClient, index: FakeDocumentIndex + ) -> None: + use_case = IndexCatalog(FakeCatalogReader(), embeddings, index) + + report = await use_case.execute() + + assert report == IndexReportDTO(indexed=0, unchanged=0, deleted=0) + assert embeddings.calls == [] + assert index.documents == {} + + +class TestReviewFixes: + async def test_r1_keeps_a_product_that_exists_but_was_not_listed( + self, embeddings: FakeEmbeddingClient, index: FakeDocumentIndex + ) -> None: + listed, skipped = _product("listed"), _product("skipped") + catalog = FakeCatalogReader([listed, skipped]) + await IndexCatalog(catalog, embeddings, index).execute() + catalog.products, catalog.unlisted = [listed], [skipped] + + report = await IndexCatalog(catalog, embeddings, index).execute() + + assert report.deleted == 0 + assert skipped.product_id in index.documents + + async def test_r1_still_deletes_a_product_that_no_longer_exists( + self, embeddings: FakeEmbeddingClient, index: FakeDocumentIndex + ) -> None: + kept, removed = _product("kept"), _product("removed") + catalog = FakeCatalogReader([kept, removed]) + await IndexCatalog(catalog, embeddings, index).execute() + catalog.products = [kept] + + report = await IndexCatalog(catalog, embeddings, index).execute() + + assert report.deleted == 1 + assert removed.product_id not in index.documents + + async def test_r2_changing_the_embedding_model_reembeds_everything( + self, use_case: IndexCatalog, embeddings: FakeEmbeddingClient + ) -> None: + await use_case.execute() + embeddings.calls.clear() + embeddings.current_model_name = "another-model" + + report = await use_case.execute() + + assert report == IndexReportDTO(indexed=3, unchanged=0, deleted=0) + assert len(embeddings.calls[0]) == 3 + + async def test_r3_embeds_and_upserts_in_batches( + self, embeddings: FakeEmbeddingClient, index: FakeDocumentIndex + ) -> None: + catalog = FakeCatalogReader([_product(f"p{i}") for i in range(5)]) + use_case = IndexCatalog(catalog, embeddings, index, batch_size=2) + + report = await use_case.execute() + + assert report.indexed == 5 + assert [len(call) for call in embeddings.calls] == [2, 2, 1] + assert index.upsert_sizes == [2, 2, 1] diff --git a/tests/unit/contexts/assistant/domain/test_product_document.py b/tests/unit/contexts/assistant/domain/test_product_document.py new file mode 100644 index 0000000..70d420c --- /dev/null +++ b/tests/unit/contexts/assistant/domain/test_product_document.py @@ -0,0 +1,71 @@ +from uuid import UUID, uuid4 + +from contexts.assistant.domain.catalog_product import CatalogProduct +from contexts.assistant.domain.product_document import ProductDocument + +MODEL = "test-model" + + +def _product( + *, + product_id: UUID | None = None, + name: str = "Blue widget", + price_cents: int = 1999, +) -> CatalogProduct: + return CatalogProduct( + product_id=product_id or uuid4(), + sku="WIDGET-001", + name=name, + price_cents=price_cents, + currency="EUR", + ) + + +class TestFromCatalogProduct: + def test_renders_sku_name_and_price(self) -> None: + document = ProductDocument.from_catalog_product(_product(), embedding_model=MODEL) + + assert document.content == "SKU: WIDGET-001\nName: Blue widget\nPrice: 19.99 EUR" + + def test_renders_cents_with_two_digits(self) -> None: + product = _product(price_cents=5) + + document = ProductDocument.from_catalog_product(product, embedding_model=MODEL) + + assert document.content.endswith("Price: 0.05 EUR") + + def test_keeps_product_id(self) -> None: + product = _product() + + document = ProductDocument.from_catalog_product(product, embedding_model=MODEL) + + assert document.product_id == product.product_id + + def test_hash_is_stable_for_the_same_product(self) -> None: + product = _product() + + first = ProductDocument.from_catalog_product(product, embedding_model=MODEL) + second = ProductDocument.from_catalog_product(product, embedding_model=MODEL) + + assert first.content_hash == second.content_hash + assert len(first.content_hash) == 64 + + def test_hash_changes_when_the_product_changes(self) -> None: + product = _product() + renamed = _product(product_id=product.product_id, name="Red widget") + + original = ProductDocument.from_catalog_product(product, embedding_model=MODEL) + changed = ProductDocument.from_catalog_product(renamed, embedding_model=MODEL) + + assert original.content_hash != changed.content_hash + + +class TestContentHashCoversTheModel: + def test_r2_same_content_with_another_model_has_another_hash(self) -> None: + product = _product() + + first = ProductDocument.from_catalog_product(product, embedding_model="model-a") + second = ProductDocument.from_catalog_product(product, embedding_model="model-b") + + assert first.content == second.content + assert first.content_hash != second.content_hash diff --git a/tests/unit/shared/infrastructure/test_settings.py b/tests/unit/shared/infrastructure/test_settings.py new file mode 100644 index 0000000..03463a4 --- /dev/null +++ b/tests/unit/shared/infrastructure/test_settings.py @@ -0,0 +1,13 @@ +import pytest +from pydantic import ValidationError + +from shared.infrastructure.settings import Settings + + +class TestRagTopK: + def test_r4_accepts_up_to_the_hnsw_default_ef_search(self) -> None: + assert Settings(rag_top_k=40).rag_top_k == 40 + + def test_r4_rejects_more_than_the_hnsw_default_ef_search(self) -> None: + with pytest.raises(ValidationError, match="rag_top_k"): + Settings(rag_top_k=41)