From c07e4e22a3c44cbac452aa6db11c05fd951ba363 Mon Sep 17 00:00:00 2001 From: Francesco Vadicamo Date: Mon, 13 Jul 2026 01:37:34 +0000 Subject: [PATCH 1/4] feat(rag): remote TEI embedding and reranker providers (FEAT-024) VEKTRA_EMBEDDING_PROVIDER=tei was documented but no provider existed: main.py unconditionally instantiated in-process SentenceTransformers, so every instance duplicates embedding compute and long-window models served by shared TEI instances (bge-m3, 8192-token window vs MiniLM's 128) were unreachable. - TEIEmbeddingProvider (vektra-index): native /embed API, batch chunking to the TEI client batch limit, dimensions from /info with an /embed probe fallback, Bearer auth, health check. - TEIRerankerService (vektra-core): TEI /rerank with the same RerankResult semantics as the in-process service (BUG-015 score propagation, sigmoid normalization when raw scores are detected); shared _build_rerank_result helper and RerankerProtocol for typing. - Dimension plumbing: the Qdrant collection is sized from the active embedding provider instead of the hardcoded 384 default, and ensure_collection fails fast with a clear message on a dimension mismatch with an existing collection. - The cohere rerank option now passes VEKTRA_RERANK_API_KEY through; it never received a key before (dead as wired). Refs: FEAT-024, ADR-0013, ARCH-035, ARCH-036 Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 1 + uv.lock | 4 + vektra-app/src/vektra_app/main.py | 37 ++++- vektra-core/pyproject.toml | 1 + .../src/vektra_core/advanced_pipeline.py | 4 +- vektra-core/src/vektra_core/reranker.py | 137 ++++++++++++++---- vektra-core/tests/test_reranker.py | 77 +++++++++- vektra-index/pyproject.toml | 1 + .../src/vektra_index/providers/qdrant.py | 18 ++- .../src/vektra_index/providers/tei.py | 121 ++++++++++++++++ vektra-index/tests/test_tei_provider.py | 137 ++++++++++++++++++ vektra-shared/src/vektra_shared/config.py | 29 +++- 12 files changed, 525 insertions(+), 42 deletions(-) create mode 100644 vektra-index/src/vektra_index/providers/tei.py create mode 100644 vektra-index/tests/test_tei_provider.py diff --git a/CHANGELOG.md b/CHANGELOG.md index fb2613fc..6e719f54 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,7 @@ Convention (Keep a Changelog 1.1.0): ### Added +- **rag**: remote embedding and reranking via HuggingFace Text Embeddings Inference (FEAT-024). `VEKTRA_EMBEDDING_PROVIDER=tei` embeds through a TEI instance (`VEKTRA_TEI_URL`/`VEKTRA_TEI_API_KEY`, native `/embed` API) instead of in-process sentence-transformers, enabling shared host inference and long-window models (bge-m3: 8192 tokens vs MiniLM's 128, which silently truncates 500-token chunks today). `VEKTRA_RERANK_PROVIDER=tei` reranks through TEI `/rerank` (`VEKTRA_RERANK_TEI_URL`/`VEKTRA_RERANK_TEI_API_KEY`). The Qdrant collection is now sized from the active embedding provider's dimensions instead of a hardcoded 384 (latent bug for any non-384 model), with a clear startup error on dimension mismatch against an existing collection. The `cohere` rerank option now actually passes its API key (`VEKTRA_RERANK_API_KEY`); it was dead as wired. - **rag**: optional retrieval-filter rescue for multi-part questions (TECH-007, `VEKTRA_RETRIEVAL_RESCUE_TOP_K` + `VEKTRA_RETRIEVAL_RESCUE_FLOOR`, default off). When `VEKTRA_MIN_RELEVANCE_SCORE` empties the candidate set, keep the top-N chunks above an absolute floor instead of refusing: the cross-encoder scores each partial-answer chunk of a comparative/multi-part question below the threshold (it answers only one part), so on the eval corpus 9/10 multi-chunk questions died at the filter with `before=5 after=0` despite 90% raw retrieval hit. With the rescue, borderline sets reach the LLM, which arbitrates via strict grounding. The `retrieval_filter` trace step now records a `rescued` count. - **rag**: optional per-namespace inline source citations (FEAT-021, `citations_enabled` in the namespace config JSONB via `PATCH /api/v1/admin/namespaces/{id}/config`, default off, advanced pipeline only). When enabled, the system prompt instructs the LLM to add inline `[n]` markers matching the `` elements, the context template carries a `title` attribute ("filename, p.N"), and each returned source includes a `title` field; the learn widget renders the markers as superscripts with a tooltip. Default-off renders byte-identical prompts; `prompt_version` changes anyway because the template files changed (trace comparability note). - **rag**: optional parent chunk expansion in the advanced query pipeline (FEAT-017, `VEKTRA_PARENT_EXPANSION_ENABLED`, default off). With `VEKTRA_CHUNKING_STRATEGY=dual`, retrieved child chunks are replaced with their parent chunk's text after the retrieval filter and before token budgeting; children of the same parent collapse into the highest-scored one. Parent-child linkage is now actually persisted (deterministic `uuid5(doc_id, position)` ids, `parent_id` in the Qdrant payload and in the pgvector column), a new `VectorStoreProvider.retrieve()` fetches chunks by id, and the trace records `children_expanded`/`siblings_merged`/`parents_fetched` in a `parent_expansion` step. diff --git a/uv.lock b/uv.lock index d0e7deed..112aae6f 100644 --- a/uv.lock +++ b/uv.lock @@ -5236,6 +5236,7 @@ version = "0.5.1" source = { editable = "vektra-core" } dependencies = [ { name = "fastapi" }, + { name = "httpx" }, { name = "jinja2" }, { name = "litellm" }, { name = "presidio-analyzer" }, @@ -5256,6 +5257,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "fastapi", specifier = ">=0.115" }, + { name = "httpx", specifier = ">=0.27" }, { name = "jinja2", specifier = ">=3.1" }, { name = "litellm", specifier = ">=1.83.10" }, { name = "presidio-analyzer", specifier = ">=2.2" }, @@ -5280,6 +5282,7 @@ source = { editable = "vektra-index" } dependencies = [ { name = "asyncpg" }, { name = "fastapi" }, + { name = "httpx" }, { name = "pgvector" }, { name = "pydantic" }, { name = "sentence-transformers" }, @@ -5310,6 +5313,7 @@ requires-dist = [ { name = "asyncpg", specifier = ">=0.29" }, { name = "fastapi", specifier = ">=0.115" }, { name = "fastembed", marker = "extra == 'sparse'", specifier = ">=0.4" }, + { name = "httpx", specifier = ">=0.27" }, { name = "pgvector", specifier = ">=0.3" }, { name = "pydantic", specifier = ">=2,<3" }, { name = "qdrant-client", marker = "extra == 'qdrant'", specifier = ">=1.12" }, diff --git a/vektra-app/src/vektra_app/main.py b/vektra-app/src/vektra_app/main.py index 25a825fa..f112f9a8 100644 --- a/vektra-app/src/vektra_app/main.py +++ b/vektra-app/src/vektra_app/main.py @@ -31,6 +31,7 @@ from vektra_shared.config import QueryPipelineConfig, VektraSettings from vektra_shared.db import init_db from vektra_shared.errors import ERR_CONFIG_001, ErrorCategory, ErrorResponse +from vektra_shared.protocols import EmbeddingProvider from vektra_shared.registry import ProviderRegistry from vektra_shared.startup import ( StartupValidationError, @@ -126,15 +127,31 @@ async def _step_5_register_providers( registry.register("llm", "default", llm_provider) # --- Embedding --- - from vektra_index.providers.sentence_transformers import ( - SentenceTransformersProvider, - ) + embedding_provider: EmbeddingProvider + if settings.embedding_provider == "tei": + from vektra_index.providers.tei import TEIEmbeddingProvider - embedding_provider = SentenceTransformersProvider( - model_name=settings.embedding_model - ) - registry.register("embedding", "default", embedding_provider) - registry.register("embedding", "sentence-transformers", embedding_provider) + embedding_provider = TEIEmbeddingProvider( + url=settings.tei_url, + api_key=settings.tei_api_key, + ) + registry.register("embedding", "default", embedding_provider) + registry.register("embedding", "tei", embedding_provider) + log.info( + "embedding_registered", + provider="tei", + url=_redact_url(settings.tei_url), + ) + else: + from vektra_index.providers.sentence_transformers import ( + SentenceTransformersProvider, + ) + + embedding_provider = SentenceTransformersProvider( + model_name=settings.embedding_model + ) + registry.register("embedding", "default", embedding_provider) + registry.register("embedding", "sentence-transformers", embedding_provider) # --- Vector store --- from vektra_index.adapters import VectorStoreServiceAdapter @@ -168,6 +185,10 @@ async def _step_5_register_providers( api_key=settings.qdrant_api_key, collection_name=settings.qdrant_collection, active_index_version=settings.active_index_version, + # FEAT-024: size the collection from the active embedding model + # instead of the hardcoded 384 default (latent bug for any + # non-384 model; bge-m3 via TEI is 1024). + dense_dimensions=embedding_provider.dimensions(), ) registry.register("vector_store", "default", qdrant_provider) registry.register("vector_store", "qdrant", qdrant_provider) diff --git a/vektra-core/pyproject.toml b/vektra-core/pyproject.toml index f1dabe5b..b8259147 100644 --- a/vektra-core/pyproject.toml +++ b/vektra-core/pyproject.toml @@ -14,6 +14,7 @@ dependencies = [ "presidio-analyzer>=2.2", "presidio-anonymizer>=2.2", "rerankers[flashrank]>=0.5", + "httpx>=0.27", ] [build-system] diff --git a/vektra-core/src/vektra_core/advanced_pipeline.py b/vektra-core/src/vektra_core/advanced_pipeline.py index 4b969f56..14dfec3d 100644 --- a/vektra-core/src/vektra_core/advanced_pipeline.py +++ b/vektra-core/src/vektra_core/advanced_pipeline.py @@ -39,7 +39,7 @@ _history_to_messages, _trace_to_dict, ) -from vektra_core.reranker import RerankerService +from vektra_core.reranker import RerankerProtocol from vektra_core.templates import TemplateRenderer from vektra_shared.config import LLMConfig, QueryPipelineConfig from vektra_shared.protocols import ( @@ -87,7 +87,7 @@ def __init__( renderer: TemplateRenderer, pipeline_config: QueryPipelineConfig, sparse_embedding: SparseEmbeddingProvider | None = None, - reranker: RerankerService | None = None, + reranker: RerankerProtocol | None = None, ) -> None: self._embedding = embedding self._vector_store = vector_store diff --git a/vektra-core/src/vektra_core/reranker.py b/vektra-core/src/vektra_core/reranker.py index 87db0be4..a8cde2bf 100644 --- a/vektra-core/src/vektra_core/reranker.py +++ b/vektra-core/src/vektra_core/reranker.py @@ -14,7 +14,9 @@ import asyncio import dataclasses import math +from typing import Protocol, runtime_checkable +import httpx import structlog from vektra_shared.config import RerankConfig @@ -49,6 +51,54 @@ class RerankResult: ] # (chunk_id, reranker_score, original_score) +@runtime_checkable +class RerankerProtocol(Protocol): + """Common interface of the in-process and remote reranker services.""" + + async def rerank( + self, + query: str, + results: list[SearchResult], + top_k: int, + ) -> RerankResult: ... + + +def _build_rerank_result( + results: list[SearchResult], + ordered: list[tuple[int, float]], + top_k: int, +) -> RerankResult: + """Build a RerankResult from (candidate_index, raw_score) pairs. + + Detects whether normalization is needed: FlashRank and TEI (with + raw_scores=false) produce sigmoid scores in [0, 1]; cross-encoder + logits can be negative or > 1. + """ + all_raw = [score for _, score in ordered] + needs_sigmoid = any(s < 0.0 or s > 1.0 for s in all_raw) + + all_scores: list[tuple[str, float, float]] = [] + reranked: list[SearchResult] = [] + + for idx, raw in ordered: + original = results[idx] + normalized = _sigmoid(raw) if needs_sigmoid else raw + all_scores.append( + (original.chunk_id, round(normalized, 4), round(original.score, 4)) + ) + + if len(reranked) < top_k: + reranked.append( + dataclasses.replace( + original, + score=normalized, + original_score=original.score, + ) + ) + + return RerankResult(top_k=reranked, all_scores=all_scores) + + class RerankerService: """Wraps the rerankers library for scoring and reordering search results.""" @@ -77,49 +127,78 @@ async def rerank( docs=docs, ) - # Score ALL candidates and detect whether normalization is needed. - # FlashRank produces sigmoid scores in [0, 1]; cross-encoder - # produces raw logits that can be negative or > 1. - all_items = ranked.results - all_raw = [float(item.score) for item in all_items] - needs_sigmoid = any(s < 0.0 or s > 1.0 for s in all_raw) - - all_scores: list[tuple[str, float, float]] = [] - reranked: list[SearchResult] = [] - - for item in all_items: - original = results[item.doc_id] - raw = float(item.score) - normalized = _sigmoid(raw) if needs_sigmoid else raw - all_scores.append( - (original.chunk_id, round(normalized, 4), round(original.score, 4)) - ) + ordered = [(item.doc_id, float(item.score)) for item in ranked.results] + return _build_rerank_result(results, ordered, top_k) - if len(reranked) < top_k: - reranked.append( - dataclasses.replace( - original, - score=normalized, - original_score=original.score, - ) - ) - return RerankResult(top_k=reranked, all_scores=all_scores) +class TEIRerankerService: + """Reranker backed by a TEI /rerank endpoint (FEAT-024). + + One TEI instance serves one reranker model (e.g. bge-reranker-v2-m3). + POST /rerank {"query", "texts", "raw_scores": false} returns + [{"index", "score"}] sorted by score descending, scores in [0, 1]. + """ + def __init__( + self, + *, + url: str, + api_key: str | None = None, + timeout_s: float = 30.0, + _client: httpx.AsyncClient | None = None, + ) -> None: + headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} + self._client = _client or httpx.AsyncClient( + base_url=url.rstrip("/"), headers=headers, timeout=timeout_s + ) + + async def rerank( + self, + query: str, + results: list[SearchResult], + top_k: int, + ) -> RerankResult: + """Rerank search results via the remote TEI cross-encoder.""" + if not results: + return RerankResult(top_k=[], all_scores=[]) + + resp = await self._client.post( + "/rerank", + json={ + "query": query, + "texts": [r.text_snippet for r in results], + "raw_scores": False, + }, + ) + resp.raise_for_status() + ranked = resp.json() + + ordered = [(int(item["index"]), float(item["score"])) for item in ranked] + return _build_rerank_result(results, ordered, top_k) -def create_reranker(config: RerankConfig) -> RerankerService | None: - """Create a RerankerService from config. Returns None if unavailable.""" + +def create_reranker(config: RerankConfig) -> RerankerProtocol | None: + """Create a reranker service from config. Returns None if unavailable.""" if not config.enabled: log.info("reranker_disabled") return None + if config.provider == "tei": + log.info("reranker_loaded", provider="tei", url=config.tei_url) + return TEIRerankerService(url=config.tei_url, api_key=config.tei_api_key) + model_type = _PROVIDER_TO_MODEL_TYPE.get(config.provider, config.provider) model_name = config.model or _default_model_for_provider(config.provider) try: from rerankers import Reranker # type: ignore[import-untyped] - ranker = Reranker(model_name, model_type=model_type, verbose=0) + # API-based providers (cohere) need the key passed through; + # without it the option was dead as wired (FEAT-024). + kwargs: dict[str, str] = {} + if config.api_key: + kwargs["api_key"] = config.api_key + ranker = Reranker(model_name, model_type=model_type, verbose=0, **kwargs) if ranker is None: log.warning( "reranker_init_failed", diff --git a/vektra-core/tests/test_reranker.py b/vektra-core/tests/test_reranker.py index ca631da4..59ef0db6 100644 --- a/vektra-core/tests/test_reranker.py +++ b/vektra-core/tests/test_reranker.py @@ -1,12 +1,17 @@ -"""Unit tests for RerankerService and create_reranker factory.""" +"""Unit tests for RerankerService, TEIRerankerService and create_reranker.""" +import json from types import SimpleNamespace from unittest.mock import MagicMock, patch from uuid import uuid4 +import httpx +import pytest + from vektra_core.reranker import ( RerankerService, RerankResult, + TEIRerankerService, _default_model_for_provider, _sigmoid, create_reranker, @@ -242,3 +247,73 @@ def test_sigmoid_extreme_values_no_overflow(): """Numerically stable sigmoid must not overflow on extreme logits.""" assert _sigmoid(1000.0) == 1.0 assert _sigmoid(-1000.0) == 0.0 + + +# --------------------------------------------------------------------------- +# TEIRerankerService (FEAT-024) +# --------------------------------------------------------------------------- + + +def _tei_client(handler) -> httpx.AsyncClient: + return httpx.AsyncClient( + transport=httpx.MockTransport(handler), base_url="http://tei.test" + ) + + +async def test_tei_rerank_empty_results(): + service = TEIRerankerService(url="http://tei.test", _client=_tei_client(None)) + result = await service.rerank("query", [], top_k=5) + assert result.top_k == [] + assert result.all_scores == [] + + +async def test_tei_rerank_orders_and_normalizes(): + results = [ + _make_result(0.5, "doc A"), + _make_result(0.3, "doc B"), + _make_result(0.9, "doc C"), + ] + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert request.url.path == "/rerank" + assert body["query"] == "test query" + assert body["texts"] == ["doc A", "doc B", "doc C"] + assert body["raw_scores"] is False + return httpx.Response( + 200, + json=[ + {"index": 2, "score": 0.95}, + {"index": 0, "score": 0.80}, + {"index": 1, "score": 0.10}, + ], + ) + + service = TEIRerankerService(url="http://tei.test", _client=_tei_client(handler)) + result = await service.rerank("test query", results, top_k=2) + + assert len(result.top_k) == 2 + assert result.top_k[0].chunk_id == results[2].chunk_id + assert result.top_k[0].score == 0.95 + assert result.top_k[0].original_score == 0.9 # BUG-015 preserved + assert result.top_k[1].chunk_id == results[0].chunk_id + assert len(result.all_scores) == 3 + + +async def test_tei_rerank_raises_on_http_error(): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(503) + + service = TEIRerankerService(url="http://tei.test", _client=_tei_client(handler)) + with pytest.raises(httpx.HTTPStatusError): + await service.rerank("q", [_make_result(0.5)], top_k=2) + + +def test_create_reranker_tei_provider(): + config = RerankConfig( + VEKTRA_RERANK_ENABLED=True, + VEKTRA_RERANK_PROVIDER="tei", + VEKTRA_RERANK_TEI_URL="http://tei.test", + ) + reranker = create_reranker(config) + assert isinstance(reranker, TEIRerankerService) diff --git a/vektra-index/pyproject.toml b/vektra-index/pyproject.toml index b891dd83..9bcb0cc9 100644 --- a/vektra-index/pyproject.toml +++ b/vektra-index/pyproject.toml @@ -13,6 +13,7 @@ dependencies = [ "sentence-transformers>=3.0", "torch>=2.2.0", "pgvector>=0.3", + "httpx>=0.27", ] [project.optional-dependencies] diff --git a/vektra-index/src/vektra_index/providers/qdrant.py b/vektra-index/src/vektra_index/providers/qdrant.py index a5028de9..71dbf39e 100644 --- a/vektra-index/src/vektra_index/providers/qdrant.py +++ b/vektra-index/src/vektra_index/providers/qdrant.py @@ -106,13 +106,29 @@ async def ensure_collection(self) -> None: """Create the collection if it doesn't exist. Called during startup validation. Configures named vectors - for dense and sparse search. + for dense and sparse search. If the collection already exists, + verifies its dense vector size matches the active embedding model + (FEAT-024): a silent mismatch would fail on every upsert/search + with an opaque Qdrant error, so fail fast with a clear message. """ from qdrant_client import models collections = await self._client.get_collections() existing = {c.name for c in collections.collections} if self._collection_name in existing: + info = await self._client.get_collection(self._collection_name) + vectors = info.config.params.vectors + dense = vectors.get("dense") if isinstance(vectors, dict) else None + existing_size = getattr(dense, "size", None) + if existing_size is not None and existing_size != self._dense_dimensions: + raise ValueError( + f"Qdrant collection '{self._collection_name}' has dense " + f"vectors of size {existing_size}, but the active embedding " + f"model produces {self._dense_dimensions} dimensions. " + "Changing the embedding model requires re-ingesting into a " + "new collection (set VEKTRA_QDRANT_COLLECTION) or deleting " + "the existing one." + ) return try: diff --git a/vektra-index/src/vektra_index/providers/tei.py b/vektra-index/src/vektra_index/providers/tei.py new file mode 100644 index 00000000..75f2b165 --- /dev/null +++ b/vektra-index/src/vektra_index/providers/tei.py @@ -0,0 +1,121 @@ +"""TEIEmbeddingProvider: EmbeddingProvider over HuggingFace Text Embeddings Inference. + +Remote embedding via a TEI instance (FEAT-024, ADR-0013). Lets deployments +reuse a shared inference service (e.g. bge-m3 on the host) instead of +duplicating in-process CPU embedding per container, and unlocks models +whose sequence window exceeds the in-process default (bge-m3: 8192 tokens +vs 128 for paraphrase-multilingual-MiniLM-L12-v2). + +Endpoints used (TEI native API): +- POST /embed {"inputs": [...]} -> [[...], ...] +- GET /info for the embedding size (with an /embed probe fallback) + +Auth: optional Bearer token (TEI --api-key). +""" + +from __future__ import annotations + +import logging +import time +from typing import Any + +import httpx + +from vektra_shared.types import HealthStatus + +logger = logging.getLogger(__name__) + +# TEI rejects batches larger than its --max-client-batch-size (default 32). +_MAX_BATCH = 32 + + +class TEIEmbeddingProvider: + """EmbeddingProvider backed by a remote TEI instance. + + Symmetric encoding: TEI applies the model's own pooling/normalization; + query and document paths use the same endpoint. + """ + + def __init__( + self, + *, + url: str, + api_key: str | None = None, + timeout_s: float = 30.0, + _client: httpx.AsyncClient | None = None, + ) -> None: + self._url = url.rstrip("/") + self._api_key = api_key + self._timeout_s = timeout_s + headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} + self._client = _client or httpx.AsyncClient( + base_url=self._url, headers=headers, timeout=timeout_s + ) + self._dimensions: int | None = None + + async def _embed_batch(self, texts: list[str]) -> list[list[float]]: + resp = await self._client.post("/embed", json={"inputs": texts}) + resp.raise_for_status() + data: list[list[float]] = resp.json() + return data + + async def embed_documents(self, texts: list[str]) -> list[list[float]]: + """Embed a batch of document passages, chunked to the TEI batch limit.""" + out: list[list[float]] = [] + for i in range(0, len(texts), _MAX_BATCH): + out.extend(await self._embed_batch(texts[i : i + _MAX_BATCH])) + return out + + async def embed_query(self, query: str) -> list[float]: + """Embed a single query string for retrieval.""" + return (await self._embed_batch([query]))[0] + + def dimensions(self) -> int: + """Return the embedding dimensionality, fetched once from the server. + + The Protocol method is synchronous, so this uses a one-off sync HTTP + call (startup/wiring path, not the query hot path). Tries /info + first; TEI versions that do not expose the size fall back to probing + /embed with a single input. + """ + if self._dimensions is not None: + return self._dimensions + + headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} + with httpx.Client( + base_url=self._url, headers=headers, timeout=self._timeout_s + ) as client: + self._dimensions = self._fetch_dimensions(client) + + logger.info("TEI embedding provider: %s (%d dims)", self._url, self._dimensions) + return self._dimensions + + @staticmethod + def _fetch_dimensions(client: httpx.Client) -> int: + try: + info: dict[str, Any] = client.get("/info").raise_for_status().json() + for key in ("embedding_size", "hidden_size"): + if isinstance(info.get(key), int): + return int(info[key]) + except httpx.HTTPStatusError: + logger.debug("TEI /info unavailable, probing /embed for dimensions") + probe = ( + client.post("/embed", json={"inputs": ["dim probe"]}) + .raise_for_status() + .json() + ) + return len(probe[0]) + + async def health_check(self) -> HealthStatus: + """Verify the TEI server responds and can produce an embedding.""" + try: + start = time.monotonic() + await self.embed_query("health check") + latency_ms = int((time.monotonic() - start) * 1000) + return HealthStatus(status="healthy", latency_ms=latency_ms) + except Exception as exc: + return HealthStatus(status="unhealthy", message=str(exc)) + + async def aclose(self) -> None: + """Release the underlying HTTP client (tests and shutdown).""" + await self._client.aclose() diff --git a/vektra-index/tests/test_tei_provider.py b/vektra-index/tests/test_tei_provider.py new file mode 100644 index 00000000..a733c48c --- /dev/null +++ b/vektra-index/tests/test_tei_provider.py @@ -0,0 +1,137 @@ +"""Unit tests for TEIEmbeddingProvider (FEAT-024). All HTTP mocked.""" + +from __future__ import annotations + +import json + +import httpx +import pytest + +from vektra_index.providers.tei import _MAX_BATCH, TEIEmbeddingProvider + + +def _embed_handler(dim: int = 4): + """MockTransport handler: /embed returns one vector per input.""" + calls: list[int] = [] + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/embed": + inputs = json.loads(request.content)["inputs"] + calls.append(len(inputs)) + return httpx.Response(200, json=[[0.1] * dim for _ in inputs]) + return httpx.Response(404) + + return handler, calls + + +def _make_provider(handler) -> TEIEmbeddingProvider: + client = httpx.AsyncClient( + transport=httpx.MockTransport(handler), base_url="http://tei.test" + ) + return TEIEmbeddingProvider(url="http://tei.test", _client=client) + + +async def test_embed_query_returns_single_vector(): + handler, _ = _embed_handler(dim=4) + provider = _make_provider(handler) + vec = await provider.embed_query("hello") + assert vec == [0.1] * 4 + await provider.aclose() + + +async def test_embed_documents_chunks_to_batch_limit(): + handler, calls = _embed_handler(dim=4) + provider = _make_provider(handler) + texts = [f"doc {i}" for i in range(_MAX_BATCH + 5)] + vectors = await provider.embed_documents(texts) + assert len(vectors) == _MAX_BATCH + 5 + assert calls == [_MAX_BATCH, 5] + await provider.aclose() + + +async def test_health_check_healthy(): + handler, _ = _embed_handler() + provider = _make_provider(handler) + status = await provider.health_check() + assert status.status == "healthy" + await provider.aclose() + + +async def test_health_check_unhealthy_on_http_error(): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(500, text="boom") + + provider = _make_provider(handler) + status = await provider.health_check() + assert status.status == "unhealthy" + await provider.aclose() + + +def test_fetch_dimensions_from_info(): + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/info": + return httpx.Response(200, json={"model_id": "m", "embedding_size": 1024}) + return httpx.Response(404) + + client = httpx.Client( + transport=httpx.MockTransport(handler), base_url="http://tei.test" + ) + assert TEIEmbeddingProvider._fetch_dimensions(client) == 1024 + + +def test_fetch_dimensions_probe_fallback_when_info_lacks_size(): + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/info": + return httpx.Response(200, json={"model_id": "m"}) + if request.url.path == "/embed": + return httpx.Response(200, json=[[0.0] * 768]) + return httpx.Response(404) + + client = httpx.Client( + transport=httpx.MockTransport(handler), base_url="http://tei.test" + ) + assert TEIEmbeddingProvider._fetch_dimensions(client) == 768 + + +def test_fetch_dimensions_probe_fallback_when_info_denied(): + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/info": + return httpx.Response(401) + if request.url.path == "/embed": + return httpx.Response(200, json=[[0.0] * 384]) + return httpx.Response(404) + + client = httpx.Client( + transport=httpx.MockTransport(handler), base_url="http://tei.test" + ) + assert TEIEmbeddingProvider._fetch_dimensions(client) == 384 + + +async def test_auth_header_sent_when_api_key_set(): + seen: dict[str, str] = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["auth"] = request.headers.get("authorization", "") + return httpx.Response(200, json=[[0.0] * 4]) + + client = httpx.AsyncClient( + transport=httpx.MockTransport(handler), + base_url="http://tei.test", + headers={"Authorization": "Bearer sekret"}, + ) + provider = TEIEmbeddingProvider( + url="http://tei.test", api_key="sekret", _client=client + ) + await provider.embed_query("q") + assert seen["auth"] == "Bearer sekret" + await provider.aclose() + + +async def test_embed_raises_on_server_error(): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(503) + + provider = _make_provider(handler) + with pytest.raises(httpx.HTTPStatusError): + await provider.embed_documents(["doc"]) + await provider.aclose() diff --git a/vektra-shared/src/vektra_shared/config.py b/vektra-shared/src/vektra_shared/config.py index d7cc1c4a..894a0030 100644 --- a/vektra-shared/src/vektra_shared/config.py +++ b/vektra-shared/src/vektra_shared/config.py @@ -88,6 +88,16 @@ class EmbeddingConfig(BaseSettings): alias="VEKTRA_SPARSE_EMBEDDING_MODEL", description="Sparse embedding model name. Phase 2 only.", ) + tei_url: str = Field( + "http://localhost:8080", + alias="VEKTRA_TEI_URL", + description="TEI server base URL (native API, no /v1 suffix). Used when embedding_provider='tei' (FEAT-024).", + ) + tei_api_key: str | None = Field( + None, + alias="VEKTRA_TEI_API_KEY", + description="Bearer token for the TEI embedding server (--api-key). Optional.", + ) model_config = SettingsConfigDict( env_prefix="", extra="ignore", populate_by_name=True @@ -158,7 +168,7 @@ class RerankConfig(BaseSettings): provider: str = Field( "cross-encoder", alias="VEKTRA_RERANK_PROVIDER", - description="Reranking provider: 'flashrank', 'cross-encoder', 'cohere'.", + description="Reranking provider: 'flashrank', 'cross-encoder', 'cohere', 'tei'.", ) model: str | None = Field( "BAAI/bge-reranker-v2-m3", @@ -171,6 +181,21 @@ class RerankConfig(BaseSettings): alias="VEKTRA_RERANK_TOP_K", description="Final top-k results after reranking.", ) + api_key: str | None = Field( + None, + alias="VEKTRA_RERANK_API_KEY", + description="API key for API-based rerank providers (e.g. 'cohere').", + ) + tei_url: str = Field( + "http://localhost:8080", + alias="VEKTRA_RERANK_TEI_URL", + description="TEI reranker server base URL (one TEI instance per model). Used when provider='tei' (FEAT-024).", + ) + tei_api_key: str | None = Field( + None, + alias="VEKTRA_RERANK_TEI_API_KEY", + description="Bearer token for the TEI reranker server. Optional.", + ) model_config = SettingsConfigDict( env_prefix="", extra="ignore", populate_by_name=True @@ -499,6 +524,8 @@ class VektraSettings(BaseSettings): sparse_embedding_model: str | None = Field( None, alias="VEKTRA_SPARSE_EMBEDDING_MODEL" ) + tei_url: str = Field("http://localhost:8080", alias="VEKTRA_TEI_URL") + tei_api_key: str | None = Field(None, alias="VEKTRA_TEI_API_KEY") # Vector store vector_store_provider: str = Field("pgvector", alias="VEKTRA_VECTOR_STORE_PROVIDER") From 5cca19cf40be525a4a7f144ef5be169e1038444e Mon Sep 17 00:00:00 2001 From: Francesco Vadicamo Date: Mon, 13 Jul 2026 01:49:40 +0000 Subject: [PATCH 2/4] fix(index): look up the active embedding provider in the startup warmup check_embedding_model resolved the provider by the hardcoded 'sentence-transformers' name, so startup failed with a registration error whenever VEKTRA_EMBEDDING_PROVIDER=tei was active (found in the FEAT-024 live smoke). The 'default' alias always points to the active provider. Co-Authored-By: Claude Fable 5 --- vektra-index/src/vektra_index/startup.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vektra-index/src/vektra_index/startup.py b/vektra-index/src/vektra_index/startup.py index 1ffd8575..ab4a6d2c 100644 --- a/vektra-index/src/vektra_index/startup.py +++ b/vektra-index/src/vektra_index/startup.py @@ -52,7 +52,9 @@ async def check_provider_registration( async def check_embedding_model(registry: Any) -> None: """ARCH-057 step 6: warm up the embedding model and verify dimensionality.""" try: - embedding_provider = registry.get("embedding", "sentence-transformers") + # "default" aliases whichever provider is active + # (sentence-transformers or tei, FEAT-024). + embedding_provider = registry.get("embedding", "default") test_embedding = await embedding_provider.embed_query("startup validation test") actual_dims = len(test_embedding) expected_dims = embedding_provider.dimensions() From 390ddcf23eab13876d74b012534cc9c858f7fa60 Mon Sep 17 00:00:00 2001 From: Francesco Vadicamo Date: Mon, 13 Jul 2026 02:16:59 +0000 Subject: [PATCH 3/4] docs(backlog): mark FEAT-024 completed with comparison numbers bge-m3 via TEI on the eval-full questions (same dual chunks): retrieval hit 82.6% -> 93.5%, MRR 0.7029 -> 0.8478, e2e stable. Configuration reference and .env.example cover the new provider options; the reindex pgvector-hardcoding discovery is recorded in the BUG-021 family note. Co-Authored-By: Claude Fable 5 --- .env.example | 8 ++++++++ .s2s/BACKLOG.md | 17 ++++++++++------- .s2s/plans/20260712-sprint3-rag-quality.md | 22 ++++++++++++++++++++++ docs/reference/configuration.md | 11 ++++++++--- 4 files changed, 48 insertions(+), 10 deletions(-) diff --git a/.env.example b/.env.example index e576fa2d..01d93810 100644 --- a/.env.example +++ b/.env.example @@ -64,6 +64,9 @@ # Embedding provider and model. # VEKTRA_EMBEDDING_PROVIDER=sentence-transformers # VEKTRA_EMBEDDING_MODEL=paraphrase-multilingual-MiniLM-L12-v2 +# Remote embedding via Text Embeddings Inference (VEKTRA_EMBEDDING_PROVIDER=tei): +# VEKTRA_TEI_URL=http://localhost:8080 +# VEKTRA_TEI_API_KEY= # -------------------------------------------------------------------------- # Vector store @@ -106,6 +109,11 @@ # VEKTRA_RERANK_PROVIDER=cross-encoder # VEKTRA_RERANK_MODEL=BAAI/bge-reranker-v2-m3 # VEKTRA_RERANK_TOP_K=5 +# API key for API-based providers (cohere): +# VEKTRA_RERANK_API_KEY= +# Remote reranking via TEI (VEKTRA_RERANK_PROVIDER=tei, one instance per model): +# VEKTRA_RERANK_TEI_URL=http://localhost:8080 +# VEKTRA_RERANK_TEI_API_KEY= # -------------------------------------------------------------------------- # Ingestion diff --git a/.s2s/BACKLOG.md b/.s2s/BACKLOG.md index 3bd8e9fb..c8f9feca 100644 --- a/.s2s/BACKLOG.md +++ b/.s2s/BACKLOG.md @@ -393,8 +393,9 @@ The `title` field would contain `filename + page` (e.g., "Costituzione italiana. ### FEAT-024: Remote embedding and reranker providers (TEI) -**Status**: planned | **Priority**: medium | **Created**: 2026-07-12 +**Status**: completed | **Priority**: medium | **Created**: 2026-07-12 | **Completed**: 2026-07-13 | **PR**: #93 **Origin**: deployment modularity review 2026-07-12 - the host workstation already serves TEI instances (bge-m3, qwen3-embedding); Vektra cannot use them. +**Evidence**: `vektra-internal/stack/20260713-feat024-tei-providers.md` (+ artifacts in `20260713-feat024-eval-artifacts/`) **Context**: `VEKTRA_EMBEDDING_PROVIDER` documents a `tei` option (config.py:74) but **no TEI provider exists**: `main.py:129-137` unconditionally instantiates in-process `SentenceTransformersProvider`; the compose even ships a `tei` profile service nobody can talk to. The reranker likewise runs in-process only (`rerankers` lib; the `cohere` path never passes an api_key, so it is dead as wired - reranker.py:122). Consequences: every Vektra instance duplicates embedding/reranker compute in-container (CPU), and shared GPU/CPU inference services on the host cannot be reused. @@ -406,14 +407,16 @@ The `title` field would contain `filename + page` (e.g., "Costituzione italiana. **Why it matters beyond dedup**: the current embedding model (paraphrase-multilingual-MiniLM-L12-v2) has **max_seq_length 128 tokens** - our 500-token chunks are silently truncated at embedding time (dense sees only the chunk head; BM25 sees the full text). bge-m3 (8192-token window, MIRACL dense nDCG@10 69.2 vs mE5-large 66.6; MiniLM sits 16-22 nDCG points below even mE5 on European-language retrieval per PL-MTEB) is the natural upgrade candidate, testable via TEI without fattening the container. +**Resolution (2026-07-13)**: implemented as designed with two deltas: (1) TEI 1.9.3 `/info` does not expose the embedding size, so `dimensions()` probes `/embed` as fallback (both paths unit-tested); (2) found and fixed in passing a startup blocker: `check_embedding_model` resolved the provider by the hardcoded `sentence-transformers` name, so startup failed with any other provider (now uses the `default` alias). Measured on eval-full questions (same dual chunks, reindexed via a fresh `eval-tei` namespace): **bge-m3 via TEI retrieval hit 93.5% / MRR 0.8478 vs MiniLM dual 82.6% / 0.7029 (+10.9pp)** - beats even the fixed-chunking MiniLM baseline (89.1%/0.8062), confirming the 128-token truncation hypothesis; e2e grounded 54/55 stable, MC answers improve in substance (MC-02 produces a real bi-document comparison; kw 7/25 vs 4/25), p50 +0.7s (TEI on CPU). TEI reranker smoke: factual query scores 0.75 (2 survive the threshold), comparative query all-below-threshold rescued by TECH-007 (`rescued=3`) - full funnel verified with both remote providers, authenticated. Discovery filed under BUG-021: `run_reindex` stores through hardcoded pgvector, so reindex-into-Qdrant silently writes nothing (worked around via fresh-namespace ingest). Switching the default embedding to bge-m3 is a separate decision (needs full corpus re-ingest and a TECH-005-grade bench). + **Traceability**: ADR-0013 (EmbeddingProvider Protocol), ARCH-035, ARCH-036, ADR-0021 **Acceptance criteria**: -- [ ] `VEKTRA_EMBEDDING_PROVIDER=tei` works end-to-end (ingest + query) against a TEI instance with api key -- [ ] Collection created with the provider's real dimensions; clear error on dimension mismatch with an existing collection -- [ ] `VEKTRA_RERANK_PROVIDER=tei` reranks via TEI /rerank with scores compatible with the threshold filter -- [ ] Embedding-model comparison (MiniLM in-process vs bge-m3 via TEI) run with the TECH-005/existing harness and recorded -- [ ] Docs: configuration.md + .env.example cover the new provider options +- [x] `VEKTRA_EMBEDDING_PROVIDER=tei` works end-to-end (ingest + query) against a TEI instance with api key +- [x] Collection created with the provider's real dimensions; clear error on dimension mismatch with an existing collection (verified live: 384-vs-1024 startup warning with remediation) +- [x] `VEKTRA_RERANK_PROVIDER=tei` reranks via TEI /rerank with scores compatible with the threshold filter +- [x] Embedding-model comparison (MiniLM in-process vs bge-m3 via TEI) run with the TECH-005/existing harness and recorded +- [x] Docs: configuration.md + .env.example cover the new provider options --- @@ -1087,7 +1090,7 @@ Three near-duplicates is the threshold where extraction starts to pay off (a fou **Resolution**: the endpoint now resolves embedding, sparse embedding, and vector store from `request.app.state.registry` (same contract as the pipeline). The per-request `SentenceTransformersProvider` instantiation and the now-unused `session` dependency were removed. Unit tests added (`vektra-index/tests/test_api_search.py`): registry resolution, hybrid→dense fallback without sparse, hybrid with sparse. -**Same family, not fixed here**: `POST /documents/{id}/chunks`, `DELETE /documents/{id}` and `GET /stats` still hardcode pgvector (the stats-vs-Qdrant mismatch was already a known issue). Track separately if needed. +**Same family, not fixed here**: `POST /documents/{id}/chunks`, `DELETE /documents/{id}` and `GET /stats` still hardcode pgvector (the stats-vs-Qdrant mismatch was already a known issue). Also `run_reindex` (`vektra-index/reindex.py`): it re-embeds with the registry's active embedding provider but stores through a hardcoded `PgvectorProvider`, so in Qdrant mode a reindex reports "completed" while the Qdrant collection receives nothing (found during the FEAT-024 live smoke, 2026-07-13: reindexing eval-full to v2 wrote to Postgres only). Track separately if needed. **Traceability**: ARCH-039 (ProviderRegistry), ARCH-051 (full-store contract), TECH-002 (eval harness) diff --git a/.s2s/plans/20260712-sprint3-rag-quality.md b/.s2s/plans/20260712-sprint3-rag-quality.md index 86ae3f93..938f9de4 100644 --- a/.s2s/plans/20260712-sprint3-rag-quality.md +++ b/.s2s/plans/20260712-sprint3-rag-quality.md @@ -323,3 +323,25 @@ the bundle and a manual smoke in the Moodle dev stack. `vektra-internal/stack/20260713-tech007-retrieval-rescue.md`; per-question artifacts in `20260713-tech007-eval-artifacts/`. Dev `.env` adds `VEKTRA_RETRIEVAL_RESCUE_TOP_K=3`, `VEKTRA_RETRIEVAL_RESCUE_FLOOR=0.005`. +- 2026-07-13: **FEAT-024 implemented and measured (post-sprint, PR #93)** — + TEI remote providers, `feat/feat-024-tei-providers`. TEIEmbeddingProvider + (native /embed, /info dimensions with /embed probe fallback — TEI 1.9.3 + does not expose the size in /info), TEIRerankerService (/rerank, same + RerankResult semantics), Qdrant collection sized from the active provider + (fixes the hardcoded-384 latent bug; live 384-vs-1024 mismatch produces a + clear startup warning), cohere api_key pass-through fixed, startup warmup + fixed to resolve the `default` embedding alias (was hardcoded to + 'sentence-transformers' and broke any alternative provider). + **Comparison on eval-full questions (same dual chunks, ns eval-tei, + bge-m3 float32 CPU via TEI): retrieval hit 82.6% -> 93.5%, MRR 0.7029 -> + 0.8478 — beats even the fixed-chunking MiniLM baseline (89.1%/0.8062), + confirming MiniLM's 128-token truncation as a real retrieval cap.** E2e + grounded 54/55 stable, multi-chunk substance improves (MC-02 real + bi-document comparison, kw 7/25 vs 4/25), p50 +0.7s (CPU TEI). Reranker + smoke: factual 0.75 survives threshold, comparative all-below rescued by + TECH-007 (rescued=3) — full funnel verified with both remote providers. + Found in passing (filed under BUG-021 family): run_reindex stores through + hardcoded pgvector — reindex-into-Qdrant silently no-ops. Dev stack + restored to reference config (MiniLM + rescue); eval-tei namespace and + vektra-tei collection (105 points, 1024-dim) left in place for follow-ups; + bge-reranker TEI cache kept at /mnt/scratch/tei-rerank-cache. diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index 1a50360d..72a32279 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -59,8 +59,10 @@ The model name must match the vLLM `--model` path exactly (e.g., `/models/qwen35 | Variable | Type | Default | Description | |----------|------|---------|-------------| -| `VEKTRA_EMBEDDING_PROVIDER` | str | `sentence-transformers` | Embedding provider implementation | -| `VEKTRA_EMBEDDING_MODEL` | str | `paraphrase-multilingual-MiniLM-L12-v2` | Model name within the selected provider | +| `VEKTRA_EMBEDDING_PROVIDER` | str | `sentence-transformers` | Embedding provider implementation: `sentence-transformers` (in-process), `tei` (remote) | +| `VEKTRA_EMBEDDING_MODEL` | str | `paraphrase-multilingual-MiniLM-L12-v2` | Model name within the selected provider (`sentence-transformers` only; a TEI instance serves one fixed model) | +| `VEKTRA_TEI_URL` | str | `http://localhost:8080` | TEI server base URL (native API, no `/v1` suffix). Used when provider is `tei`. The Qdrant collection is sized from the served model's dimensions at startup | +| `VEKTRA_TEI_API_KEY` | str | - | Bearer token for the TEI embedding server (`--api-key`). Optional | | `VEKTRA_SPARSE_EMBEDDING_PROVIDER` | str | - | Sparse embedding provider: `fastembed-bm25`, `splade` | | `VEKTRA_SPARSE_EMBEDDING_MODEL` | str | - | Sparse embedding model name | @@ -101,9 +103,12 @@ The model name must match the vLLM `--model` path exactly (e.g., `/models/qwen35 | Variable | Type | Default | Description | |----------|------|---------|-------------| | `VEKTRA_RERANK_ENABLED` | bool | `true` | Enable cross-encoder reranking after retrieval | -| `VEKTRA_RERANK_PROVIDER` | str | `cross-encoder` | Reranking provider: `flashrank`, `cross-encoder`, `cohere` | +| `VEKTRA_RERANK_PROVIDER` | str | `cross-encoder` | Reranking provider: `flashrank`, `cross-encoder`, `cohere`, `tei` (remote) | | `VEKTRA_RERANK_MODEL` | str | `BAAI/bge-reranker-v2-m3` | Multilingual reranking model. For English-only lightweight deployments: provider=`flashrank`, model=`ms-marco-MiniLM-L-12-v2` | | `VEKTRA_RERANK_TOP_K` | int | `5` | Final top-k results after reranking | +| `VEKTRA_RERANK_API_KEY` | str | - | API key for API-based providers (`cohere`) | +| `VEKTRA_RERANK_TEI_URL` | str | `http://localhost:8080` | TEI reranker server base URL (one TEI instance per model, e.g. serving `BAAI/bge-reranker-v2-m3`). Used when provider is `tei`. Scores are sigmoid-normalized like the in-process path | +| `VEKTRA_RERANK_TEI_API_KEY` | str | - | Bearer token for the TEI reranker server. Optional | ## Ingestion From dfbc961bab40b8167f96041208d8ff6ca2d257cc Mon Sep 17 00:00:00 2001 From: Francesco Vadicamo Date: Mon, 13 Jul 2026 02:32:49 +0000 Subject: [PATCH 4/4] fix(rag): harden TEI providers per review - create_reranker: the tei branch now degrades to None on init/URL errors like the in-process path, and logs the URL redacted (coderabbit #3567707446, #3567707458) - TEIRerankerService gains aclose(); the app lifespan teardown closes remote provider clients best-effort via a new reranker registry entry (coderabbit #3567707442, #3567707464) - _embed_batch warms the dimensions cache so the startup warmup makes the synchronous dimensions() call free (coderabbit #3567707461) - _fetch_dimensions tolerates non-JSON /info responses and falls back to the /embed probe (gemini #3567695402) - ensure_collection reads the dense vector params via getattr fallback for non-dict vector configs (gemini #3567695404) Co-Authored-By: Claude Fable 5 --- vektra-app/src/vektra_app/main.py | 15 ++++++++++ vektra-core/src/vektra_core/reranker.py | 24 ++++++++++++++-- vektra-core/tests/test_reranker.py | 10 +++++++ .../src/vektra_index/providers/qdrant.py | 6 +++- .../src/vektra_index/providers/tei.py | 28 +++++++++++-------- 5 files changed, 68 insertions(+), 15 deletions(-) diff --git a/vektra-app/src/vektra_app/main.py b/vektra-app/src/vektra_app/main.py index f112f9a8..30e8f1bb 100644 --- a/vektra-app/src/vektra_app/main.py +++ b/vektra-app/src/vektra_app/main.py @@ -267,6 +267,9 @@ async def _step_5_register_providers( pipeline_config = QueryPipelineConfig() reranker = create_reranker(pipeline_config.rerank) + if reranker is not None: + # Registered so the lifespan teardown can close remote clients (FEAT-024) + registry.register("reranker", "default", reranker) # --- Query pipeline (Phase 2: select simple or advanced) --- templates_dir = ( @@ -559,6 +562,18 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: # Shutdown log.info("shutdown_started") + # Close remote provider HTTP clients (TEI, FEAT-024) best-effort + for category in ("embedding", "reranker"): + try: + provider = registry.get(category, "default") + except Exception: + continue + aclose = getattr(provider, "aclose", None) + if aclose is not None: + try: + await aclose() + except Exception as exc: + log.warning("provider_close_failed", category=category, error=str(exc)) from vektra_shared.db import get_engine engine = get_engine() diff --git a/vektra-core/src/vektra_core/reranker.py b/vektra-core/src/vektra_core/reranker.py index a8cde2bf..eba02e74 100644 --- a/vektra-core/src/vektra_core/reranker.py +++ b/vektra-core/src/vektra_core/reranker.py @@ -15,6 +15,7 @@ import dataclasses import math from typing import Protocol, runtime_checkable +from urllib.parse import urlparse import httpx import structlog @@ -32,6 +33,14 @@ } +def _redact_url(url: str) -> str: + """Return scheme://hostname:port only, stripping credentials and path.""" + parsed = urlparse(url) + host = parsed.hostname or "" + port = f":{parsed.port}" if parsed.port else "" + return f"{parsed.scheme}://{host}{port}" + + def _sigmoid(x: float) -> float: """Numerically stable sigmoid for cross-encoder logits.""" if x >= 0: @@ -176,6 +185,10 @@ async def rerank( ordered = [(int(item["index"]), float(item["score"])) for item in ranked] return _build_rerank_result(results, ordered, top_k) + async def aclose(self) -> None: + """Release the underlying HTTP client (tests and shutdown).""" + await self._client.aclose() + def create_reranker(config: RerankConfig) -> RerankerProtocol | None: """Create a reranker service from config. Returns None if unavailable.""" @@ -184,8 +197,15 @@ def create_reranker(config: RerankConfig) -> RerankerProtocol | None: return None if config.provider == "tei": - log.info("reranker_loaded", provider="tei", url=config.tei_url) - return TEIRerankerService(url=config.tei_url, api_key=config.tei_api_key) + try: + service = TEIRerankerService(url=config.tei_url, api_key=config.tei_api_key) + log.info("reranker_loaded", provider="tei", url=_redact_url(config.tei_url)) + except Exception as exc: + # Same graceful-degradation contract as the in-process path: + # a bad reranker config must not abort startup. + log.warning("reranker_init_failed", provider="tei", error=str(exc)) + return None + return service model_type = _PROVIDER_TO_MODEL_TYPE.get(config.provider, config.provider) model_name = config.model or _default_model_for_provider(config.provider) diff --git a/vektra-core/tests/test_reranker.py b/vektra-core/tests/test_reranker.py index 59ef0db6..9621653a 100644 --- a/vektra-core/tests/test_reranker.py +++ b/vektra-core/tests/test_reranker.py @@ -317,3 +317,13 @@ def test_create_reranker_tei_provider(): ) reranker = create_reranker(config) assert isinstance(reranker, TEIRerankerService) + + +def test_create_reranker_tei_invalid_url_returns_none(): + """A malformed TEI URL degrades to None instead of aborting startup.""" + config = RerankConfig( + VEKTRA_RERANK_ENABLED=True, + VEKTRA_RERANK_PROVIDER="tei", + VEKTRA_RERANK_TEI_URL="http://[invalid", + ) + assert create_reranker(config) is None diff --git a/vektra-index/src/vektra_index/providers/qdrant.py b/vektra-index/src/vektra_index/providers/qdrant.py index 71dbf39e..9e073c16 100644 --- a/vektra-index/src/vektra_index/providers/qdrant.py +++ b/vektra-index/src/vektra_index/providers/qdrant.py @@ -118,7 +118,11 @@ async def ensure_collection(self) -> None: if self._collection_name in existing: info = await self._client.get_collection(self._collection_name) vectors = info.config.params.vectors - dense = vectors.get("dense") if isinstance(vectors, dict) else None + dense = ( + vectors.get("dense") + if isinstance(vectors, dict) + else getattr(vectors, "dense", None) + ) existing_size = getattr(dense, "size", None) if existing_size is not None and existing_size != self._dense_dimensions: raise ValueError( diff --git a/vektra-index/src/vektra_index/providers/tei.py b/vektra-index/src/vektra_index/providers/tei.py index 75f2b165..ffead977 100644 --- a/vektra-index/src/vektra_index/providers/tei.py +++ b/vektra-index/src/vektra_index/providers/tei.py @@ -57,6 +57,10 @@ async def _embed_batch(self, texts: list[str]) -> list[list[float]]: resp = await self._client.post("/embed", json={"inputs": texts}) resp.raise_for_status() data: list[list[float]] = resp.json() + if data and self._dimensions is None: + # Warm the dimensions cache so the startup warmup (embed_query) + # makes the later synchronous dimensions() call free. + self._dimensions = len(data[0]) return data async def embed_documents(self, texts: list[str]) -> list[list[float]]: @@ -93,18 +97,18 @@ def dimensions(self) -> int: @staticmethod def _fetch_dimensions(client: httpx.Client) -> int: try: - info: dict[str, Any] = client.get("/info").raise_for_status().json() - for key in ("embedding_size", "hidden_size"): - if isinstance(info.get(key), int): - return int(info[key]) - except httpx.HTTPStatusError: - logger.debug("TEI /info unavailable, probing /embed for dimensions") - probe = ( - client.post("/embed", json={"inputs": ["dim probe"]}) - .raise_for_status() - .json() - ) - return len(probe[0]) + resp = client.get("/info") + resp.raise_for_status() + info: Any = resp.json() + if isinstance(info, dict): + for key in ("embedding_size", "hidden_size"): + if isinstance(info.get(key), int): + return int(info[key]) + except (httpx.HTTPError, ValueError): + logger.debug("TEI /info unavailable or invalid, probing /embed") + resp = client.post("/embed", json={"inputs": ["dim probe"]}) + resp.raise_for_status() + return len(resp.json()[0]) async def health_check(self) -> HealthStatus: """Verify the TEI server responds and can produce an embedding."""