diff --git a/REFACTORING_DECISION_LOG.md b/REFACTORING_DECISION_LOG.md index 6a7e2e352..4685f4539 100644 --- a/REFACTORING_DECISION_LOG.md +++ b/REFACTORING_DECISION_LOG.md @@ -548,6 +548,102 @@ producing `AttributeError: 'TranscriberConfig' object has no attribute --- +## Phase 6B — vLLM inference clients + legacy shims (2026-05-07) + +**1. `VLLMVision(VLLMClient, VLM)` — multiple inheritance kept for nominal typing.** +`VLLMClient` provides the full implementation (httpx pool, retry, +circuit breaker, `aclose()`). `VLM` is a pure abstract mixin with no +conflicting methods, so the MRO is linear and clean. Adds only +`_max_tokens`, `caption_image()`, and `caption_images_batch()`. +- Why: VLM and LLM talk to the same vLLM OpenAI-compatible + chat/completions endpoint, so `VLLMClient` is the right concrete + base. Keeping `VLM` in the bases preserves nominal typing — + `isinstance(vision, VLM)` works, and any future code that type-checks + against the VLM ABC will accept `VLLMVision` without a cast. +- Alternative considered: single inheritance `VLLMVision(VLLMClient)` + only, relying on structural/duck typing for registry lookup. Rejected + — the registry is currently structurally typed, but explicit ABC + conformance is cheap here (no diamond, no conflicting methods) and + makes the intent clear to readers. + +**2. `LLM.generate()` and `LLM.chat()` return `dict` (full OpenAI-compatible response body), not `str`.** +The original ABC typed both methods as `→ str`, which forced callers to +re-construct the surrounding OpenAI envelope when building RAG answers +(losing `model`, `usage`, `finish_reason`, etc.). The concrete vLLM +implementation already returned the full `httpx` JSON body; the `str` +annotation was aspirational, not real. +- Why: RAG answers are ultimately forwarded to the client in OpenAI format. + Stripping to plain text at the LLM boundary means the pipeline has to + re-wrap the content into `{"choices": [{"message": {"content": …}}]}` + further up — metadata (token counts, model id, stop reason) is lost in + the process. Returning `dict` preserves the full payload and keeps + back-ends interchangeable without wrapping shims. Using bare `dict` (not + a `TypedDict`) is a deliberate first step: it is backward-compatible with + all current callers and eases compat-shim re-exports while the + refactoring is still ongoing. +- Alternative considered: introduce typed response models (`ChatCompletion`, + `CompletionResponse`, `ChatCompletionChunk`) immediately. Rejected as + premature — Phase 6 adds the concrete client; Phase 10 (API layer + clean-up) is the right time to freeze the contract with typed models. + The `dict` annotation signals intent without coupling every caller to a + model definition that will evolve. +- `stream_chat` stays `AsyncIterator[str]` yielding raw SSE lines + (`data: {…}` strings). Parsing SSE chunks into typed dicts is Phase 10+ + work; the current shape keeps the streaming path consistent with + OpenAI's SDK behaviour. +- Future normalisation — when the typed models land, callers will migrate + to this pattern (TypedDict shown; Pydantic models are equally valid and + would expose `chat_content` / `completion_text` as properties instead): + +```python +from typing import TypedDict + +class _Message(TypedDict): + role: str + content: str + +class _Choice(TypedDict): + index: int + message: _Message # chat completions + finish_reason: str | None + +class _CompletionChoice(TypedDict): + index: int + text: str # text completions + finish_reason: str | None + +class _Usage(TypedDict): + prompt_tokens: int + completion_tokens: int + total_tokens: int + +class ChatCompletion(TypedDict): + id: str + object: str # "chat.completion" + model: str + choices: list[_Choice] + usage: _Usage + +class Completion(TypedDict): + id: str + object: str # "text_completion" + model: str + choices: list[_CompletionChoice] + usage: _Usage + +# Convenience extractors at the pipeline boundary: +def chat_content(resp: ChatCompletion) -> str: + return resp["choices"][0]["message"]["content"] + +def completion_text(resp: Completion) -> str: + return resp["choices"][0]["text"] +``` + + Until then, callers that need the text can use + `resp["choices"][0]["message"]["content"]` directly. + +--- + ## Template for future entries ``` diff --git a/openrag/components/indexer/chunker/chunker.py b/openrag/components/indexer/chunker/chunker.py index ec11da4a3..939afa125 100644 --- a/openrag/components/indexer/chunker/chunker.py +++ b/openrag/components/indexer/chunker/chunker.py @@ -6,7 +6,7 @@ Scheduled for removal in Phase 12. """ -from typing import ClassVar, Literal +from typing import Any, ClassVar, Literal # Side-effect import: pre-loads the indexer-utils submodule so the legacy # circular import between `components.utils` and `components.indexer.utils.files` @@ -52,6 +52,9 @@ async def chat(self, messages: list[dict[str, str]], **kwargs) -> str: out = await self._llm.ainvoke(lc_msgs) return out.content + async def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> Any: + pass # Not implemented since the contextualizer never streams. + def _chunks_to_documents(chunks: list, base_metadata: dict) -> list[Document]: """Convert a list of core domain Chunks into legacy LangChain Documents. diff --git a/openrag/components/indexer/embeddings/__init__.py b/openrag/components/indexer/embeddings/__init__.py index 69850e74f..a1d35852c 100644 --- a/openrag/components/indexer/embeddings/__init__.py +++ b/openrag/components/indexer/embeddings/__init__.py @@ -1,5 +1,7 @@ +from services.inference.vllm_client import VLLMEmbedder # noqa: F401 + from .base import BaseEmbedding -from .openai import OpenAIEmbedding +from .openai import _ShimOpenAIEmbedding as OpenAIEmbedding EMBEDDER_MAPPING = { "openai": OpenAIEmbedding, diff --git a/openrag/components/indexer/embeddings/openai.py b/openrag/components/indexer/embeddings/openai.py index 23aafd904..2c23e6f90 100644 --- a/openrag/components/indexer/embeddings/openai.py +++ b/openrag/components/indexer/embeddings/openai.py @@ -1,6 +1,17 @@ +"""Backward-compatibility shim — delegates to services.inference.vllm_client. + +All new code should import directly from ``services.inference.vllm_client``. +""" + +import asyncio +from concurrent.futures import ThreadPoolExecutor + +import httpx import openai +from core.config.endpoints import EmbedderConfig from langchain_core.documents.base import Document from openai import OpenAI +from services.inference.vllm_client import VLLMEmbedder # noqa: F401 from utils.exceptions.embeddings import * from utils.logger import get_logger @@ -9,7 +20,75 @@ logger = get_logger() +_SYNC_POOL = ThreadPoolExecutor(max_workers=1) + + +def _run_sync(coro): + """Run an async coroutine from sync code, safe inside a running event loop (e.g. Ray).""" + return _SYNC_POOL.submit(asyncio.run, coro).result() + + +def _normalize_texts(texts: list[str | Document]) -> list[str]: + return [item.page_content if isinstance(item, Document) else item for item in texts] + + +class _ShimOpenAIEmbedding(BaseEmbedding): + """Legacy shim — delegates to ``VLLMEmbedder`` for actual HTTP transport. + + Preserves the sync ``embed_documents``/``embed_query`` contract expected by + ``vectordb.py`` (via LangChain's ``aembed_documents`` thread wrapper) while + using VLLMEmbedder's long-lived async httpx pool under the hood. + """ + + def __init__(self, embeddings_config: EmbedderConfig): + self._delegate = VLLMEmbedder( + endpoint=embeddings_config.base_url, + model_name=embeddings_config.model_name, + max_model_len=embeddings_config.max_model_len, + api_key=embeddings_config.api_key, + ) + + @property + def embedding_dimension(self) -> int: + # Probe once if unknown — legacy callers (e.g. MilvusDB schema creation) read + # this before any embed() call, but VLLMEmbedder only learns its dimension + # from a real response. The probe must run on a one-off sync httpx.Client: + # asyncio.run() here would tear down the loop and leave the delegate's + # long-lived AsyncClient pool with stale connections, breaking the next + # real async call with "Event loop is closed". + try: + return self._delegate.dimension + except RuntimeError: + pass + body: dict = {"model": self._delegate._model, "input": ["dim-probe"]} + if self._delegate._max_model_len is not None: + body["truncate_prompt_tokens"] = self._delegate._max_model_len + with httpx.Client(timeout=30.0, headers=dict(self._delegate._client.headers)) as client: + resp = client.post(f"{self._delegate._endpoint}/embeddings", json=body) + resp.raise_for_status() + self._delegate._dimension = len(resp.json()["data"][0]["embedding"]) + return self._delegate._dimension + + def embed_documents(self, texts: list[str | Document]) -> list[list[float]]: + if not texts: + return [] + return _run_sync(self._delegate.embed(_normalize_texts(texts))) + + async def aembed_documents(self, texts: list[str | Document]) -> list[list[float]]: + if not texts: + return [] + return await self._delegate.embed(_normalize_texts(texts)) + + def embed_query(self, text: str) -> list[float]: + return _run_sync(self._delegate.embed_single(text)) + + async def aembed_query(self, text: str) -> list[float]: + return await self._delegate.embed_single(text) + + class OpenAIEmbedding(BaseEmbedding): + """Legacy OpenAI embedding wrapper. New code should use VLLMEmbedder (via DI).""" + def __init__(self, embeddings_config): self.embedding_model = embeddings_config.model_name self.base_url = embeddings_config.base_url @@ -20,16 +99,12 @@ def __init__(self, embeddings_config): @property def embedding_dimension(self) -> int: try: - # Test call to get embedding dimension output = self.embed_documents([Document(page_content="test")]) return len(output[0]) except Exception: raise def embed_documents(self, texts: list[str | Document]) -> list[list[float]]: - """ - Embed documents using the configured embedder. - """ if isinstance(texts[0], Document): texts = [doc.page_content for doc in texts] @@ -69,9 +144,6 @@ def embed_documents(self, texts: list[str | Document]) -> list[list[float]]: ) def embed_query(self, text: str) -> list[float]: - """ - Embed a query using the configured embedder. - """ try: output = self.embed_documents([Document(page_content=text)]) return output[0] diff --git a/openrag/components/llm.py b/openrag/components/llm.py index bfaf4ed76..e4151879f 100644 --- a/openrag/components/llm.py +++ b/openrag/components/llm.py @@ -1,13 +1,59 @@ +"""Backward-compatibility shim — delegates to services.inference.vllm_client. + +All new code should import directly from ``services.inference.vllm_client``. +""" + import copy import json +import warnings import httpx +from config.models import LLMConfig +from services.inference.vllm_client import VLLMClient # noqa: F401 from utils.logger import get_logger logger = get_logger() +class _LLMShim: + """Legacy shim — delegates to ``VLLMClient`` for retry, circuit breaker, + and connection pooling while preserving the generator-based interface.""" + + def __init__(self, llm_config: LLMConfig, logger=None): + warnings.warn( + "components.llm.LLM is deprecated — use services.inference.vllm_client.VLLMClient", + DeprecationWarning, + stacklevel=2, + ) + self.logger = logger + config_kwargs = {k: v for k, v in llm_config.model_dump().items() if k not in ("api_key", "base_url", "model")} + self._delegate = VLLMClient( + endpoint=llm_config.base_url, + model_name=llm_config.model, + api_key=llm_config.api_key, + **config_kwargs, + ) + + async def completions(self, request: dict): + prompt = request.pop("prompt") + response = await self._delegate.generate(prompt, **request) + yield response + + async def chat_completion(self, request: dict): + messages = request.pop("messages") + stream = request.pop("stream", False) + + if stream: + async for line in self._delegate.stream_chat(messages, **request): + yield line + else: + resp_dict = await self._delegate.chat(messages, **request) + yield resp_dict + + class LLM: + """Legacy LLM wrapper. New code should use VLLMClient (via DI) instead.""" + def __init__(self, llm_config, logger=None): self.logger = logger default_llm_config = llm_config.model_dump() @@ -21,7 +67,6 @@ def __init__(self, llm_config, logger=None): } def _extract_llm_overrides(self, request: dict): - """Extract and apply LLM overrides from metadata.llm_override.""" metadata = request.get("metadata") or {} llm_override = metadata.pop("llm_override", None) or {} @@ -87,7 +132,7 @@ async def chat_completion(self, request: dict): logger.error(f"Error while streaming chat completion: {str(e)}") raise - else: # Handle non-streaming response + else: try: response = await client.post( url=f"{base_url}/chat/completions", diff --git a/openrag/components/pipeline.py b/openrag/components/pipeline.py index 11471d7ee..89c339ccb 100644 --- a/openrag/components/pipeline.py +++ b/openrag/components/pipeline.py @@ -27,7 +27,7 @@ from pydantic import ValidationError from utils.logger import get_logger -from .llm import LLM +from .llm import _LLMShim as LLM from .map_reduce import RAGMapReduce from .reranker import BaseReranker, RerankerFactory from .retriever import BaseRetriever, RetrieverFactory @@ -90,7 +90,7 @@ def __init__(self) -> None: self.allow_filterless_fallback = config.retriever.allow_filterless_fallback self.reranker_enabled = config.reranker.enabled - self.reranker: BaseReranker = RerankerFactory.get_reranker(config) + self.reranker: BaseReranker = RerankerFactory.get_reranker(config.reranker) logger.debug("Reranker", enabled=self.reranker_enabled, provider=config.reranker.provider) self.reranker_top_k = config.reranker.top_k diff --git a/openrag/components/reranker/__init__.py b/openrag/components/reranker/__init__.py index aa7e719ac..2b9742509 100644 --- a/openrag/components/reranker/__init__.py +++ b/openrag/components/reranker/__init__.py @@ -1,17 +1,43 @@ +import asyncio + +import services.inference.reranker_clients # noqa: F401 — registers "infinity"/"openai" +from core.config.retrieval import RerankerConfig +from core.rerankers import reranker_registry + from .base import BaseReranker +class _RerankerShim(BaseReranker): + """Wraps a core ``Reranker`` (str-in / (idx, score)-out) behind the + legacy ``BaseReranker`` interface (Document-in / Document-out).""" + + def __init__(self, delegate, semaphore: int = 3): + self._delegate = delegate + self._semaphore = asyncio.Semaphore(semaphore) + + async def rerank(self, query, documents, top_k=None): + async with self._semaphore: + texts = [doc.page_content for doc in documents] + ranked = await self._delegate.rerank(query, texts, top_k=top_k) + output = [] + for index, score in ranked: + if not 0 <= index < len(documents): + continue + doc = documents[index] + doc.metadata["relevance_score"] = score + output.append(doc) + return output + + class RerankerFactory: @staticmethod - def get_reranker(config) -> BaseReranker: - provider = config.reranker.provider - if provider == "infinity": - from .infinity import InfinityReranker - - return InfinityReranker(config) - elif provider == "openai": - from .openai import OpenAIReranker - - return OpenAIReranker(config) - else: - raise ValueError(f"Unsupported reranker provider: {provider}") + def get_reranker(reranker_config: RerankerConfig) -> BaseReranker: + provider = reranker_config.provider + delegate = reranker_registry.create( + provider, + endpoint=reranker_config.base_url, + model_name=reranker_config.model_name, + api_key=reranker_config.api_key, + timeout=reranker_config.timeout, + ) + return _RerankerShim(delegate, semaphore=reranker_config.semaphore) diff --git a/openrag/components/reranker/base.py b/openrag/components/reranker/base.py index 4a0c2367e..19a3defa3 100644 --- a/openrag/components/reranker/base.py +++ b/openrag/components/reranker/base.py @@ -1,5 +1,6 @@ from abc import ABC, abstractmethod +from core.retrieval.rrf import rrf_reranking from langchain_core.documents.base import Document @@ -10,31 +11,8 @@ async def rerank(self, query: str, documents: list[Document], top_k: int | None @staticmethod def rrf_reranking(doc_lists: list[list[Document]], k: int = 60) -> list[Document]: - """Reciprocal_rank_fusion that takes multiple lists of ranked documents - and an optional parameter k used in the RRF formula - RRF formula: \\sum_{i=1}^{n} \frac{1}{k + rank_i} - where rank_i is the rank of the document in the i-th list and n is the number of lists. - - k small: High sensitivity to top ranks - k large: More balanced sensitivity across ranks - k = 60 a common and balanced choice in practice. - """ - - if len(doc_lists) == 1: - return doc_lists[0] - - # Initialize a dictionary to hold fused scores for each unique document - fused_scores = {} - - for doc_list in doc_lists: - doc_list: list[Document] - for rank, doc in enumerate(doc_list, start=1): - doc_id = doc.metadata.get("_id") - doc_key = ("id", doc_id) if doc_id is not None else ("object", id(doc)) - - score, d = fused_scores.get(doc_key, (0, doc)) - fused_scores[doc_key] = (score + 1 / (rank + k), d) - - # sort the docs - reranked_docs = [doc for _, doc in sorted(fused_scores.values(), key=lambda x: x[0], reverse=True)] - return reranked_docs + return rrf_reranking( + doc_lists, + key_fn=lambda doc: doc.metadata.get("_id", id(doc)), + k=k, + ) diff --git a/openrag/components/reranker/infinity.py b/openrag/components/reranker/infinity.py index f55cda05a..32626466d 100644 --- a/openrag/components/reranker/infinity.py +++ b/openrag/components/reranker/infinity.py @@ -1,9 +1,15 @@ +"""Backward-compatibility shim — delegates to services.inference.reranker_clients. + +All new code should import directly from ``services.inference.reranker_clients``. +""" + import asyncio from infinity_client import Client from infinity_client.api.default import rerank from infinity_client.models import RerankInput, ReRankResult from langchain_core.documents.base import Document +from services.inference.reranker_clients import InfinityReranker as InfinityRerankerAdapter # noqa: F401 from utils.logger import get_logger from .base import BaseReranker @@ -12,6 +18,8 @@ class InfinityReranker(BaseReranker): + """Legacy InfinityReranker. New code should use InfinityRerankerAdapter (via DI).""" + def __init__(self, config): self.model_name = config.reranker.model_name self.client = Client( @@ -33,7 +41,7 @@ async def rerank(self, query: str, documents: list[Document], top_k: int | None "documents": [doc.page_content for doc in documents], "top_n": top_k, "return_documents": True, - "raw_scores": True, # Normalized score between 0 and 1 + "raw_scores": True, } ) try: diff --git a/openrag/components/reranker/openai.py b/openrag/components/reranker/openai.py index 43f9d0acf..7bca4fbcf 100644 --- a/openrag/components/reranker/openai.py +++ b/openrag/components/reranker/openai.py @@ -1,7 +1,13 @@ +"""Backward-compatibility shim — delegates to services.inference.reranker_clients. + +All new code should import directly from ``services.inference.reranker_clients``. +""" + import asyncio import httpx from langchain_core.documents.base import Document +from services.inference.reranker_clients import OpenAIReranker as OpenAIRerankerAdapter # noqa: F401 from utils.logger import get_logger from .base import BaseReranker @@ -10,6 +16,8 @@ class OpenAIReranker(BaseReranker): + """Legacy OpenAIReranker. New code should use OpenAIRerankerAdapter (via DI).""" + def __init__(self, config): self.model_name = config.reranker.model_name base_url = config.reranker.base_url.rstrip("/") diff --git a/openrag/components/reranker/test_rrf_reranking.py b/openrag/components/reranker/test_rrf_reranking.py index 13f098baa..40efb456d 100644 --- a/openrag/components/reranker/test_rrf_reranking.py +++ b/openrag/components/reranker/test_rrf_reranking.py @@ -10,10 +10,11 @@ def make_doc(doc_id: str, content: str = "", **metadata) -> Document: class TestRrfRerankingSingleList: - def test_single_list_returned_as_is(self): + def test_single_list_returned_as_list_copy(self): docs = [make_doc("a"), make_doc("b"), make_doc("c")] result = BaseReranker.rrf_reranking([docs]) - assert result is docs + assert result == docs + assert result is not docs class TestRrfRerankingMultipleLists: diff --git a/openrag/components/retriever.py b/openrag/components/retriever.py index c18bcc031..b02e76e1d 100644 --- a/openrag/components/retriever.py +++ b/openrag/components/retriever.py @@ -66,6 +66,9 @@ async def chat(self, messages: list[dict[str, str]], **kwargs) -> str: out = await self._llm.ainvoke(lc_msgs) return out.content + async def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> Any: + pass # Not implemented since legacy code doesn't use streaming; core retriever won't call this method. + def _searcher() -> MilvusRayShim: """Wrap the legacy Vectordb Ray actor as a core ``RetrievalSearcher``.""" diff --git a/openrag/components/utils.py b/openrag/components/utils.py index 1c5f99e2a..a68215001 100644 --- a/openrag/components/utils.py +++ b/openrag/components/utils.py @@ -6,11 +6,14 @@ from collections import deque from typing import ClassVar -import ray from config import load_config from fast_langdetect import LangDetectConfig, LangDetector from langchain_core.documents.base import Document from langchain_openai import ChatOpenAI +from services.inference.distributed_semaphore import ( + DistributedSemaphore, # noqa: F401 + DistributedSemaphoreActor, # noqa: F401 +) from utils.logger import get_logger SOURCE_SEPARATOR = "-" * 10 + "\n\n" @@ -33,56 +36,6 @@ def __call__(cls, *args, **kwargs): return cls._instances[cls] -@ray.remote(max_restarts=5, max_concurrency=config.ray.semaphore.concurrency) -class DistributedSemaphoreActor: - def __init__(self, max_concurrent_ops: int): - self.semaphore = asyncio.Semaphore(max_concurrent_ops) - - async def acquire(self): - await self.semaphore.acquire() - - def release(self): - self.semaphore.release() - - -class DistributedSemaphore: - # https://chat.deepseek.com/a/chat/s/890dbcc0-2d3f-4819-af9d-774b892905bc - def __init__( - self, - name: str = "llmSemaphore", - namespace="openrag", - max_concurrent_ops: int = 10, - ): - self._name = name - self._namespace = namespace - self._max_concurrent_ops = max_concurrent_ops - - def _get_or_create_actor(self): - try: - # reuse existing actor if it exists - _actor = ray.get_actor(self._name, namespace=self._namespace) - except ValueError: - # create new actor if it doesn't exist - _actor = DistributedSemaphoreActor.options( - name=self._name, - namespace=self._namespace, - lifetime="detached", - ).remote(self._max_concurrent_ops) - except Exception: - raise - - return _actor - - async def __aenter__(self): - semaphore_actor = self._get_or_create_actor() - await semaphore_actor.acquire.remote() - return self - - async def __aexit__(self, exc_type, exc, tb): - semaphore_actor = self._get_or_create_actor() - await semaphore_actor.release.remote() - - _cached_length_function = None diff --git a/openrag/core/llm/llm.py b/openrag/core/llm/llm.py index ac0c8a657..54c0bc275 100644 --- a/openrag/core/llm/llm.py +++ b/openrag/core/llm/llm.py @@ -2,63 +2,29 @@ from __future__ import annotations -import json from abc import ABC, abstractmethod from collections.abc import AsyncIterator -from openrag.core.utils.exceptions import LLMParsingError - class LLM(ABC): """Base class for all LLM providers.""" @abstractmethod - async def generate(self, prompt: str, **kwargs) -> str: - """Generate a completion for a prompt.""" + async def generate(self, prompt: str, **kwargs) -> dict: + """Generate a text completion for a prompt.""" ... @abstractmethod - async def chat(self, messages: list[dict[str, str]], **kwargs) -> str: + async def chat(self, messages: list[dict[str, str]], **kwargs) -> dict: """Chat completion with message list.""" ... - async def generate_json(self, prompt: str, **kwargs) -> dict: - """Generate a JSON response. Default: parse generate() output. - - Raises LLMParsingError if the LLM output is not valid JSON - or if the result is not a dict. - """ - response = await self.generate(prompt, **kwargs) - text = response.strip() - try: - result = json.loads(text) - except json.JSONDecodeError as exc: - raise LLMParsingError( - raw_response=text, - parse_error=str(exc), - ) from exc - if not isinstance(result, dict): - raise LLMParsingError( - raw_response=text, - parse_error=f"Expected JSON object, got {type(result).__name__}", - ) - return result - - async def chat_with_tools( - self, - messages: list[dict[str, str]], - tools: list[dict], - tool_choice: str | dict = "required", - **kwargs, - ) -> dict: - """Chat completion with function calling. + @abstractmethod + def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> AsyncIterator[str]: + """Stream chat completion as raw SSE lines. - Only supported by backends with function calling (e.g., vLLM). - Default raises NotImplementedError. + Implementations must be ``async def`` generators yielding ``str`` chunks. + Declared without ``async def`` here so the abstract signature matches the + ``AsyncIterator[str]`` return type without forcing an empty ``yield``. """ - raise NotImplementedError(f"{type(self).__name__} does not support tool calling") - - async def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> AsyncIterator[str]: - """Stream chat completion. Default falls back to non-streaming.""" - result = await self.chat(messages, **kwargs) - yield result + ... diff --git a/openrag/core/retrieval/rrf.py b/openrag/core/retrieval/rrf.py index eadffdbe9..13e7f5baf 100644 --- a/openrag/core/retrieval/rrf.py +++ b/openrag/core/retrieval/rrf.py @@ -40,7 +40,7 @@ def rrf_reranking( Returns: A single ranked list, best first. Empty input -> empty list. - Single input list is returned as-is. + Single input list is shallow-copied so callers always get a ``list``. Raises: ValueError: if ``k < 0`` (would produce a zero or negative diff --git a/openrag/core/utils/test_external_errors.py b/openrag/core/utils/test_external_errors.py new file mode 100644 index 000000000..f0ee3b290 --- /dev/null +++ b/openrag/core/utils/test_external_errors.py @@ -0,0 +1,129 @@ +""" +Tests for external resource error detection utilities. +Related to: https://github.com/linagora/openrag/issues/182 +""" + +import pytest +from core.utils.external_errors import is_external_resource_error + + +class TestIsExternalResourceError: + """Test suite for is_external_resource_error function.""" + + @pytest.mark.parametrize( + "error_msg,expected_code,url_contains", + [ + # Issue #182 + ( + "aiohttp.client_exceptions.ClientResponseError: 403, message='Forbidden', " + "url='https://upload.wikimedia.org/wikipedia/commons/thumb/d/d5/Logo.png'", + "403", + "upload.wikimedia.org", + ), + # Other HTTP status codes + ( + "ClientResponseError: 404, url='https://example.com/missing.png'", + "404", + "example.com", + ), + ( + "HTTPError: 401 Unauthorized for url: https://api.example.com/image.jpg", + "401", + "api.example.com", + ), + ( + "ClientResponseError: 429 Too Many Requests - https://cdn.example.com/img.png", + "429", + "cdn.example.com", + ), + # 5xx gateway errors + ( + "502 Bad Gateway: https://api.example.com/image.png", + "502", + "api.example.com", + ), + ( + "ClientResponseError: 503 Service Unavailable - https://cdn.example.com/img.png", + "503", + "cdn.example.com", + ), + # vLLM wrapped error (the real-world scenario) + ( + "openai.InternalServerError: Error code: 500 - {'error': {'message': " + "'litellm.InternalServerError: aiohttp.client_exceptions.ClientResponseError: " + "403, message=Forbidden, url=https://example.com/path/to/image.png'}}", + "403", + "example.com/path/to/image.png", + ), + ], + ) + def test_detects_http_errors_with_urls(self, error_msg, expected_code, url_contains): + """Test detection of HTTP errors with URL extraction.""" + is_external, status_code, url = is_external_resource_error(Exception(error_msg)) + + assert is_external is True + assert status_code == expected_code + assert url_contains in url + + @pytest.mark.parametrize( + "error_msg", + [ + "TimeoutError: Connection timed out while fetching resource", + "SSLError: Certificate verification failed", + "ConnectionError: Failed to connect to server", + "aiohttp.client_exceptions.ClientResponseError: some error", + "requests.exceptions.HTTPError: 500 Server Error", + ], + ) + def test_detects_error_indicators(self, error_msg): + """Test detection via error type indicators.""" + is_external, _, _ = is_external_resource_error(Exception(error_msg)) + assert is_external is True + + @pytest.mark.parametrize( + "error", + [ + Exception("ValueError: Invalid input parameter"), + Exception("Something went wrong during processing"), + TypeError("'NoneType' object is not subscriptable"), + AttributeError("'dict' object has no attribute 'content'"), + Exception(""), + # vLLM error without external cause details + Exception( + "openai.InternalServerError: Error code: 500 - {'error': {'message': " + "'litellm.InternalServerError: InternalServerError: OpenAIException'}}" + ), + ], + ) + def test_does_not_flag_internal_errors(self, error): + """Test that internal/generic errors are not flagged as external.""" + is_external, status_code, url = is_external_resource_error(error) + + assert is_external is False + assert status_code == "" + assert url == "" + + def test_extracts_url_with_query_params(self): + """Test URL extraction with query parameters.""" + error = Exception("403 Forbidden: https://api.example.com/image?id=123&size=large") + _, _, url = is_external_resource_error(error) + + assert "api.example.com/image?id=123" in url + + def test_indicator_substring_causes_false_positive(self): + """Document known limitation: indicator substrings cause false positives. + + This test documents that internal errors mentioning HTTP error class names + will be incorrectly classified as external. This is accepted because: + 1. Real error messages use these as exception class names, not prose + 2. Stricter matching (word boundaries) would break legitimate matches + like 'aiohttp.client_exceptions.ClientResponseError' + 3. This scenario is unlikely in practice + """ + error = Exception("InternalServerError: Failed to handle ClientResponseError in retry logic") + is_external, status_code, url = is_external_resource_error(error) + + # This IS classified as external (false positive) due to substring match + assert is_external is True + assert status_code == "" # No HTTP status code + assert url == "" # No URL diff --git a/openrag/di/container.py b/openrag/di/container.py new file mode 100644 index 000000000..c322eda48 --- /dev/null +++ b/openrag/di/container.py @@ -0,0 +1,38 @@ +"""Service container — wires registries and exposes component factories.""" + +from __future__ import annotations + +from core.embeddings import embedder_registry +from core.llm import llm_registry +from core.rerankers import reranker_registry +from core.vlm import vlm_registry +from di.embedders import register_embedders +from di.llms import register_llms +from di.rerankers import register_rerankers +from di.vlms import register_vlms + + +class ServiceContainer: + """Populates registries and provides typed factory access.""" + + def __init__(self) -> None: + register_embedders() + register_llms() + register_rerankers() + register_vlms() + + @staticmethod + def create_embedder(name: str = "vllm", **kwargs): + return embedder_registry.create(name, **kwargs) + + @staticmethod + def create_llm(name: str = "vllm", **kwargs): + return llm_registry.create(name, **kwargs) + + @staticmethod + def create_reranker(name: str = "infinity", **kwargs): + return reranker_registry.create(name, **kwargs) + + @staticmethod + def create_vlm(name: str = "vllm", **kwargs): + return vlm_registry.create(name, **kwargs) diff --git a/openrag/di/embedders.py b/openrag/di/embedders.py new file mode 100644 index 000000000..5f09ecda2 --- /dev/null +++ b/openrag/di/embedders.py @@ -0,0 +1,5 @@ +"""Register embedder implementations with the core registry.""" + + +def register_embedders() -> None: + import services.inference.vllm_client # noqa: F401 diff --git a/openrag/di/inference.py b/openrag/di/inference.py new file mode 100644 index 000000000..8447f4e42 --- /dev/null +++ b/openrag/di/inference.py @@ -0,0 +1,16 @@ +"""Convenience wrapper — registers all inference adapters at once. + +Delegates to the per-domain registration modules. +""" + +from di.embedders import register_embedders +from di.llms import register_llms +from di.rerankers import register_rerankers +from di.vlms import register_vlms + + +def register_inference() -> None: + register_embedders() + register_llms() + register_rerankers() + register_vlms() diff --git a/openrag/di/llms.py b/openrag/di/llms.py new file mode 100644 index 000000000..8620f7159 --- /dev/null +++ b/openrag/di/llms.py @@ -0,0 +1,5 @@ +"""Register LLM implementations with the core registry.""" + + +def register_llms() -> None: + import services.inference.vllm_client # noqa: F401 diff --git a/openrag/di/rerankers.py b/openrag/di/rerankers.py new file mode 100644 index 000000000..d16aab1bb --- /dev/null +++ b/openrag/di/rerankers.py @@ -0,0 +1,5 @@ +"""Register reranker implementations with the core registry.""" + + +def register_rerankers() -> None: + import services.inference.reranker_clients # noqa: F401 diff --git a/openrag/di/test_inference.py b/openrag/di/test_inference.py new file mode 100644 index 000000000..f38de06b2 --- /dev/null +++ b/openrag/di/test_inference.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +from core.embeddings import embedder_registry +from core.llm import llm_registry +from core.rerankers import reranker_registry +from core.vlm import vlm_registry +from di.container import ServiceContainer +from di.inference import register_inference + + +class TestRegisterInference: + def test_registries_populated(self): + register_inference() + + assert "vllm" in llm_registry + assert "vllm" in embedder_registry + assert "vllm" in vlm_registry + assert "infinity" in reranker_registry + assert "openai" in reranker_registry + + def test_idempotent(self): + register_inference() + register_inference() + + +class TestServiceContainer: + def test_container_populates_all_registries(self): + ServiceContainer() + + assert "vllm" in llm_registry + assert "vllm" in embedder_registry + assert "vllm" in vlm_registry + assert "infinity" in reranker_registry + assert "openai" in reranker_registry + + def test_create_llm(self): + container = ServiceContainer() + client = container.create_llm(endpoint="http://vllm:8000/v1", model_name="m") + assert client is not None + + def test_create_embedder(self): + container = ServiceContainer() + client = container.create_embedder(endpoint="http://vllm:8000/v1", model_name="m") + assert client is not None + + def test_create_reranker(self): + container = ServiceContainer() + client = container.create_reranker(endpoint="http://reranker:7997", model_name="m") + assert client is not None + + def test_create_vlm(self): + container = ServiceContainer() + client = container.create_vlm(endpoint="http://vllm:8000/v1", model_name="m") + assert client is not None diff --git a/openrag/di/vlms.py b/openrag/di/vlms.py new file mode 100644 index 000000000..873c78b81 --- /dev/null +++ b/openrag/di/vlms.py @@ -0,0 +1,5 @@ +"""Register VLM implementations with the core registry.""" + + +def register_vlms() -> None: + import services.inference.vllm_client # noqa: F401 diff --git a/openrag/services/inference/__init__.py b/openrag/services/inference/__init__.py index e69de29bb..efbe4f41a 100644 --- a/openrag/services/inference/__init__.py +++ b/openrag/services/inference/__init__.py @@ -0,0 +1,25 @@ +"""Inference service layer — clients, resilience, and concurrency primitives. + +Importing this package registers implementations in the core registries +(``llm_registry``, ``embedder_registry``, ``vlm_registry``, ``reranker_registry``) +so they can be created via ``registry.create("name", **kwargs)``. +""" + +from ._circuit_breaker import get_breaker, with_circuit_breaker +from ._retry import with_retry +from .distributed_semaphore import DistributedSemaphore, DistributedSemaphoreActor +from .reranker_clients import InfinityReranker, OpenAIReranker +from .vllm_client import VLLMClient, VLLMEmbedder, VLLMVision + +__all__ = [ + "DistributedSemaphore", + "DistributedSemaphoreActor", + "InfinityReranker", + "OpenAIReranker", + "VLLMClient", + "VLLMEmbedder", + "VLLMVision", + "get_breaker", + "with_circuit_breaker", + "with_retry", +] diff --git a/openrag/services/inference/_circuit_breaker.py b/openrag/services/inference/_circuit_breaker.py new file mode 100644 index 000000000..2d332699c --- /dev/null +++ b/openrag/services/inference/_circuit_breaker.py @@ -0,0 +1,85 @@ +from datetime import timedelta +from functools import wraps + +import httpx +from aiobreaker import CircuitBreaker, CircuitBreakerError, CircuitBreakerListener +from core.utils.exceptions import InferenceConnectionError, LLMParsingError, OpenRAGError +from prometheus_client import Gauge +from utils.logger import get_logger + +logger = get_logger() + +_breakers: dict[str, CircuitBreaker] = {} +_breaker_config: dict[str, tuple[int, float]] = {} + +try: + CIRCUIT_BREAKER_STATE = Gauge( + "openrag_circuit_breaker_state", + "Circuit breaker state (0=closed, 1=open, 2=half-open)", + ["name"], + ) +except ValueError: + from prometheus_client import REGISTRY + + CIRCUIT_BREAKER_STATE = REGISTRY._names_to_collectors["openrag_circuit_breaker_state"] + +_STATE_VALUES = {"ClosedState": 0, "OpenState": 1, "HalfOpenState": 2} + + +def _is_client_error(exc: Exception) -> bool: + if isinstance(exc, httpx.HTTPStatusError): + return 400 <= exc.response.status_code < 500 + if isinstance(exc, OpenRAGError): + return 400 <= exc.status_code < 500 + return False + + +def _is_excluded(exc: Exception) -> bool: + if _is_client_error(exc): + return True + if isinstance(exc, LLMParsingError): + return True + return False + + +class _LoggingListener(CircuitBreakerListener): + def state_change(self, breaker, old, new): + state_name = type(new).__name__ + logger.warning( + "Circuit breaker '{name}' state: {old} -> {new}", + name=breaker.name, + old=type(old).__name__, + new=state_name, + ) + CIRCUIT_BREAKER_STATE.labels(name=breaker.name).set(_STATE_VALUES.get(state_name, -1)) + + +def get_breaker(name: str, fail_max: int = 50, timeout_duration: float = 60.0) -> CircuitBreaker: + requested = (fail_max, timeout_duration) + if name not in _breakers: + _breakers[name] = CircuitBreaker( + fail_max=fail_max, + timeout_duration=timedelta(seconds=timeout_duration), + name=name, + exclude=[_is_excluded], + listeners=[_LoggingListener()], + ) + _breaker_config[name] = requested + elif _breaker_config.get(name) != requested: + raise ValueError(f"Breaker '{name}' already exists with config={_breaker_config[name]}, requested={requested}") + return _breakers[name] + + +def with_circuit_breaker(name: str, fail_max: int = 50, timeout_duration: float = 60.0): + def decorator(fn): + @wraps(fn) + async def wrapper(*args, **kwargs): + breaker = get_breaker(name, fail_max, timeout_duration) + try: + return await breaker.call_async(fn, *args, **kwargs) + except CircuitBreakerError: + raise InferenceConnectionError(f"Circuit open for '{name}'") + + return wrapper + + return decorator diff --git a/openrag/services/inference/_retry.py b/openrag/services/inference/_retry.py new file mode 100644 index 000000000..f74ff6bc1 --- /dev/null +++ b/openrag/services/inference/_retry.py @@ -0,0 +1,44 @@ +import httpx +from core.utils.exceptions import OpenRAGError +from tenacity import ( + RetryCallState, + retry, + retry_if_exception, + stop_after_attempt, + wait_exponential_jitter, +) +from utils.logger import get_logger + +logger = get_logger() + +_RETRYABLE_STATUS_CODES = {429, 502, 503, 504} + + +def _is_retryable(exc: BaseException) -> bool: + if isinstance(exc, (httpx.TimeoutException, httpx.ConnectError)): + return True + if isinstance(exc, httpx.HTTPStatusError): + return exc.response.status_code in _RETRYABLE_STATUS_CODES + if isinstance(exc, OpenRAGError): + return exc.status_code in _RETRYABLE_STATUS_CODES + return False + + +def _log_retry(state: RetryCallState) -> None: + exc = state.outcome.exception() if state.outcome else None + logger.warning( + "Retrying after transient failure (attempt {attempt}/{max}): {exc}", + attempt=state.attempt_number, + max=state.retry_object.stop.max_attempt_number, # type: ignore[union-attr] + exc=repr(exc), + ) + + +def with_retry(max_attempts: int = 3, base_wait: float = 1.0, max_wait: float = 30.0): + return retry( + stop=stop_after_attempt(max_attempts), + wait=wait_exponential_jitter(initial=base_wait, max=max_wait, exp_base=2), + retry=retry_if_exception(_is_retryable), + before_sleep=_log_retry, + reraise=True, + ) diff --git a/openrag/services/inference/distributed_semaphore.py b/openrag/services/inference/distributed_semaphore.py new file mode 100644 index 000000000..c535e2bda --- /dev/null +++ b/openrag/services/inference/distributed_semaphore.py @@ -0,0 +1,61 @@ +"""Ray-based distributed semaphore for cluster-wide concurrency limiting. + +Extracted from ``components/utils.py``. The actor handles acquire/release; +``DistributedSemaphore`` locates (or creates) the actor and wraps it as an +async context manager. +""" + +from __future__ import annotations + +import asyncio + +import ray + + +@ray.remote(max_restarts=5) +class DistributedSemaphoreActor: + def __init__(self, max_concurrent_ops: int): + self.semaphore = asyncio.Semaphore(max_concurrent_ops) + + async def acquire(self): + await self.semaphore.acquire() + + def release(self): + self.semaphore.release() + + +class DistributedSemaphore: + """Async context manager backed by a detached Ray actor. + + The actor is created on first use (get-or-create) and survives across + callers within the same Ray cluster. + """ + + def __init__( + self, + name: str = "llmSemaphore", + namespace: str = "openrag", + max_concurrent_ops: int = 10, + ): + self._name = name + self._namespace = namespace + self._max_concurrent_ops = max_concurrent_ops + + def _get_or_create_actor(self): + try: + return ray.get_actor(self._name, namespace=self._namespace) + except ValueError: + return DistributedSemaphoreActor.options( + name=self._name, + namespace=self._namespace, + lifetime="detached", + ).remote(self._max_concurrent_ops) + + async def __aenter__(self): + semaphore_actor = self._get_or_create_actor() + await semaphore_actor.acquire.remote() + return self + + async def __aexit__(self, exc_type, exc, tb): + semaphore_actor = self._get_or_create_actor() + await semaphore_actor.release.remote() diff --git a/openrag/services/inference/healthcheck.py b/openrag/services/inference/healthcheck.py new file mode 100644 index 000000000..7cf259006 --- /dev/null +++ b/openrag/services/inference/healthcheck.py @@ -0,0 +1,90 @@ +"""Probe inference endpoints for readiness. + +Used at container startup (fail-fast) and by the ``/health_check`` route. +Uses raw ``httpx`` — no OpenAI SDK dependency. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass, field +from enum import Enum + +import httpx +from utils.logger import get_logger + +logger = get_logger() + + +class EndpointStatus(str, Enum): + HEALTHY = "healthy" + UNHEALTHY = "unhealthy" + UNREACHABLE = "unreachable" + + +@dataclass +class HealthResult: + url: str + status: EndpointStatus + latency_ms: float = 0.0 + models: list[str] = field(default_factory=list) + http_status: int | None = None + error: str | None = None + + +async def check_endpoint_health(endpoint: str, *, timeout: float = 5.0) -> HealthResult: + """Probe a vLLM / OpenAI-compatible server via ``GET /v1/models``.""" + url = endpoint.rstrip("/") + start = time.monotonic() + try: + async with httpx.AsyncClient(timeout=timeout) as client: + resp = await client.get(f"{url}/v1/models") + latency = (time.monotonic() - start) * 1000 + if resp.status_code == 200: + models = [m["id"] for m in resp.json().get("data", [])] + return HealthResult(url=url, status=EndpointStatus.HEALTHY, latency_ms=latency, models=models) + return HealthResult(url=url, status=EndpointStatus.UNHEALTHY, latency_ms=latency, http_status=resp.status_code) + except (httpx.ConnectError, httpx.TimeoutException) as exc: + latency = (time.monotonic() - start) * 1000 + return HealthResult(url=url, status=EndpointStatus.UNREACHABLE, latency_ms=latency, error=str(exc)) + except Exception as exc: + latency = (time.monotonic() - start) * 1000 + logger.warning("Unexpected error probing endpoint", url=url, error=str(exc)) + return HealthResult(url=url, status=EndpointStatus.UNREACHABLE, latency_ms=latency, error=str(exc)) + + +async def check_infinity(endpoint: str, *, timeout: float = 5.0) -> HealthResult: + """Probe an Infinity reranker server via ``GET /health``.""" + url = endpoint.rstrip("/") + start = time.monotonic() + try: + async with httpx.AsyncClient(timeout=timeout) as client: + resp = await client.get(f"{url}/health") + latency = (time.monotonic() - start) * 1000 + if resp.status_code == 200: + return HealthResult(url=url, status=EndpointStatus.HEALTHY, latency_ms=latency) + return HealthResult(url=url, status=EndpointStatus.UNHEALTHY, latency_ms=latency, http_status=resp.status_code) + except (httpx.ConnectError, httpx.TimeoutException) as exc: + latency = (time.monotonic() - start) * 1000 + return HealthResult(url=url, status=EndpointStatus.UNREACHABLE, latency_ms=latency, error=str(exc)) + except Exception as exc: + latency = (time.monotonic() - start) * 1000 + logger.warning("Unexpected error probing infinity endpoint", url=url, error=str(exc)) + return HealthResult(url=url, status=EndpointStatus.UNREACHABLE, latency_ms=latency, error=str(exc)) + + +async def check_model_available(endpoint: str, model: str, *, timeout: float = 5.0) -> HealthResult: + """Probe an OpenAI-compatible endpoint and verify a specific model is served.""" + result = await check_endpoint_health(endpoint, timeout=timeout) + if result.status != EndpointStatus.HEALTHY: + return result + if model not in result.models: + available = ", ".join(result.models) if result.models else "(none)" + return HealthResult( + url=result.url, + status=EndpointStatus.UNHEALTHY, + latency_ms=result.latency_ms, + models=result.models, + error=f"Model '{model}' not found. Available: {available}", + ) + return result diff --git a/openrag/services/inference/reranker_clients.py b/openrag/services/inference/reranker_clients.py new file mode 100644 index 000000000..2e6366f3e --- /dev/null +++ b/openrag/services/inference/reranker_clients.py @@ -0,0 +1,127 @@ +"""Reranker inference clients. + +Two classes — Infinity and OpenAI-compatible — both implementing the +``Reranker`` ABC. Both talk to a ``/rerank`` endpoint with the same +payload shape, differing only in the base URL and transport library +the old code used. Now both use ``httpx`` directly. +""" + +from __future__ import annotations + +import httpx +from core.rerankers import Reranker, reranker_registry +from core.utils.exceptions import InferenceConnectionError, InferenceTimeoutError +from utils.logger import get_logger + +from ._circuit_breaker import with_circuit_breaker +from ._retry import with_retry + +logger = get_logger() + + +@reranker_registry.register("infinity") +class InfinityReranker(Reranker): + """Reranker backed by an Infinity server.""" + + def __init__( + self, + endpoint: str, + model_name: str, + *, + api_key: str = "", + timeout: float = 30.0, + **_kwargs, + ): + self._endpoint = endpoint.rstrip("/") + self._model = model_name + headers = {} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + self._client = httpx.AsyncClient(timeout=timeout, headers=headers) + + @with_circuit_breaker("reranker") + @with_retry(max_attempts=2) + async def rerank(self, query: str, documents: list[str], top_k: int | None = None) -> list[tuple[int, float]]: + top_k = min(top_k, len(documents)) if top_k is not None else len(documents) + try: + resp = await self._client.post( + f"{self._endpoint}/rerank", + json={ + "model": self._model, + "query": query, + "documents": documents, + "top_n": top_k, + "return_documents": False, + "raw_scores": True, + }, + ) + resp.raise_for_status() + except httpx.ConnectError as exc: + raise InferenceConnectionError(f"Cannot reach reranker at {self._endpoint}") from exc + except httpx.TimeoutException as exc: + raise InferenceTimeoutError(f"Reranker request timed out at {self._endpoint}") from exc + except httpx.HTTPStatusError as exc: + raise InferenceConnectionError( + f"Reranker at {self._endpoint} returned HTTP {exc.response.status_code}" + ) from exc + try: + results = resp.json()["results"] + return [(r["index"], r["relevance_score"]) for r in results] + except (KeyError, TypeError, ValueError) as exc: + raise InferenceConnectionError(f"Unexpected reranker response format from {self._endpoint}") from exc + + async def aclose(self) -> None: + await self._client.aclose() + + +@reranker_registry.register("openai") +class OpenAIReranker(Reranker): + """Reranker backed by an OpenAI-compatible reranking endpoint.""" + + def __init__( + self, + endpoint: str, + model_name: str, + *, + api_key: str = "", + timeout: float = 30.0, + **_kwargs, + ): + self._endpoint = endpoint.rstrip("/") + self._model = model_name + headers = {} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + self._client = httpx.AsyncClient(timeout=timeout, headers=headers) + + @with_circuit_breaker("reranker") + @with_retry(max_attempts=2) + async def rerank(self, query: str, documents: list[str], top_k: int | None = None) -> list[tuple[int, float]]: + top_k = min(top_k, len(documents)) if top_k is not None else len(documents) + try: + resp = await self._client.post( + f"{self._endpoint}/rerank", + json={ + "model": self._model, + "query": query, + "documents": documents, + "top_n": top_k, + }, + ) + resp.raise_for_status() + except httpx.ConnectError as exc: + raise InferenceConnectionError(f"Cannot reach reranker at {self._endpoint}") from exc + except httpx.TimeoutException as exc: + raise InferenceTimeoutError(f"Reranker request timed out at {self._endpoint}") from exc + except httpx.HTTPStatusError as exc: + raise InferenceConnectionError( + f"Reranker at {self._endpoint} returned HTTP {exc.response.status_code}" + ) from exc + try: + results = resp.json()["results"] + return [(r["index"], r["relevance_score"]) for r in results] + except (KeyError, TypeError, ValueError) as exc: + raise InferenceConnectionError(f"Unexpected reranker response format from {self._endpoint}") from exc + + async def aclose(self) -> None: + await self._client.aclose() diff --git a/openrag/services/inference/test_circuit_breaker.py b/openrag/services/inference/test_circuit_breaker.py new file mode 100644 index 000000000..dedb3a58d --- /dev/null +++ b/openrag/services/inference/test_circuit_breaker.py @@ -0,0 +1,116 @@ +import httpx +import pytest +from core.utils.exceptions import InferenceConnectionError, LLMParsingError +from services.inference._circuit_breaker import ( + _breaker_config, + _breakers, + get_breaker, + with_circuit_breaker, +) + + +@pytest.fixture(autouse=True) +def _clean_breakers(): + for breaker in _breakers.values(): + breaker.close() + _breakers.clear() + _breaker_config.clear() + yield + for breaker in _breakers.values(): + breaker.close() + _breakers.clear() + _breaker_config.clear() + + +class TestGetBreaker: + def test_returns_same_instance(self): + b1 = get_breaker("llm") + b2 = get_breaker("llm") + assert b1 is b2 + + def test_different_names_different_instances(self): + b1 = get_breaker("llm") + b2 = get_breaker("embedder") + assert b1 is not b2 + + def test_default_fail_max_is_50(self): + b = get_breaker("test-default") + assert b.fail_max == 50 + + +class TestExclusions: + @pytest.mark.asyncio + async def test_client_4xx_excluded(self): + breaker = get_breaker("test-4xx", fail_max=2, timeout_duration=1.0) + + async def fail_4xx(): + req = httpx.Request("GET", "http://test") + raise httpx.HTTPStatusError("bad request", request=req, response=httpx.Response(400, request=req)) + + for _ in range(5): + with pytest.raises(httpx.HTTPStatusError): + await breaker.call_async(fail_4xx) + + assert "Closed" in type(breaker.state).__name__ + + @pytest.mark.asyncio + async def test_llm_parsing_error_excluded(self): + breaker = get_breaker("test-parse", fail_max=2, timeout_duration=1.0) + + async def fail_parse(): + raise LLMParsingError(raw_response="not json") + + for _ in range(5): + with pytest.raises(LLMParsingError): + await breaker.call_async(fail_parse) + + assert "Closed" in type(breaker.state).__name__ + + @pytest.mark.asyncio + async def test_server_5xx_trips_breaker(self): + breaker = get_breaker("test-5xx", fail_max=2, timeout_duration=1.0) + + async def fail_5xx(): + req = httpx.Request("GET", "http://test") + raise httpx.HTTPStatusError("bad gateway", request=req, response=httpx.Response(502, request=req)) + + with pytest.raises(httpx.HTTPStatusError): + await breaker.call_async(fail_5xx) + + from aiobreaker import CircuitBreakerError + + with pytest.raises(CircuitBreakerError): + await breaker.call_async(fail_5xx) + + assert "Open" in type(breaker.state).__name__ + + +class TestWithCircuitBreaker: + @pytest.mark.asyncio + async def test_passes_through_on_success(self): + @with_circuit_breaker("test-ok", fail_max=3, timeout_duration=1.0) + async def ok(): + return "result" + + assert await ok() == "result" + + @pytest.mark.asyncio + async def test_raises_inference_connection_error_when_open(self): + call_count = 0 + + @with_circuit_breaker("test-open", fail_max=2, timeout_duration=60.0) + async def always_fail(): + nonlocal call_count + call_count += 1 + raise ConnectionError("down") + + with pytest.raises(ConnectionError): + await always_fail() + + with pytest.raises(InferenceConnectionError, match="Circuit open"): + await always_fail() + + with pytest.raises(InferenceConnectionError, match="Circuit open"): + await always_fail() + + assert call_count == 2 diff --git a/openrag/services/inference/test_distributed_semaphore.py b/openrag/services/inference/test_distributed_semaphore.py new file mode 100644 index 000000000..cc77c50cc --- /dev/null +++ b/openrag/services/inference/test_distributed_semaphore.py @@ -0,0 +1,18 @@ +from services.inference.distributed_semaphore import DistributedSemaphore, DistributedSemaphoreActor + + +class TestDistributedSemaphore: + def test_default_params(self): + sem = DistributedSemaphore() + assert sem._name == "llmSemaphore" + assert sem._namespace == "openrag" + assert sem._max_concurrent_ops == 10 + + def test_custom_params(self): + sem = DistributedSemaphore(name="vlm", namespace="test", max_concurrent_ops=5) + assert sem._name == "vlm" + assert sem._namespace == "test" + assert sem._max_concurrent_ops == 5 + + def test_actor_class_exists(self): + assert hasattr(DistributedSemaphoreActor, "remote") diff --git a/openrag/services/inference/test_healthcheck.py b/openrag/services/inference/test_healthcheck.py new file mode 100644 index 000000000..24ff2054e --- /dev/null +++ b/openrag/services/inference/test_healthcheck.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +from .healthcheck import ( + EndpointStatus, + check_endpoint_health, + check_infinity, + check_model_available, +) + + +@pytest.fixture +def models_response(): + return httpx.Response(200, json={"data": [{"id": "mistral-small"}, {"id": "bge-m3"}]}) + + +@pytest.fixture +def health_ok(): + return httpx.Response(200, json={"status": "ok"}) + + +class TestCheckOpenAICompatible: + @pytest.mark.asyncio + async def test_healthy(self, models_response): + transport = httpx.MockTransport(lambda req: models_response) + with patch( + "services.inference.healthcheck.httpx.AsyncClient", return_value=httpx.AsyncClient(transport=transport) + ): + result = await check_endpoint_health("http://vllm:8000") + assert result.status == EndpointStatus.HEALTHY + assert "mistral-small" in result.models + assert "bge-m3" in result.models + assert result.latency_ms > 0 + + @pytest.mark.asyncio + async def test_strips_trailing_slash(self, models_response): + transport = httpx.MockTransport(lambda req: models_response) + with patch( + "services.inference.healthcheck.httpx.AsyncClient", return_value=httpx.AsyncClient(transport=transport) + ): + result = await check_endpoint_health("http://vllm:8000/") + assert result.url == "http://vllm:8000" + + @pytest.mark.asyncio + async def test_unhealthy_status(self): + transport = httpx.MockTransport(lambda req: httpx.Response(503)) + with patch( + "services.inference.healthcheck.httpx.AsyncClient", return_value=httpx.AsyncClient(transport=transport) + ): + result = await check_endpoint_health("http://vllm:8000") + assert result.status == EndpointStatus.UNHEALTHY + assert result.http_status == 503 + + @pytest.mark.asyncio + async def test_connection_error(self): + async def raise_connect_error(*a, **kw): + raise httpx.ConnectError("Connection refused") + + client = AsyncMock() + client.get = raise_connect_error + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=False) + with patch("services.inference.healthcheck.httpx.AsyncClient", return_value=client): + result = await check_endpoint_health("http://vllm:8000") + assert result.status == EndpointStatus.UNREACHABLE + assert "Connection refused" in result.error + + @pytest.mark.asyncio + async def test_timeout(self): + async def raise_timeout(*a, **kw): + raise httpx.TimeoutException("timed out") + + client = AsyncMock() + client.get = raise_timeout + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=False) + with patch("services.inference.healthcheck.httpx.AsyncClient", return_value=client): + result = await check_endpoint_health("http://vllm:8000") + assert result.status == EndpointStatus.UNREACHABLE + assert "timed out" in result.error + + +class TestCheckInfinity: + @pytest.mark.asyncio + async def test_healthy(self, health_ok): + transport = httpx.MockTransport(lambda req: health_ok) + with patch( + "services.inference.healthcheck.httpx.AsyncClient", return_value=httpx.AsyncClient(transport=transport) + ): + result = await check_infinity("http://reranker:7997") + assert result.status == EndpointStatus.HEALTHY + assert result.latency_ms > 0 + + @pytest.mark.asyncio + async def test_unhealthy(self): + transport = httpx.MockTransport(lambda req: httpx.Response(500)) + with patch( + "services.inference.healthcheck.httpx.AsyncClient", return_value=httpx.AsyncClient(transport=transport) + ): + result = await check_infinity("http://reranker:7997") + assert result.status == EndpointStatus.UNHEALTHY + assert result.http_status == 500 + + +class TestCheckModelAvailable: + @pytest.mark.asyncio + async def test_model_found(self, models_response): + transport = httpx.MockTransport(lambda req: models_response) + with patch( + "services.inference.healthcheck.httpx.AsyncClient", return_value=httpx.AsyncClient(transport=transport) + ): + result = await check_model_available("http://vllm:8000", "mistral-small") + assert result.status == EndpointStatus.HEALTHY + + @pytest.mark.asyncio + async def test_model_not_found(self, models_response): + transport = httpx.MockTransport(lambda req: models_response) + with patch( + "services.inference.healthcheck.httpx.AsyncClient", return_value=httpx.AsyncClient(transport=transport) + ): + result = await check_model_available("http://vllm:8000", "nonexistent-model") + assert result.status == EndpointStatus.UNHEALTHY + assert "nonexistent-model" in result.error + assert "mistral-small" in result.error + + @pytest.mark.asyncio + async def test_endpoint_unreachable_skips_model_check(self): + async def raise_connect_error(*a, **kw): + raise httpx.ConnectError("refused") + + client = AsyncMock() + client.get = raise_connect_error + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=False) + with patch("services.inference.healthcheck.httpx.AsyncClient", return_value=client): + result = await check_model_available("http://vllm:8000", "mistral-small") + assert result.status == EndpointStatus.UNREACHABLE diff --git a/openrag/services/inference/test_reranker_clients.py b/openrag/services/inference/test_reranker_clients.py new file mode 100644 index 000000000..860dd94c8 --- /dev/null +++ b/openrag/services/inference/test_reranker_clients.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +from unittest.mock import AsyncMock + +import httpx +import pytest +from core.utils.exceptions import InferenceConnectionError, InferenceTimeoutError + +from .reranker_clients import InfinityReranker, OpenAIReranker + + +def _rerank_response(results: list[dict] | None = None) -> httpx.Response: + results = results or [ + {"index": 0, "relevance_score": 0.9}, + {"index": 2, "relevance_score": 0.7}, + {"index": 1, "relevance_score": 0.3}, + ] + return httpx.Response(200, json={"results": results}) + + +DOCS = ["doc zero", "doc one", "doc two"] + + +class TestInfinityReranker: + @pytest.fixture + def reranker(self): + return InfinityReranker(endpoint="http://reranker:7997", model_name="gte-reranker") + + @pytest.mark.asyncio + async def test_rerank(self, reranker): + transport = httpx.MockTransport(lambda req: _rerank_response()) + reranker._client = httpx.AsyncClient(transport=transport) + result = await reranker.rerank("query", DOCS) + assert result == [(0, 0.9), (2, 0.7), (1, 0.3)] + + @pytest.mark.asyncio + async def test_rerank_with_top_k(self, reranker): + captured = {} + + def capture(req): + import json + + captured.update(json.loads(req.content)) + return _rerank_response([{"index": 0, "relevance_score": 0.9}]) + + transport = httpx.MockTransport(capture) + reranker._client = httpx.AsyncClient(transport=transport) + result = await reranker.rerank("query", DOCS, top_k=1) + assert captured["top_n"] == 1 + assert len(result) == 1 + + @pytest.mark.asyncio + async def test_top_k_clamped_to_doc_count(self, reranker): + captured = {} + + def capture(req): + import json + + captured.update(json.loads(req.content)) + return _rerank_response() + + transport = httpx.MockTransport(capture) + reranker._client = httpx.AsyncClient(transport=transport) + await reranker.rerank("query", DOCS, top_k=100) + assert captured["top_n"] == 3 + + @pytest.mark.asyncio + async def test_sends_raw_scores(self, reranker): + captured = {} + + def capture(req): + import json + + captured.update(json.loads(req.content)) + return _rerank_response() + + transport = httpx.MockTransport(capture) + reranker._client = httpx.AsyncClient(transport=transport) + await reranker.rerank("query", DOCS) + assert captured["raw_scores"] is True + assert captured["return_documents"] is False + + @pytest.mark.asyncio + async def test_connection_error(self, reranker): + async def fail(*a, **kw): + raise httpx.ConnectError("refused") + + reranker._client = AsyncMock() + reranker._client.post = fail + with pytest.raises(InferenceConnectionError): + await reranker.rerank("query", DOCS) + + @pytest.mark.asyncio + async def test_timeout(self, reranker): + async def fail(*a, **kw): + raise httpx.TimeoutException("timeout") + + reranker._client = AsyncMock() + reranker._client.post = fail + with pytest.raises(InferenceTimeoutError): + await reranker.rerank("query", DOCS) + + @pytest.mark.asyncio + async def test_trailing_slash_stripped(self): + r = InfinityReranker(endpoint="http://reranker:7997/", model_name="m") + assert r._endpoint == "http://reranker:7997" + await r.aclose() + + +class TestOpenAIReranker: + @pytest.fixture + def reranker(self): + return OpenAIReranker(endpoint="http://reranker:8000/v1", model_name="gte-reranker", api_key="k") + + @pytest.mark.asyncio + async def test_rerank(self, reranker): + transport = httpx.MockTransport(lambda req: _rerank_response()) + reranker._client = httpx.AsyncClient(transport=transport) + result = await reranker.rerank("query", DOCS) + assert result == [(0, 0.9), (2, 0.7), (1, 0.3)] + + @pytest.mark.asyncio + async def test_connection_error(self, reranker): + async def fail(*a, **kw): + raise httpx.ConnectError("refused") + + reranker._client = AsyncMock() + reranker._client.post = fail + with pytest.raises(InferenceConnectionError): + await reranker.rerank("query", DOCS) + + @pytest.mark.asyncio + async def test_timeout(self, reranker): + async def fail(*a, **kw): + raise httpx.TimeoutException("timeout") + + reranker._client = AsyncMock() + reranker._client.post = fail + with pytest.raises(InferenceTimeoutError): + await reranker.rerank("query", DOCS) + + +class TestRegistryIntegration: + def test_infinity_registered(self): + from core.rerankers import reranker_registry + + assert "infinity" in reranker_registry + + def test_openai_registered(self): + from core.rerankers import reranker_registry + + assert "openai" in reranker_registry diff --git a/openrag/services/inference/test_retry.py b/openrag/services/inference/test_retry.py new file mode 100644 index 000000000..222e6c5d4 --- /dev/null +++ b/openrag/services/inference/test_retry.py @@ -0,0 +1,127 @@ +import httpx +import pytest +from core.utils.exceptions import OpenRAGError, ServiceUnavailableError +from services.inference._retry import _is_retryable, with_retry + + +class TestIsRetryable: + def test_timeout_exception(self): + assert _is_retryable(httpx.ReadTimeout("timeout")) + + def test_connect_error(self): + assert _is_retryable(httpx.ConnectError("refused")) + + @pytest.mark.parametrize("status", [429, 502, 503, 504]) + def test_retryable_http_status(self, status): + req = httpx.Request("GET", "http://test") + exc = httpx.HTTPStatusError("err", request=req, response=httpx.Response(status, request=req)) + assert _is_retryable(exc) + + @pytest.mark.parametrize("status", [400, 401, 403, 404, 422, 500]) + def test_non_retryable_http_status(self, status): + req = httpx.Request("GET", "http://test") + exc = httpx.HTTPStatusError("err", request=req, response=httpx.Response(status, request=req)) + assert not _is_retryable(exc) + + def test_openrag_error_retryable(self): + assert _is_retryable(ServiceUnavailableError("down")) # 503 + + def test_openrag_error_non_retryable(self): + assert not _is_retryable(OpenRAGError("bad", status_code=404)) + + def test_unrelated_exception(self): + assert not _is_retryable(ValueError("nope")) + + +class TestWithRetry: + @pytest.mark.asyncio + async def test_success_no_retry(self): + call_count = 0 + + @with_retry(max_attempts=3, base_wait=0.01, max_wait=0.1) + async def succeed(): + nonlocal call_count + call_count += 1 + return "ok" + + result = await succeed() + assert result == "ok" + assert call_count == 1 + + @pytest.mark.asyncio + async def test_retries_on_connect_error(self): + call_count = 0 + + @with_retry(max_attempts=3, base_wait=0.01, max_wait=0.1) + async def fail_connect(): + nonlocal call_count + call_count += 1 + raise httpx.ConnectError("refused") + + with pytest.raises(httpx.ConnectError): + await fail_connect() + + assert call_count == 3 + + @pytest.mark.asyncio + async def test_retries_on_429(self): + call_count = 0 + req = httpx.Request("GET", "http://test") + + @with_retry(max_attempts=3, base_wait=0.01, max_wait=0.1) + async def fail_429(): + nonlocal call_count + call_count += 1 + raise httpx.HTTPStatusError("rate limited", request=req, response=httpx.Response(429, request=req)) + + with pytest.raises(httpx.HTTPStatusError): + await fail_429() + + assert call_count == 3 + + @pytest.mark.asyncio + async def test_no_retry_on_400(self): + call_count = 0 + req = httpx.Request("GET", "http://test") + + @with_retry(max_attempts=3, base_wait=0.01, max_wait=0.1) + async def fail_400(): + nonlocal call_count + call_count += 1 + raise httpx.HTTPStatusError("bad request", request=req, response=httpx.Response(400, request=req)) + + with pytest.raises(httpx.HTTPStatusError): + await fail_400() + + assert call_count == 1 + + @pytest.mark.asyncio + async def test_succeeds_after_transient_failure(self): + call_count = 0 + + @with_retry(max_attempts=3, base_wait=0.01, max_wait=0.1) + async def flaky(): + nonlocal call_count + call_count += 1 + if call_count < 3: + raise httpx.ConnectError("transient") + return "recovered" + + result = await flaky() + assert result == "recovered" + assert call_count == 3 + + @pytest.mark.asyncio + async def test_retries_openrag_503(self): + call_count = 0 + + @with_retry(max_attempts=3, base_wait=0.01, max_wait=0.1) + async def fail_503(): + nonlocal call_count + call_count += 1 + raise ServiceUnavailableError("down") + + with pytest.raises(ServiceUnavailableError): + await fail_503() + + assert call_count == 3 diff --git a/openrag/services/inference/test_vllm_client.py b/openrag/services/inference/test_vllm_client.py new file mode 100644 index 000000000..797bb2760 --- /dev/null +++ b/openrag/services/inference/test_vllm_client.py @@ -0,0 +1,439 @@ +from __future__ import annotations + +import json +from unittest.mock import AsyncMock + +import httpx +import pytest +from core.utils.exceptions import ( + EmbeddingAPIError, + EmbeddingResponseError, + InferenceConnectionError, + InferenceError, + InferenceTimeoutError, +) +from services.inference._circuit_breaker import _breakers + +from .vllm_client import VLLMClient, VLLMEmbedder, VLLMVision + + +@pytest.fixture(autouse=True) +def _clean_breakers(): + yield + for breaker in _breakers.values(): + breaker.close() + _breakers.clear() + + +def _make_transport(handler): + return httpx.MockTransport(handler) + + +def _chat_response(content: str = "hello") -> httpx.Response: + return httpx.Response(200, json={"choices": [{"message": {"content": content}}]}) + + +def _completions_response(text: str = "result") -> httpx.Response: + return httpx.Response(200, json={"choices": [{"text": text}]}) + + +def _embed_response(vectors: list[list[float]] | None = None) -> httpx.Response: + vectors = vectors or [[0.1, 0.2, 0.3]] + data = [{"index": i, "embedding": v} for i, v in enumerate(vectors)] + return httpx.Response(200, json={"data": data}) + + +# --------------------------------------------------------------------------- +# VLLMClient (LLM) +# --------------------------------------------------------------------------- + + +class TestVLLMClient: + def _make_client(self, handler, **kwargs): + client = VLLMClient( + endpoint="http://vllm:8000/v1", + model_name="test-model", + api_key="test-key", + temperature=0.3, + **kwargs, + ) + client._client = httpx.AsyncClient(transport=_make_transport(handler)) + return client + + @pytest.mark.asyncio + async def test_chat_returns_full_response(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert "/chat/completions" in str(request.url) + assert body["model"] == "test-model" + assert body["stream"] is False + assert body["temperature"] == 0.3 + return _chat_response("world") + + result = await self._make_client(handler).chat([{"role": "user", "content": "hi"}]) + assert result["choices"][0]["message"]["content"] == "world" + + @pytest.mark.asyncio + async def test_generate_returns_full_response(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert "/completions" in str(request.url) + assert "/chat/" not in str(request.url) + assert body["prompt"] == "say something" + return _completions_response("done") + + result = await self._make_client(handler).generate("say something") + assert result["choices"][0]["text"] == "done" + + @pytest.mark.asyncio + async def test_stream_chat_yields_raw_sse_lines(self): + sse_body = ( + 'data: {"choices":[{"delta":{"content":"Hello"}}]}\n' + 'data: {"choices":[{"delta":{"content":" world"}}]}\n' + "data: [DONE]\n" + ) + + def handler(request: httpx.Request) -> httpx.Response: + assert json.loads(request.content)["stream"] is True + return httpx.Response(200, text=sse_body) + + client = self._make_client(handler) + lines = [line async for line in client.stream_chat([{"role": "user", "content": "hi"}])] + assert 'data: {"choices":[{"delta":{"content":"Hello"}}]}' in lines + assert 'data: {"choices":[{"delta":{"content":" world"}}]}' in lines + + @pytest.mark.asyncio + async def test_stream_chat_error_raises(self): + client = self._make_client(lambda req: httpx.Response(503, text="unavailable")) + with pytest.raises(InferenceError): + async for _ in client.stream_chat([{"role": "user", "content": "hi"}]): + pass + + @pytest.mark.asyncio + async def test_chat_connection_error(self): + async def fail(*a, **kw): + raise httpx.ConnectError("refused") + + client = VLLMClient(endpoint="http://vllm:8000/v1", model_name="m") + client._client = AsyncMock() + client._client.post = fail + with pytest.raises(InferenceConnectionError): + await client.chat([{"role": "user", "content": "hi"}]) + + @pytest.mark.asyncio + async def test_chat_timeout(self): + async def fail(*a, **kw): + raise httpx.TimeoutException("timeout") + + client = VLLMClient(endpoint="http://vllm:8000/v1", model_name="m") + client._client = AsyncMock() + client._client.post = fail + with pytest.raises(InferenceTimeoutError): + await client.chat([{"role": "user", "content": "hi"}]) + + @pytest.mark.asyncio + async def test_defaults_forwarded(self): + captured: dict = {} + + def capture(req: httpx.Request) -> httpx.Response: + captured.update(json.loads(req.content)) + return _chat_response() + + await self._make_client(capture).chat([{"role": "user", "content": "hi"}]) + assert captured["temperature"] == 0.3 + + @pytest.mark.asyncio + async def test_per_request_kwargs_override_defaults(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert body["temperature"] == 0.9 + assert body["max_tokens"] == 100 + return _chat_response() + + await self._make_client(handler).chat([{"role": "user", "content": "hi"}], temperature=0.9, max_tokens=100) + + @pytest.mark.asyncio + async def test_trailing_slash_stripped(self): + c = VLLMClient(endpoint="http://vllm:8000/v1/", model_name="m") + assert c._endpoint == "http://vllm:8000/v1" + await c.aclose() + + @pytest.mark.asyncio + async def test_aclose(self): + client = VLLMClient(endpoint="http://vllm:8000/v1", model_name="m") + client._client = AsyncMock() + await client.aclose() + client._client.aclose.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# VLLMClientOverrides +# --------------------------------------------------------------------------- + + +class TestVLLMClientOverrides: + """Tests for _resolve_overrides (partition-level model selection).""" + + def _make_client(self): + return VLLMClient( + endpoint="http://default:8000/v1", + model_name="default-model", + api_key="default-key", + ) + + def test_no_override_uses_defaults(self): + client = self._make_client() + kwargs: dict = {} + base_url, model, headers = client._resolve_overrides(kwargs) + assert base_url == "http://default:8000/v1" + assert model == "default-model" + assert headers is None + + def test_llm_override_in_metadata(self): + client = self._make_client() + original_metadata = { + "llm_override": { + "base_url": "http://custom:9000/v1/", + "api_key": "custom-key", + "model": "custom-model", + }, + } + kwargs: dict = {"metadata": original_metadata} + base_url, model, headers = client._resolve_overrides(kwargs) + assert base_url == "http://custom:9000/v1" + assert model == "custom-model" + assert headers is not None + assert headers["Authorization"] == "Bearer custom-key" + # kwargs must not be mutated — retries depend on llm_override surviving. + assert kwargs["metadata"] is original_metadata + assert "llm_override" in kwargs["metadata"] + + def test_llm_override_partial(self): + client = self._make_client() + original_metadata = { + "llm_override": {"model": "override-model"}, + "use_map_reduce": True, + } + kwargs: dict = {"metadata": original_metadata} + base_url, model, headers = client._resolve_overrides(kwargs) + assert base_url == "http://default:8000/v1" + assert model == "override-model" + assert headers is None + assert kwargs["metadata"] is original_metadata + assert kwargs["metadata"] == { + "llm_override": {"model": "override-model"}, + "use_map_reduce": True, + } + + def test_trailing_slash_stripped(self): + client = self._make_client() + kwargs: dict = {"metadata": {"llm_override": {"base_url": "http://custom:9000/v1///"}}} + base_url, _, _ = client._resolve_overrides(kwargs) + assert base_url == "http://custom:9000/v1" + + +# --------------------------------------------------------------------------- +# VLLMEmbedder +# --------------------------------------------------------------------------- + + +class TestVLLMEmbedder: + def _make_embedder(self, handler, **kwargs): + embedder = VLLMEmbedder( + endpoint="http://vllm:8000/v1", + model_name="bge-m3", + api_key="test-key", + **kwargs, + ) + embedder._client = httpx.AsyncClient(transport=_make_transport(handler)) + return embedder + + @pytest.mark.asyncio + async def test_embed_returns_sorted_vectors(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert body["model"] == "bge-m3" + assert body["input"] == ["hello", "world"] + return httpx.Response( + 200, + json={"data": [{"index": 1, "embedding": [0.3, 0.4]}, {"index": 0, "embedding": [0.1, 0.2]}]}, + ) + + result = await self._make_embedder(handler).embed(["hello", "world"]) + assert result == [[0.1, 0.2], [0.3, 0.4]] + + @pytest.mark.asyncio + async def test_embed_single(self): + def handler(req: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.5, 0.6, 0.7]}]}) + + result = await self._make_embedder(handler).embed_single("test") + assert result == [0.5, 0.6, 0.7] + + @pytest.mark.asyncio + async def test_dimension_auto_detected(self): + def handler(req: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.1, 0.2, 0.3]}]}) + + embedder = self._make_embedder(handler) + + with pytest.raises(RuntimeError, match="unknown"): + _ = embedder.dimension + + await embedder.embed(["test"]) + assert embedder.dimension == 3 + + def test_dimension_from_init(self): + assert VLLMEmbedder(endpoint="http://x", model_name="m", dimension=768).dimension == 768 + + @pytest.mark.asyncio + async def test_truncate_prompt_tokens_included(self): + def handler(request: httpx.Request) -> httpx.Response: + assert json.loads(request.content)["truncate_prompt_tokens"] == 8192 + return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.1]}]}) + + await self._make_embedder(handler, max_model_len=8192).embed(["test"]) + + @pytest.mark.asyncio + async def test_truncate_prompt_tokens_absent_when_none(self): + def handler(request: httpx.Request) -> httpx.Response: + assert "truncate_prompt_tokens" not in json.loads(request.content) + return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.1]}]}) + + await self._make_embedder(handler).embed(["test"]) + + @pytest.mark.asyncio + async def test_embed_connection_error(self): + async def fail(*a, **kw): + raise httpx.ConnectError("refused") + + embedder = VLLMEmbedder(endpoint="http://vllm:8000/v1", model_name="bge-m3") + embedder._client = AsyncMock() + embedder._client.post = fail + with pytest.raises(EmbeddingAPIError): + await embedder.embed(["text"]) + + @pytest.mark.asyncio + async def test_embed_timeout(self): + async def fail(*a, **kw): + raise httpx.TimeoutException("timeout") + + embedder = VLLMEmbedder(endpoint="http://vllm:8000/v1", model_name="bge-m3") + embedder._client = AsyncMock() + embedder._client.post = fail + with pytest.raises(EmbeddingAPIError): + await embedder.embed(["text"]) + + @pytest.mark.asyncio + async def test_embed_bad_response_format(self): + def handler(_req: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"wrong": "shape"}) + + with pytest.raises(EmbeddingResponseError): + await self._make_embedder(handler).embed(["text"]) + + +# --------------------------------------------------------------------------- +# VLLMVision +# --------------------------------------------------------------------------- + + +class TestVLLMVision: + def _make_vision(self, handler, **kwargs): + vision = VLLMVision( + endpoint="http://vllm:8000/v1", + model_name="qwen-vl", + api_key="test-key", + **kwargs, + ) + vision._client = httpx.AsyncClient(transport=_make_transport(handler)) + return vision + + @pytest.mark.asyncio + async def test_caption_image(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert body["model"] == "qwen-vl" + assert body["max_tokens"] == 1024 + msg = body["messages"][0] + assert msg["content"][0]["type"] == "image_url" + assert msg["content"][0]["image_url"]["url"].startswith("data:image/png;base64,") + assert msg["content"][1]["type"] == "text" + return _chat_response("A red car") + + result = await self._make_vision(handler).caption_image(b"\x89PNG\r\n\x1a\n", prompt="What is this?") + assert result == "A red car" + + @pytest.mark.asyncio + async def test_caption_image_default_prompt(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert body["messages"][0]["content"][1]["text"] == "Describe this image in detail." + return _chat_response("An image") + + await self._make_vision(handler).caption_image(b"\x89PNG\r\n\x1a\n") + + @pytest.mark.asyncio + async def test_caption_images_batch(self): + call_count = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal call_count + call_count += 1 + return _chat_response(f"Caption {call_count}") + + results = await self._make_vision(handler).caption_images_batch([b"img1", b"img2", b"img3"]) + assert len(results) == 3 + assert call_count == 3 + + @pytest.mark.asyncio + async def test_custom_max_tokens(self): + def handler(request: httpx.Request) -> httpx.Response: + assert json.loads(request.content)["max_tokens"] == 512 + return _chat_response("ok") + + await self._make_vision(handler, max_tokens=512).caption_image(b"img") + + @pytest.mark.asyncio + async def test_caption_connection_error(self): + async def fail(*a, **kw): + raise httpx.ConnectError("refused") + + vision = VLLMVision(endpoint="http://vllm:8000/v1", model_name="qwen-vl") + vision._client = AsyncMock() + vision._client.post = fail + with pytest.raises(InferenceConnectionError): + await vision.caption_image(b"img") + + @pytest.mark.asyncio + async def test_caption_timeout(self): + async def fail(*a, **kw): + raise httpx.TimeoutException("timeout") + + vision = VLLMVision(endpoint="http://vllm:8000/v1", model_name="qwen-vl") + vision._client = AsyncMock() + vision._client.post = fail + with pytest.raises(InferenceTimeoutError): + await vision.caption_image(b"img") + + +# --------------------------------------------------------------------------- +# Registry integration +# --------------------------------------------------------------------------- + + +class TestRegistryIntegration: + def test_llm_registered(self): + from core.llm import llm_registry + + assert "vllm" in llm_registry + + def test_embedder_registered(self): + from core.embeddings import embedder_registry + + assert "vllm" in embedder_registry + + def test_vlm_registered(self): + from core.vlm import vlm_registry + + assert "vllm" in vlm_registry diff --git a/openrag/services/inference/vllm_client.py b/openrag/services/inference/vllm_client.py new file mode 100644 index 000000000..1b6fd9e4b --- /dev/null +++ b/openrag/services/inference/vllm_client.py @@ -0,0 +1,323 @@ +"""vLLM / OpenAI-compatible inference clients. + +Three classes grouped by server — all talk to the same OpenAI-compatible API: + +* ``VLLMClient`` → ``LLM`` (chat completions) +* ``VLLMEmbedder`` → ``Embedder`` (embeddings) +* ``VLLMVision`` → ``VLM`` (image captioning via chat completions) + +Each class has its own circuit breaker instance so an embedder outage +doesn't trip the LLM breaker. +""" + +from __future__ import annotations + +import asyncio +import base64 +from collections.abc import AsyncIterator + +import httpx +from core.embeddings import Embedder, embedder_registry +from core.llm import LLM, llm_registry +from core.utils.exceptions import ( + EmbeddingAPIError, + EmbeddingResponseError, + InferenceConnectionError, + InferenceError, + InferenceTimeoutError, +) +from core.vlm import VLM, vlm_registry +from utils.logger import get_logger + +from ._circuit_breaker import with_circuit_breaker +from ._retry import with_retry + +logger = get_logger() + + +def _parse_response(resp: httpx.Response) -> dict: + try: + return resp.json() + except ValueError as e: + raise InferenceError(f"Invalid JSON from inference server ({resp.url}): {e}", status_code=502) from e + + +# --------------------------------------------------------------------------- +# LLM +# --------------------------------------------------------------------------- + + +@llm_registry.register("vllm") +class VLLMClient(LLM): + """OpenAI-compatible LLM client backed by vLLM. + + *endpoint* should include the version prefix, e.g. ``http://vllm:8000/v1``. + A single long-lived ``httpx.AsyncClient`` is reused across requests for + connection pooling. + """ + + def __init__( + self, + endpoint: str, + model_name: str, + *, + api_key: str = "", + timeout: float = 240.0, + **kwargs, + ) -> None: + self._endpoint = endpoint.rstrip("/") + self._model = model_name + self._api_key = api_key + self._defaults: dict = kwargs + headers: dict[str, str] = {"Content-Type": "application/json"} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + self._client = httpx.AsyncClient(timeout=timeout, headers=headers) + + def _resolve_overrides(self, kwargs: dict) -> tuple[str, str, dict[str, str] | None]: + """Read ``metadata.llm_override`` from *kwargs* without mutating caller data. + + Pure read: ``kwargs`` is untouched so retries see the original override on + every attempt. The caller strips ``metadata`` from the outbound payload via + ``_payload_kwargs`` — every ``metadata`` key is OpenRAG-internal and never + belongs on the wire. + """ + base_url = self._endpoint + model = self._model + override_headers: dict[str, str] | None = None + + llm_override = (kwargs.get("metadata") or {}).get("llm_override") or {} + if llm_override: + if llm_override.get("base_url"): + base_url = llm_override["base_url"].rstrip("/") + if llm_override.get("model"): + model = llm_override["model"] + if llm_override.get("api_key"): + override_headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {llm_override['api_key']}", + } + + return base_url, model, override_headers + + @with_circuit_breaker("llm") + @with_retry(max_attempts=3) + async def generate(self, prompt: str, **kwargs) -> dict: + base_url, model, headers = self._resolve_overrides(kwargs) + kwargs.pop("metadata", None) + payload = {**self._defaults, **kwargs, "model": model, "prompt": prompt} + try: + resp = await self._client.post(f"{base_url}/completions", json=payload, headers=headers) + resp.raise_for_status() + except httpx.ConnectError as exc: + raise InferenceConnectionError(f"Cannot reach LLM at {base_url}") from exc + except httpx.TimeoutException as exc: + raise InferenceTimeoutError(f"LLM request timed out at {base_url}") from exc + except httpx.HTTPStatusError as exc: + raise InferenceError( + f"LLM error ({exc.response.status_code}): {exc.response.text[:500]}", + status_code=exc.response.status_code, + ) from exc + return _parse_response(resp) + + @with_circuit_breaker("llm") + @with_retry(max_attempts=3) + async def chat(self, messages: list[dict[str, str]], **kwargs) -> dict: + base_url, model, headers = self._resolve_overrides(kwargs) + kwargs.pop("metadata", None) + payload = {**self._defaults, **kwargs, "model": model, "messages": messages, "stream": False} + try: + resp = await self._client.post(f"{base_url}/chat/completions", json=payload, headers=headers) + resp.raise_for_status() + except httpx.ConnectError as exc: + raise InferenceConnectionError(f"Cannot reach LLM at {base_url}") from exc + except httpx.TimeoutException as exc: + raise InferenceTimeoutError(f"LLM request timed out at {base_url}") from exc + except httpx.HTTPStatusError as exc: + raise InferenceError( + f"LLM error ({exc.response.status_code}): {exc.response.text[:500]}", + status_code=exc.response.status_code, + ) from exc + return _parse_response(resp) + + async def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> AsyncIterator[str]: + base_url, model, headers = self._resolve_overrides(kwargs) + kwargs.pop("metadata", None) + payload = {**self._defaults, **kwargs, "model": model, "messages": messages, "stream": True} + try: + async with self._client.stream( + "POST", f"{base_url}/chat/completions", json=payload, headers=headers + ) as resp: + if resp.status_code >= 400: + await resp.aread() + raise InferenceError( + f"LLM streaming error ({resp.status_code}): {resp.text[:500]}", + status_code=resp.status_code, + ) + async for line in resp.aiter_lines(): + yield line + except httpx.ConnectError as exc: + raise InferenceConnectionError(f"Cannot reach LLM at {base_url}") from exc + except httpx.TimeoutException as exc: + raise InferenceTimeoutError(f"LLM streaming request timed out at {base_url}") from exc + + async def aclose(self) -> None: + await self._client.aclose() + + +# --------------------------------------------------------------------------- +# Embedder +# --------------------------------------------------------------------------- + + +@embedder_registry.register("vllm") +class VLLMEmbedder(Embedder): + """OpenAI-compatible embedding client backed by vLLM. + + Replaces the sync ``openai.OpenAI`` SDK with an async ``httpx`` client. + """ + + def __init__( + self, + endpoint: str, + model_name: str, + *, + max_model_len: int | None = None, + dimension: int | None = None, + timeout: float = 60.0, + api_key: str = "", + **_kwargs, + ) -> None: + self._endpoint = endpoint.rstrip("/") + self._model = model_name + self._max_model_len = max_model_len + self._dimension: int | None = dimension + headers: dict[str, str] = {} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + self._client = httpx.AsyncClient(timeout=timeout, headers=headers) + + @with_circuit_breaker("embedder") + @with_retry(max_attempts=3) + async def embed(self, texts: list[str]) -> list[list[float]]: + body: dict = {"model": self._model, "input": texts} + if self._max_model_len is not None: + body["truncate_prompt_tokens"] = self._max_model_len + try: + resp = await self._client.post(f"{self._endpoint}/embeddings", json=body) + resp.raise_for_status() + except httpx.ConnectError as exc: + raise EmbeddingAPIError( + f"Cannot reach embedder at {self._endpoint}", + model_name=self._model, + base_url=self._endpoint, + error=str(exc), + ) from exc + except httpx.TimeoutException as exc: + raise EmbeddingAPIError( + f"Embedder request timed out at {self._endpoint}", + model_name=self._model, + base_url=self._endpoint, + error=str(exc), + ) from exc + except httpx.HTTPStatusError as exc: + raise EmbeddingAPIError( + f"Embedder API error ({exc.response.status_code})", + model_name=self._model, + base_url=self._endpoint, + error=exc.response.text, + ) from exc + + try: + data = resp.json()["data"] + embeddings = [item["embedding"] for item in sorted(data, key=lambda x: x["index"])] + except (ValueError, KeyError, IndexError, TypeError) as exc: + raise EmbeddingResponseError( + "Unexpected embedding response format", + model_name=self._model, + base_url=self._endpoint, + error=str(exc), + ) from exc + + if self._dimension is None and embeddings: + self._dimension = len(embeddings[0]) + return embeddings + + async def embed_single(self, text: str) -> list[float]: + result = await self.embed([text]) + return result[0] + + @property + def dimension(self) -> int: + if self._dimension is None: + raise RuntimeError("Embedding dimension unknown — call embed() first") + return self._dimension + + async def aclose(self) -> None: + await self._client.aclose() + + +# --------------------------------------------------------------------------- +# VLM (Vision-Language Model) +# --------------------------------------------------------------------------- + + +@vlm_registry.register("vllm") +class VLLMVision(VLLMClient, VLM): + """OpenAI-compatible vision client backed by vLLM. + + Inherits connection pooling, retry, and circuit breaker from VLLMClient. + Adds image captioning via the same OpenAI-compatible chat/completions endpoint. + """ + + def __init__( + self, + endpoint: str, + model_name: str, + *, + timeout: float = 60.0, + api_key: str = "", + max_tokens: int = 1024, + **kwargs, + ) -> None: + super().__init__(endpoint=endpoint, model_name=model_name, api_key=api_key, timeout=timeout, **kwargs) + self._max_tokens = max_tokens + + @with_circuit_breaker("vlm") + @with_retry(max_attempts=2) + async def caption_image(self, image_bytes: bytes, prompt: str | None = None) -> str: + image_b64 = base64.b64encode(image_bytes).decode() + messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{image_b64}"}, + }, + { + "type": "text", + "text": prompt or "Describe this image in detail.", + }, + ], + } + ] + try: + resp = await self._client.post( + f"{self._endpoint}/chat/completions", + json={"model": self._model, "messages": messages, "max_tokens": self._max_tokens}, + ) + resp.raise_for_status() + except httpx.ConnectError as exc: + raise InferenceConnectionError(f"Cannot reach VLM at {self._endpoint}") from exc + except httpx.TimeoutException as exc: + raise InferenceTimeoutError(f"VLM request timed out at {self._endpoint}") from exc + except httpx.HTTPStatusError as exc: + raise InferenceError( + f"VLM error ({exc.response.status_code}): {exc.response.text[:500]}", + status_code=exc.response.status_code, + ) from exc + return _parse_response(resp)["choices"][0]["message"]["content"] + + async def caption_images_batch(self, images: list[bytes], prompt: str | None = None) -> list[str]: + return list(await asyncio.gather(*(self.caption_image(img, prompt) for img in images))) diff --git a/openrag/utils/dependencies.py b/openrag/utils/dependencies.py index ee21cd004..9141edebb 100644 --- a/openrag/utils/dependencies.py +++ b/openrag/utils/dependencies.py @@ -7,8 +7,8 @@ from components.indexer.loaders.pdf_loaders.marker import MarkerPool from components.indexer.loaders.serializer import DocSerializer from components.indexer.vectordb.vectordb import ConnectorFactory -from components.utils import DistributedSemaphoreActor from config import load_config +from services.inference.distributed_semaphore import DistributedSemaphoreActor from utils.logger import get_logger # load config diff --git a/openrag/utils/external_resource_errors.py b/openrag/utils/external_resource_errors.py index 3752a0e3f..49f930a26 100644 --- a/openrag/utils/external_resource_errors.py +++ b/openrag/utils/external_resource_errors.py @@ -1,65 +1,13 @@ -""" -Utilities for detecting external resource access errors. +"""Re-export from canonical location for backwards compatibility.""" -When VLM models fetch external image URLs, HTTP errors (403, 404, etc.) from -remote servers get wrapped in InternalServerError, which is misleading. -This module detects such errors for better logging. -""" - -import re - -# HTTP error codes indicating external resource issues -EXTERNAL_ERROR_CODES = frozenset( - { - # 4xx client errors - "400", - "401", - "403", - "404", - "405", - "408", - "410", - "429", - "451", - # 5xx gateway errors - "502", - "503", - "504", - } +from core.utils.external_errors import ( + EXTERNAL_ERROR_CODES, + EXTERNAL_ERROR_INDICATORS, + is_external_resource_error, ) -# Error type indicators for external fetch failures -EXTERNAL_ERROR_INDICATORS = ( - "ClientResponseError", - "HTTPError", - "ConnectionError", - "TimeoutError", - "SSLError", -) - - -def is_external_resource_error(error: Exception) -> tuple[bool, str, str]: - """ - Check if an error is caused by an external resource access issue. - - Returns: - (is_external_error, status_code, url) - status_code and url are empty - strings if not detected. - """ - error_str = str(error) - - # Find HTTP 4xx/5xx status code (first one that's in our allowed set) - status_code = "" - for match in re.finditer(r"\b([45]\d{2})\b", error_str): - if match.group(1) in EXTERNAL_ERROR_CODES: - status_code = match.group(1) - break - - # Extract URL - url_match = re.search(r"https?://[^\s'\"\)>]+", error_str) - url = url_match.group(0) if url_match else "" - - # Check for error type indicators - has_indicator = any(ind in error_str for ind in EXTERNAL_ERROR_INDICATORS) - - return bool(status_code) or has_indicator, status_code, url +__all__ = [ + "EXTERNAL_ERROR_CODES", + "EXTERNAL_ERROR_INDICATORS", + "is_external_resource_error", +] diff --git a/pyproject.toml b/pyproject.toml index 38ee95b43..782007a32 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,8 @@ dependencies = [ "authlib>=1.3", "itsdangerous>=2.2", "cryptography>=42", + "tenacity>=8.2.0", + "aiobreaker>=1.2.0", ] [dependency-groups] diff --git a/uv.lock b/uv.lock index eb51ee9ae..c259e8700 100644 --- a/uv.lock +++ b/uv.lock @@ -6,6 +6,15 @@ resolution-markers = [ "python_full_version < '3.13'", ] +[[package]] +name = "aiobreaker" +version = "1.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7f/eb/749ef48d3227fd62d500ff01fcd451f10111e00d822c200eb51782ba076a/aiobreaker-1.2.0.tar.gz", hash = "sha256:217a9cfa12e520bb2dd1934bace281d1d7deb8d7630dd183a6295fd22e323ce7", size = 15947, upload-time = "2021-05-17T11:58:05.278Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ff/c8/4cd4b2834012ffc71ae3fd69187f08a17f01f3937527b6b5e077f4f5d0db/aiobreaker-1.2.0-py3-none-any.whl", hash = "sha256:f275decad78bdd161715afeee67e5dde7967de54c836648b44f4eea1b5e41d60", size = 20700, upload-time = "2021-05-17T11:58:04.192Z" }, +] + [[package]] name = "aiofile" version = "3.9.0" @@ -2635,6 +2644,7 @@ name = "openrag" version = "1.1.10" source = { editable = "." } dependencies = [ + { name = "aiobreaker" }, { name = "aiopath" }, { name = "alembic" }, { name = "asyncpg" }, @@ -2683,6 +2693,7 @@ dependencies = [ { name = "ruff" }, { name = "spire-doc" }, { name = "sqlalchemy-utils" }, + { name = "tenacity" }, { name = "torch" }, { name = "umap-learn" }, ] @@ -2699,6 +2710,7 @@ lint = [ [package.metadata] requires-dist = [ + { name = "aiobreaker", specifier = ">=1.2.0" }, { name = "aiopath", specifier = ">=0.7.7" }, { name = "alembic", specifier = ">=1.17.0" }, { name = "asyncpg", specifier = ">=0.30.0" }, @@ -2747,6 +2759,7 @@ requires-dist = [ { name = "ruff", specifier = ">=0.14.1" }, { name = "spire-doc", specifier = ">=13.1.0" }, { name = "sqlalchemy-utils" }, + { name = "tenacity", specifier = ">=8.2.0" }, { name = "torch", specifier = ">=2.4.1" }, { name = "umap-learn", specifier = ">=0.5.9.post2" }, ]