From 998f4130c81576e8b28a5588386ae8336c5f72e2 Mon Sep 17 00:00:00 2001 From: Giuseppe Zileni <6775950+gzileni@users.noreply.github.com> Date: Tue, 4 Aug 2026 11:50:45 +0200 Subject: [PATCH] feat(llm): llama.cpp come motore predefinito per estrazione ed embedding Il server di inferenza self-hosted espone la forma OpenAI, quindi LlamaCppProvider usa /v1/chat/completions con response_format json_object: llama.cpp lo implementa vincolando il sampling con una grammatica JSON, quindi l'output e' ben formato per costruzione e non per buona volonta' del modello. Conta, visto che il chiamante fa un json.loads secco. Gli embedding sono selezionati separatamente (EMBEDDING_PROVIDER) da KG_LLM_PROVIDER: l'estrazione si cambia liberamente, gli embedding no, perche' dimensionano l'indice vettoriale. Il batch parte ora in un'unica richiesta invece di un round-trip per testo, e i vettori sono riordinati per "index" invece che per posizione: un batch permutato in silenzio attaccherebbe ogni vettore al chunk sbagliato. CAMBIO CHE RICHIEDE REINDICIZZAZIONE: Qwen3-Embedding-0.6B produce 1024 dimensioni contro le 768 di nomic-embed-text, quindi REDIS_VECTOR_DIM passa a 1024 e un'installazione con dati esistenti va ricostruita. create_index inghiottiva ogni eccezione: con l'indice gia' esistente a 768 la creazione era un no-op silenzioso e il disallineamento emergeva come query che non restituivano nulla. Ora la larghezza viene confrontata e l'avvio si interrompe stampando il FT.DROPINDEX da eseguire. Chi vuole restare su Ollama dichiara EMBEDDING_PROVIDER=ollama e REDIS_VECTOR_DIM=768. La validazione della forma della risposta sta fuori dal retry: una risposta troncata e' una violazione di contratto, non un guasto transitorio, e ritentarla ritardava soltanto l'errore di ~15 secondi. Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 10 +- docker-compose.ghcr.yml | 12 +++ docker-compose.yml | 9 +- knowledge-graph-api/.env.example | 37 +++++++- knowledge-graph-api/README.md | 40 +++++++- knowledge-graph-api/config/settings.py | 37 +++++++- knowledge-graph-api/pipeline/embedder.py | 95 +++++++++++++++++-- knowledge-graph-api/pipeline/llm/__init__.py | 13 ++- knowledge-graph-api/pipeline/llm/factory.py | 7 +- .../pipeline/llm/llamacpp_provider.py | 53 +++++++++++ knowledge-graph-api/storage/redis_vector.py | 56 ++++++++++- knowledge-graph-api/tests/conftest.py | 44 +++++++++ knowledge-graph-api/tests/test_embedder.py | 86 ++++++++++++++++- knowledge-graph-api/tests/test_ingest.py | 10 +- .../tests/test_llm_providers.py | 39 +++++++- 15 files changed, 510 insertions(+), 38 deletions(-) create mode 100644 knowledge-graph-api/pipeline/llm/llamacpp_provider.py diff --git a/README.md b/README.md index 67658f0..73cdd15 100644 --- a/README.md +++ b/README.md @@ -129,7 +129,7 @@ Data flows through three main paths: | **Vector Store** | Redis for AI (RedisSearch + RedisJSON) | latest | | **LLM Inference** | Ollama (local, no API key) | latest | | **LLM Model** | Llama 3 | latest | -| **Embedding Model** | nomic-embed-text (768 dim) | latest | +| **Embedding Model** | Qwen3-Embedding-0.6B (1024 dim) via llama.cpp | Q8_0 | | **REST API** | FastAPI + uvicorn | 0.115+ | | **Data Models** | Pydantic v2 + pydantic-settings | 2.7+ | | **Multi-Agent** | Microsoft Agent Framework (MAF) | latest | @@ -837,7 +837,7 @@ cp .env.example .env | --- | --- | --- | | `REDIS_URL` | `redis://redis:6379` | Connection URL (use `localhost` for local dev) | | `REDIS_INDEX_NAME` | `kg_vectors` | Vector index name | -| `REDIS_VECTOR_DIM` | `768` | Vector dimension (depends on embedding model) | +| `REDIS_VECTOR_DIM` | `1024` | Vector dimension — **must** match the embedding model (1024 Qwen3-Embedding, 768 nomic-embed-text). Cambiarlo impone il rebuild dell'indice | ### Ollama @@ -845,7 +845,11 @@ cp .env.example .env | --- | --- | --- | | `OLLAMA_BASE_URL` | `http://ollama:11434` | Ollama URL (use `localhost` for local dev) | | `OLLAMA_LLM_MODEL` | `llama3` | Text generation model | -| `OLLAMA_EMBEDDING_MODEL` | `nomic-embed-text` | Embedding model (768 dim) | +| `OLLAMA_EMBEDDING_MODEL` | `nomic-embed-text` | Embedding model (768 dim), solo con `EMBEDDING_PROVIDER=ollama` | +| `KG_LLM_PROVIDER` | `llamacpp` | Estrazione: `llamacpp` / `ollama` / `anthropic` | +| `EMBEDDING_PROVIDER` | `llamacpp` | Embedding: `llamacpp` / `ollama` | +| `LLAMACPP_BASE_URL` | `http://localhost:8080` | llama-swap, chat compatibile OpenAI | +| `LLAMACPP_EMBEDDING_BASE_URL` | `http://localhost:8081` | Server embedding dedicato | ### Chunking diff --git a/docker-compose.ghcr.yml b/docker-compose.ghcr.yml index c3493c1..63955b7 100644 --- a/docker-compose.ghcr.yml +++ b/docker-compose.ghcr.yml @@ -122,9 +122,21 @@ services: NEO4J_USER: "${NEO4J_USER:-neo4j}" NEO4J_PASSWORD: "${NEO4J_PASSWORD:-password}" REDIS_URL: "redis://redis:6379" + # Default engine: self-hosted llama.cpp. REDIS_VECTOR_DIM must match the + # embedding model's width (1024 for Qwen3-Embedding-0.6B). + KG_LLM_PROVIDER: "${KG_LLM_PROVIDER:-llamacpp}" + EMBEDDING_PROVIDER: "${EMBEDDING_PROVIDER:-llamacpp}" + REDIS_VECTOR_DIM: "${REDIS_VECTOR_DIM:-1024}" + LLAMACPP_BASE_URL: "${LLAMACPP_BASE_URL:-http://host.docker.internal:8080}" + LLAMACPP_LLM_MODEL: "${LLAMACPP_LLM_MODEL:-qwen3.5-9b}" + LLAMACPP_EMBEDDING_BASE_URL: "${LLAMACPP_EMBEDDING_BASE_URL:-http://host.docker.internal:8081}" + LLAMACPP_EMBEDDING_MODEL: "${LLAMACPP_EMBEDDING_MODEL:-qwen3-embedding}" + # Only used with KG_LLM_PROVIDER/EMBEDDING_PROVIDER=ollama. OLLAMA_BASE_URL: "${OLLAMA_BASE_URL:-http://host.docker.internal:11434}" OLLAMA_LLM_MODEL: "${OLLAMA_LLM_MODEL:-llama3}" OLLAMA_EMBEDDING_MODEL: "${OLLAMA_EMBEDDING_MODEL:-nomic-embed-text}" + extra_hosts: + - "host.docker.internal:host-gateway" depends_on: neo4j: condition: service_healthy diff --git a/docker-compose.yml b/docker-compose.yml index ec43964..63b1486 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -119,11 +119,18 @@ services: dockerfile: Dockerfile ports: - "8000:8000" - # OLLAMA_BASE_URL is read from .env: + # Default engine: the self-hosted llama.cpp server on the host, reached via + # host.docker.internal — chat on :8080 (llama-swap), embeddings on :8081. + # Both are read from .env (LLAMACPP_BASE_URL / LLAMACPP_EMBEDDING_BASE_URL). + # + # OLLAMA_BASE_URL only matters with KG_LLM_PROVIDER/EMBEDDING_PROVIDER=ollama: # default prod → http://host.docker.internal:11434 (Ollama on host) # profile gpu → http://ollama-gpu:11434 # profile cpu → http://ollama-cpu:11434 env_file: .env + # host.docker.internal is not resolvable on Linux without this. + extra_hosts: + - "host.docker.internal:host-gateway" depends_on: neo4j: condition: service_healthy diff --git a/knowledge-graph-api/.env.example b/knowledge-graph-api/.env.example index 78f0784..474c178 100644 --- a/knowledge-graph-api/.env.example +++ b/knowledge-graph-api/.env.example @@ -7,13 +7,40 @@ NEO4J_DATABASE=neo4j # Redis for AI (Vector Store) REDIS_URL=redis://redis:6379 REDIS_INDEX_NAME=kg_vectors -REDIS_VECTOR_DIM=768 +# DEVE combaciare con l'ampiezza del modello di embedding scelto: +# llamacpp + Qwen3-Embedding-0.6B → 1024 (misurato) +# ollama + nomic-embed-text → 768 +# Cambiarlo invalida un indice esistente: l'avvio si interrompe con l'istruzione +# di FT.DROPINDEX, e va poi rieseguita l'ingestione di TUTTI i documenti. +REDIS_VECTOR_DIM=1024 -# LLM provider for extraction stage: "ollama" (default) or "anthropic". -# Embeddings always go to Ollama regardless of this setting. -KG_LLM_PROVIDER=ollama +# Provider LLM per la fase di estrazione: +# "llamacpp" (default) → server di inferenza self-hosted +# "ollama" → comportamento legacy +# "anthropic" → Claude, opzionale +KG_LLM_PROVIDER=llamacpp -# Ollama (Inference API locale) +# Provider degli embedding, scelto separatamente dall'estrazione: e' l'unica +# scelta non reversibile a costo zero, perche' dimensiona l'indice vettoriale. +# "llamacpp" (default) oppure "ollama". +EMBEDDING_PROVIDER=llamacpp + +# llama.cpp / llama-swap — compatibile OpenAI, self-hosted. +# Due server distinti di proposito: l'embedding deve rispondere in millisecondi +# e non puo' stare dietro a uno swap di modello che dura minuti. +LLAMACPP_BASE_URL=http://host.docker.internal:8080 +LLAMACPP_LLM_MODEL=qwen3.5-9b +# Vuoto → usa LLAMACPP_LLM_MODEL +LLAMACPP_EXTRACTION_MODEL= +LLAMACPP_EMBEDDING_BASE_URL=http://host.docker.internal:8081 +LLAMACPP_EMBEDDING_MODEL=qwen3-embedding +# Solo se llama-server e' avviato con --api-key +LLAMACPP_API_KEY= +# Alto di proposito: una richiesta che arriva durante uno swap attende il +# caricamento, e il modello batch da 17,7 GB su disco rotante impiega ~2 minuti. +LLAMACPP_TIMEOUT_SECONDS=600 + +# Ollama (Inference API locale) — usato solo con KG_LLM_PROVIDER/EMBEDDING_PROVIDER=ollama OLLAMA_BASE_URL=http://ollama:11434 OLLAMA_LLM_MODEL=llama3 OLLAMA_EMBEDDING_MODEL=nomic-embed-text diff --git a/knowledge-graph-api/README.md b/knowledge-graph-api/README.md index daecf62..665c365 100644 --- a/knowledge-graph-api/README.md +++ b/knowledge-graph-api/README.md @@ -30,13 +30,48 @@ Handles document ingestion pipeline, Redis vector store, Neo4j graph database an | Data models | Pydantic v2 + pydantic-settings | | Graph DB | Neo4j 5.18 (async driver) | | Vector store | Redis Stack (RedisSearch + RedisJSON) | -| LLM / Embedding | Ollama (llama3 + nomic-embed-text) | +| LLM / Embedding | llama.cpp self-hosted (qwen3.5-9b + Qwen3-Embedding-0.6B, 1024 dim). Ollama e Claude selezionabili via `KG_LLM_PROVIDER` / `EMBEDDING_PROVIDER` | | Logging | structlog (JSON in prod, console in dev) | | Testing | pytest + pytest-asyncio + pytest-mock | | Linting | ruff | --- +## ⚠️ Migrazione: gli embedding passano a 1024 dimensioni + +Il provider predefinito è ora il **server di inferenza llama.cpp self-hosted**. +Per l'estrazione il cambio è trasparente. Per gli **embedding no**: + +| | prima | ora | +| --- | --- | --- | +| provider | Ollama | llama.cpp | +| modello | `nomic-embed-text` | `Qwen3-Embedding-0.6B-Q8_0` | +| dimensioni | 768 | **1024** | + +Vettori prodotti da modelli diversi non sono confrontabili, e l'indice Redis è +dimensionato sull'ampiezza del modello. **Un'installazione con dati esistenti +va reindicizzata:** + +```bash +redis-cli FT.DROPINDEX kg_vectors DD # DD cancella anche i documenti +# poi rieseguire l'ingestione di tutti i documenti +``` + +Non serve ricordarselo: `RedisVectorStore.create_index` confronta la larghezza +dell'indice vivo con `REDIS_VECTOR_DIM` e **si rifiuta di partire** su +disallineamento, stampando il comando qui sopra. Prima non lo faceva — la +creazione era un no-op silenzioso e il disallineamento emergeva come query che +non restituivano nulla. + +Per restare su Ollama senza reindicizzare, basta dichiararlo: + +```bash +EMBEDDING_PROVIDER=ollama +REDIS_VECTOR_DIM=768 +``` + +--- + ## Prerequisites - Python 3.11+ @@ -144,7 +179,8 @@ Document [1] File Routing → MIME type detection (PDF / DOCX / TXT) [2] Content Extraction → raw text + page count [3] Text Chunking → 1024-char chunks, 128-char overlap - [4] Embedding → 768-D vectors via Ollama (nomic-embed-text) + [4] Embedding → 1024-D vectors via llama.cpp (Qwen3-Embedding-0.6B), + l'intero batch in una sola richiesta [5] Deduplication → SHA-256 hash, skip existing chunks [6] Entity Extraction → LLM extracts entities and relations [7] Vector Storage → upsert into Redis (RedisSearch) diff --git a/knowledge-graph-api/config/settings.py b/knowledge-graph-api/config/settings.py index 0de7225..d67581e 100644 --- a/knowledge-graph-api/config/settings.py +++ b/knowledge-graph-api/config/settings.py @@ -15,13 +15,40 @@ class Settings(BaseSettings): # Redis (Vector Store) REDIS_URL: str = "redis://localhost:6379" REDIS_INDEX_NAME: str = "kg_vectors" - REDIS_VECTOR_DIM: int = 768 + # Must match the output width of EMBEDDING_PROVIDER's model: + # llamacpp + Qwen3-Embedding-0.6B → 1024 (measured, not declared) + # ollama + nomic-embed-text → 768 + # Changing this invalidates an existing index — see RedisVectorStore.create_index, + # which refuses to run against an index built at a different width. + REDIS_VECTOR_DIM: int = 1024 # LLM provider selection for the extraction stage. - # "ollama" (default) keeps the legacy local-inference behaviour. - # "anthropic" routes extraction to Claude (Haiku by default). - # Embeddings always go to Ollama regardless of this setting. - KG_LLM_PROVIDER: str = "ollama" + # "llamacpp" (default) → the self-hosted llama.cpp inference server. + # "ollama" → the legacy local-inference behaviour. + # "anthropic" → Claude (Haiku by default), opt-in. + KG_LLM_PROVIDER: str = "llamacpp" + + # Embeddings are selected separately from extraction: they are the one + # thing that cannot be swapped freely, because the vector index is sized + # to the model's output width. + # "llamacpp" (default) or "ollama". + EMBEDDING_PROVIDER: str = "llamacpp" + + # llama.cpp / llama-swap — OpenAI-compatible, self-hosted. + # Chat and embeddings are two separate servers on purpose: embeddings must + # answer in milliseconds and cannot sit behind a model swap that takes + # minutes. + LLAMACPP_BASE_URL: str = "http://localhost:8080" + LLAMACPP_LLM_MODEL: str = "qwen3.5-9b" + # Empty → falls back to LLAMACPP_LLM_MODEL. + LLAMACPP_EXTRACTION_MODEL: str = "" + LLAMACPP_EMBEDDING_BASE_URL: str = "http://localhost:8081" + LLAMACPP_EMBEDDING_MODEL: str = "qwen3-embedding" + # Set only if llama-server was started with --api-key. + LLAMACPP_API_KEY: str = "" + # A request arriving mid-swap waits for the load to finish; the 17.7 GB + # batch model needs ~2 minutes from spinning disk just to reach VRAM. + LLAMACPP_TIMEOUT_SECONDS: float = 600.0 # Ollama (Inference API locale) OLLAMA_BASE_URL: str = "http://localhost:11434" diff --git a/knowledge-graph-api/pipeline/embedder.py b/knowledge-graph-api/pipeline/embedder.py index 0cd1757..261d3e3 100644 --- a/knowledge-graph-api/pipeline/embedder.py +++ b/knowledge-graph-api/pipeline/embedder.py @@ -1,4 +1,17 @@ -"""Embedding generation via Ollama HTTP API.""" +"""Embedding generation. + +Two backends, selected by ``EMBEDDING_PROVIDER``: + +* ``llamacpp`` (default) — the self-hosted llama.cpp embedding server, an + OpenAI-compatible ``/v1/embeddings`` endpoint that accepts a whole batch in + one request. +* ``ollama`` — the legacy path, one HTTP round-trip per text. + +Embeddings are configured separately from the extraction LLM +(``KG_LLM_PROVIDER``) because they are the one choice that cannot be changed +freely: the Redis index is sized to the model's output width, so switching +provider or model means rebuilding it. +""" import httpx from tenacity import retry, stop_after_attempt, wait_exponential @@ -8,27 +21,97 @@ class Embedder: - """Generates embeddings using Ollama's nomic-embed-text model.""" + """Generates embeddings from the configured provider.""" def __init__(self) -> None: - self.base_url = settings.OLLAMA_BASE_URL - self.model = settings.OLLAMA_EMBEDDING_MODEL + self.provider = settings.EMBEDDING_PROVIDER.lower().strip() + if self.provider == "llamacpp": + self.base_url = settings.LLAMACPP_EMBEDDING_BASE_URL + self.model = settings.LLAMACPP_EMBEDDING_MODEL + elif self.provider == "ollama": + self.base_url = settings.OLLAMA_BASE_URL + self.model = settings.OLLAMA_EMBEDDING_MODEL + else: + raise ValueError( + f"Unknown EMBEDDING_PROVIDER={settings.EMBEDDING_PROVIDER!r} " + "(supported: 'llamacpp', 'ollama')" + ) async def embed(self, texts: list[str]) -> list[list[float]]: - """Embed a list of texts sequentially. + """Embed a list of texts. Args: texts: Texts to embed. Returns: - List of 768-dimensional vectors. + One vector per input text, in the same order, each of + ``settings.REDIS_VECTOR_DIM`` dimensions. """ + if not texts: + return [] + if self.provider == "llamacpp": + # One request for the whole batch: the server processes them in a + # single forward pass, which is far cheaper than a round-trip per + # text (measured: 64 texts / 1142 tokens in 0.58 s). + data = await self._post_embeddings(texts) + # Parsing is deliberately outside the retry: a malformed or + # truncated response is a contract violation, not a transient + # fault, and retrying it just delays the error by ~15 seconds. + return self._parse_embeddings(data, texts) vectors: list[list[float]] = [] for text in texts: vec = await self._embed_single(text) vectors.append(vec) return vectors + @retry( + stop=stop_after_attempt(settings.MAX_RETRIES), + wait=wait_exponential(multiplier=1, min=1, max=30), + reraise=True, + ) + async def _post_embeddings(self, texts: list[str]) -> dict: + """POST a batch to the OpenAI-compatible ``/v1/embeddings`` endpoint.""" + headers = ( + {"Authorization": f"Bearer {settings.LLAMACPP_API_KEY}"} + if settings.LLAMACPP_API_KEY + else None + ) + async with httpx.AsyncClient(timeout=settings.LLAMACPP_TIMEOUT_SECONDS) as client: + response = await client.post( + f"{self.base_url.rstrip('/')}/v1/embeddings", + headers=headers, + json={"model": self.model, "input": texts}, + ) + response.raise_for_status() + return response.json() + + def _parse_embeddings(self, data: dict, texts: list[str]) -> list[list[float]]: + """Validate the response and return one vector per input, in order.""" + try: + rows = data["data"] + except (KeyError, TypeError) as exc: + raise RuntimeError(f"unexpected embeddings response shape: {exc}") from exc + + if len(rows) != len(texts): + raise RuntimeError( + f"embeddings endpoint returned {len(rows)} vectors for {len(texts)} inputs" + ) + # The OpenAI shape carries an explicit index; sort by it rather than + # trusting positional order. A silently permuted batch would attach + # every vector to the wrong chunk — a corruption nothing downstream + # would catch. + rows = sorted(rows, key=lambda r: r.get("index", 0)) + vectors = [list(row["embedding"]) for row in rows] + + logger.debug( + "embeddings_generated", + provider=self.provider, + model=self.model, + n=len(vectors), + dim=len(vectors[0]) if vectors else 0, + ) + return vectors + @retry( stop=stop_after_attempt(settings.MAX_RETRIES), wait=wait_exponential(multiplier=1, min=1, max=30), diff --git a/knowledge-graph-api/pipeline/llm/__init__.py b/knowledge-graph-api/pipeline/llm/__init__.py index d21967b..dd497c4 100644 --- a/knowledge-graph-api/pipeline/llm/__init__.py +++ b/knowledge-graph-api/pipeline/llm/__init__.py @@ -2,12 +2,15 @@ The extraction stage of the pipeline calls an LLM to convert text chunks into structured JSON. The provider is pluggable so the same pipeline can -run against a local Ollama model, a hosted Anthropic Claude model, or -future providers without touching the call sites. +run against the self-hosted llama.cpp server (the default), a local Ollama +model, or a hosted Anthropic Claude model, without touching the call sites. -Embeddings are NOT covered by this abstraction — they remain wired to -``OLLAMA_EMBEDDING_MODEL`` via ``pipeline.embedder`` because the vector -index is sized to that model's output (768 dims for ``nomic-embed-text``). +Embeddings are NOT covered by this abstraction: they are selected on their +own by ``EMBEDDING_PROVIDER`` in ``pipeline.embedder``. The split is +deliberate — extraction can be swapped freely, embeddings cannot, because +the Redis index is sized to the model's output width (1024 for +``Qwen3-Embedding-0.6B``, 768 for ``nomic-embed-text``) and changing it +means rebuilding the index and re-ingesting every document. """ from __future__ import annotations diff --git a/knowledge-graph-api/pipeline/llm/factory.py b/knowledge-graph-api/pipeline/llm/factory.py index 17c551b..2c5e56a 100644 --- a/knowledge-graph-api/pipeline/llm/factory.py +++ b/knowledge-graph-api/pipeline/llm/factory.py @@ -5,18 +5,21 @@ from config.settings import settings from pipeline.llm.anthropic_provider import AnthropicProvider from pipeline.llm.base import LLMProvider +from pipeline.llm.llamacpp_provider import LlamaCppProvider from pipeline.llm.ollama_provider import OllamaProvider def get_llm_provider() -> LLMProvider: - """Return the provider selected by ``KG_LLM_PROVIDER`` (default: ollama).""" + """Return the provider selected by ``KG_LLM_PROVIDER`` (default: llamacpp).""" choice = settings.KG_LLM_PROVIDER.lower().strip() + if choice == "llamacpp": + return LlamaCppProvider() if choice == "anthropic": return AnthropicProvider() if choice == "ollama": return OllamaProvider() raise ValueError( f"Unknown KG_LLM_PROVIDER={settings.KG_LLM_PROVIDER!r} " - "(supported: 'ollama', 'anthropic')" + "(supported: 'llamacpp', 'ollama', 'anthropic')" ) diff --git a/knowledge-graph-api/pipeline/llm/llamacpp_provider.py b/knowledge-graph-api/pipeline/llm/llamacpp_provider.py new file mode 100644 index 0000000..f1d3347 --- /dev/null +++ b/knowledge-graph-api/pipeline/llm/llamacpp_provider.py @@ -0,0 +1,53 @@ +"""llama.cpp-backed LLM provider (self-hosted inference server, default).""" + +from __future__ import annotations + +import httpx + +from config.settings import settings +from pipeline.llm.base import LLMProvider + + +class LlamaCppProvider(LLMProvider): + """Calls llama.cpp's OpenAI-compatible chat endpoint. + + ``response_format: json_object`` is not a polite request here: llama.cpp + implements it by constraining sampling with a JSON grammar, so the output + is well-formed by construction rather than by the model's good behaviour. + That matters for the extraction stage, whose caller does a bare + ``json.loads`` on whatever comes back. + """ + + def __init__(self, *, base_url: str | None = None, model: str | None = None) -> None: + self._base_url = (base_url or settings.LLAMACPP_BASE_URL).rstrip("/") + self._model = ( + model or settings.LLAMACPP_EXTRACTION_MODEL or settings.LLAMACPP_LLM_MODEL + ) + self._api_key = settings.LLAMACPP_API_KEY + self._timeout = settings.LLAMACPP_TIMEOUT_SECONDS + + async def chat_json(self, system: str, user: str) -> str: + headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + async with httpx.AsyncClient(timeout=self._timeout) as client: + response = await client.post( + f"{self._base_url}/v1/chat/completions", + headers=headers, + json={ + "model": self._model, + "messages": [ + {"role": "system", "content": system}, + {"role": "user", "content": user}, + ], + "stream": False, + "temperature": 0, + "response_format": {"type": "json_object"}, + }, + ) + response.raise_for_status() + data = response.json() + try: + return str(data["choices"][0]["message"]["content"]) + except (KeyError, IndexError, TypeError) as exc: + raise RuntimeError( + f"unexpected llama.cpp response shape: {exc}" + ) from exc diff --git a/knowledge-graph-api/storage/redis_vector.py b/knowledge-graph-api/storage/redis_vector.py index f2e90a6..36fa53e 100644 --- a/knowledge-graph-api/storage/redis_vector.py +++ b/knowledge-graph-api/storage/redis_vector.py @@ -60,11 +60,61 @@ async def create_index(self) -> None: await self._client.ft(self._index_name).create_index( schema, definition=definition ) - logger.info("redis_index_created", index=self._index_name) + logger.info( + "redis_index_created", + index=self._index_name, + dim=settings.REDIS_VECTOR_DIM, + ) + return except Exception: - # Index already exists + # Index already exists — the only expected failure here. logger.debug("redis_index_exists", index=self._index_name) + # An index built at a different width does NOT fail here: creation is a + # no-op and every later KNN query fails, or worse returns nonsense. So + # check the width explicitly and refuse to run on a mismatch. + existing = await self._existing_vector_dim() + if existing is not None and existing != settings.REDIS_VECTOR_DIM: + raise RuntimeError( + f"Redis index {self._index_name!r} was built with DIM={existing}, " + f"but REDIS_VECTOR_DIM={settings.REDIS_VECTOR_DIM}. Embeddings from " + "different models are not interchangeable: the index has to be " + "rebuilt and every document re-ingested.\n" + f" redis-cli FT.DROPINDEX {self._index_name} DD\n" + "then re-run the ingestion pipeline." + ) + + async def _existing_vector_dim(self) -> int | None: + """Return the vector width of the live index, or None if undetectable. + + ``FT.INFO`` attribute rows differ across redis-py and Redis Stack + versions (flat list vs mapping, bytes vs str), so this scans for the + key/value pair rather than indexing a fixed position. Undetectable is + not an error: it must not block startup on a shape we don't recognise. + """ + try: + info = await self._client.ft(self._index_name).info() + except Exception: # noqa: BLE001 — pragma: no cover + # Deliberately broad: this is a best-effort safety check, and no + # failure to introspect may block startup. + return None + + def _norm(value: object) -> str: + return ( + value.decode("utf-8", "replace") if isinstance(value, bytes) else str(value) + ).lower() + + attributes = info.get("attributes") if isinstance(info, dict) else None + for attr in attributes or []: + row = list(attr) if isinstance(attr, (list, tuple)) else [] + for i, item in enumerate(row[:-1]): + if _norm(item) == "dim": + try: + return int(_norm(row[i + 1])) + except ValueError: + return None + return None + # ── CRUD ───────────────────────────────────────────────────────── async def upsert(self, doc: VectorDocument) -> None: @@ -110,7 +160,7 @@ async def vector_search( """Run a KNN vector similarity search. Args: - query_vector: The 768-D query embedding. + query_vector: The query embedding, of settings.REDIS_VECTOR_DIM dims. top_k: Number of results to return. namespace: Optional thread_id filter. diff --git a/knowledge-graph-api/tests/conftest.py b/knowledge-graph-api/tests/conftest.py index 2f871f8..d0e6668 100644 --- a/knowledge-graph-api/tests/conftest.py +++ b/knowledge-graph-api/tests/conftest.py @@ -6,6 +6,7 @@ import pytest +from config.settings import settings from models.base import VectorDocument from models.graph_node import GraphNode from models.relation import Relation @@ -153,3 +154,46 @@ def mock_redis_client(): mock_from_url.return_value = client yield client + + +@pytest.fixture +def mock_llamacpp_client(monkeypatch): + """Mock for the llama.cpp OpenAI-compatible endpoints via httpx. + + Embeddings answer with the OpenAI shape and 1024 dims (the measured width + of Qwen3-Embedding-0.6B), one row per input. Rows are returned in + DELIBERATELY SHUFFLED order with explicit ``index`` fields, so a caller + that trusts positional order instead of the index fails the test. + """ + monkeypatch.setattr(settings, "EMBEDDING_PROVIDER", "llamacpp") + monkeypatch.setattr(settings, "KG_LLM_PROVIDER", "llamacpp") + + with patch("httpx.AsyncClient") as mock_cls: + client = AsyncMock() + + def _post(url, **kwargs): + payload = kwargs.get("json") or {} + response = MagicMock() + response.status_code = 200 + response.raise_for_status = MagicMock() + if url.endswith("/v1/embeddings"): + items = payload.get("input") or [] + rows = [ + {"index": i, "embedding": [float(i) / 100] * 1024} + for i in range(len(items)) + ] + response.json.return_value = {"data": list(reversed(rows))} + else: + response.json.return_value = { + "choices": [ + {"message": {"content": '{"entities": [], "relations": []}'}} + ] + } + return response + + client.post = AsyncMock(side_effect=_post) + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=False) + + mock_cls.return_value = client + yield client diff --git a/knowledge-graph-api/tests/test_embedder.py b/knowledge-graph-api/tests/test_embedder.py index 2cf7025..9656a43 100644 --- a/knowledge-graph-api/tests/test_embedder.py +++ b/knowledge-graph-api/tests/test_embedder.py @@ -2,11 +2,22 @@ import pytest +from config.settings import settings from pipeline.embedder import Embedder +@pytest.fixture +def ollama_embeddings(monkeypatch, mock_ollama_client): + """Force the legacy Ollama embedding path (the default is llama.cpp).""" + monkeypatch.setattr(settings, "EMBEDDING_PROVIDER", "ollama") + return mock_ollama_client + + +# ── Ollama (legacy path) ───────────────────────────────────────────── + + @pytest.mark.asyncio -async def test_embed_single(mock_ollama_client) -> None: +async def test_embed_single_ollama(ollama_embeddings) -> None: """Embedder should return a 768-D vector for a single text.""" embedder = Embedder() vectors = await embedder.embed(["Hello world"]) @@ -15,9 +26,80 @@ async def test_embed_single(mock_ollama_client) -> None: @pytest.mark.asyncio -async def test_embed_batch(mock_ollama_client) -> None: +async def test_embed_batch_ollama(ollama_embeddings) -> None: """Embedder should return one vector per input text.""" embedder = Embedder() texts = ["First sentence.", "Second sentence.", "Third sentence."] vectors = await embedder.embed(texts) assert len(vectors) == len(texts) + # The legacy path is one round-trip per text. + assert ollama_embeddings.post.await_count == len(texts) + + +# ── llama.cpp (default) ────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_embed_single_llamacpp(mock_llamacpp_client) -> None: + """Default provider yields the 1024-D Qwen3-Embedding vector.""" + embedder = Embedder() + vectors = await embedder.embed(["Hello world"]) + assert len(vectors) == 1 + assert len(vectors[0]) == 1024 + + +@pytest.mark.asyncio +async def test_embed_batch_llamacpp_is_a_single_request(mock_llamacpp_client) -> None: + """The whole batch goes out in ONE request, not one per text.""" + embedder = Embedder() + texts = ["First.", "Second.", "Third.", "Fourth."] + vectors = await embedder.embed(texts) + + assert len(vectors) == len(texts) + assert mock_llamacpp_client.post.await_count == 1 + + call = mock_llamacpp_client.post.await_args + assert call.args[0].endswith("/v1/embeddings") + assert call.kwargs["json"]["input"] == texts + + +@pytest.mark.asyncio +async def test_embed_batch_llamacpp_respects_index_order(mock_llamacpp_client) -> None: + """Vectors must be reordered by ``index``, not taken as they arrive. + + The fixture returns rows reversed on purpose: trusting positional order + would attach every vector to the wrong chunk, silently poisoning the index. + """ + embedder = Embedder() + vectors = await embedder.embed(["a", "b", "c"]) + # The fixture encodes the input position in the vector's value. + assert [v[0] for v in vectors] == [0.0, 0.01, 0.02] + + +@pytest.mark.asyncio +async def test_embed_rejects_truncated_response(mock_llamacpp_client) -> None: + """Fewer vectors than inputs must raise, not silently misalign.""" + from unittest.mock import AsyncMock, MagicMock + + short = MagicMock() + short.status_code = 200 + short.raise_for_status = MagicMock() + short.json.return_value = {"data": [{"index": 0, "embedding": [0.0] * 1024}]} + mock_llamacpp_client.post = AsyncMock(return_value=short) + + embedder = Embedder() + with pytest.raises(RuntimeError, match="1 vectors for 3 inputs"): + await embedder.embed(["a", "b", "c"]) + + +@pytest.mark.asyncio +async def test_embed_empty_list_makes_no_request(mock_llamacpp_client) -> None: + embedder = Embedder() + assert await embedder.embed([]) == [] + assert mock_llamacpp_client.post.await_count == 0 + + +def test_unknown_provider_is_rejected(monkeypatch) -> None: + monkeypatch.setattr(settings, "EMBEDDING_PROVIDER", "wrong") + with pytest.raises(ValueError, match="Unknown EMBEDDING_PROVIDER"): + Embedder() diff --git a/knowledge-graph-api/tests/test_ingest.py b/knowledge-graph-api/tests/test_ingest.py index d9a9183..4f4d13c 100644 --- a/knowledge-graph-api/tests/test_ingest.py +++ b/knowledge-graph-api/tests/test_ingest.py @@ -12,8 +12,14 @@ @pytest.mark.asyncio -async def test_ingest_txt_file(mock_ollama_client, mock_neo4j_driver, mock_redis_client) -> None: - """Ingestion of a plain text file should produce chunks.""" +async def test_ingest_txt_file( + mock_llamacpp_client, mock_neo4j_driver, mock_redis_client +) -> None: + """Ingestion of a plain text file should produce chunks. + + Runs against the shipped default (llama.cpp for both extraction and + embeddings), so this covers the configuration that actually ships. + """ # Create a temp txt file with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False, encoding="utf-8") as f: f.write("Redis is great. Neo4j is great. Ollama is great.") diff --git a/knowledge-graph-api/tests/test_llm_providers.py b/knowledge-graph-api/tests/test_llm_providers.py index 9975672..bec6234 100644 --- a/knowledge-graph-api/tests/test_llm_providers.py +++ b/knowledge-graph-api/tests/test_llm_providers.py @@ -1,4 +1,4 @@ -"""Tests for the LLM provider abstraction (Ollama + Anthropic + factory).""" +"""Tests for the LLM provider abstraction (llama.cpp + Ollama + Anthropic + factory).""" from __future__ import annotations @@ -9,6 +9,7 @@ from config.settings import settings from pipeline.llm.anthropic_provider import AnthropicProvider, _strip_code_fences from pipeline.llm.factory import get_llm_provider +from pipeline.llm.llamacpp_provider import LlamaCppProvider from pipeline.llm.ollama_provider import OllamaProvider @@ -94,7 +95,13 @@ def test_anthropic_provider_rejects_missing_key(monkeypatch) -> None: AnthropicProvider() -def test_factory_returns_ollama_by_default(monkeypatch) -> None: +def test_factory_returns_llamacpp_by_default() -> None: + """The shipped default is the self-hosted inference server.""" + assert settings.KG_LLM_PROVIDER == "llamacpp" + assert isinstance(get_llm_provider(), LlamaCppProvider) + + +def test_factory_returns_ollama_when_configured(monkeypatch) -> None: monkeypatch.setattr(settings, "KG_LLM_PROVIDER", "ollama") assert isinstance(get_llm_provider(), OllamaProvider) @@ -107,6 +114,34 @@ def test_factory_returns_anthropic_when_configured(monkeypatch) -> None: assert isinstance(provider, AnthropicProvider) +@pytest.mark.asyncio +async def test_llamacpp_provider_calls_openai_chat_endpoint(mock_llamacpp_client) -> None: + """LlamaCppProvider POSTs the OpenAI shape and returns the message content.""" + provider = LlamaCppProvider(base_url="http://inference:8080", model="qwen3.5-9b") + result = await provider.chat_json("sys", "user") + + assert result == '{"entities": [], "relations": []}' + call = mock_llamacpp_client.post.await_args + assert call.args[0] == "http://inference:8080/v1/chat/completions" + payload = call.kwargs["json"] + assert payload["model"] == "qwen3.5-9b" + assert payload["messages"][0] == {"role": "system", "content": "sys"} + assert payload["messages"][1] == {"role": "user", "content": "user"} + # Grammar-constrained JSON: the caller does a bare json.loads on this. + assert payload["response_format"] == {"type": "json_object"} + assert payload["temperature"] == 0 + + +@pytest.mark.asyncio +async def test_llamacpp_provider_falls_back_to_llm_model(monkeypatch, mock_llamacpp_client) -> None: + """An empty LLAMACPP_EXTRACTION_MODEL means 'use the chat model'.""" + monkeypatch.setattr(settings, "LLAMACPP_EXTRACTION_MODEL", "") + monkeypatch.setattr(settings, "LLAMACPP_LLM_MODEL", "qwen3.6-35b") + provider = LlamaCppProvider() + await provider.chat_json("sys", "user") + assert mock_llamacpp_client.post.await_args.kwargs["json"]["model"] == "qwen3.6-35b" + + def test_factory_rejects_unknown_provider(monkeypatch) -> None: monkeypatch.setattr(settings, "KG_LLM_PROVIDER", "wrong") with pytest.raises(ValueError, match="Unknown KG_LLM_PROVIDER"):