Skip to content

refactor(phase 6): inference adapters - #375

Merged
Ahmath-Gadji merged 12 commits into
refactor/hexagonalfrom
refactor/phase-6-inference-adapters
May 11, 2026
Merged

refactor(phase 6): inference adapters#375
Ahmath-Gadji merged 12 commits into
refactor/hexagonalfrom
refactor/phase-6-inference-adapters

Conversation

@Ahmath-Gadji

@Ahmath-Gadji Ahmath-Gadji commented May 7, 2026

Copy link
Copy Markdown
Collaborator

Summary by CodeRabbit

  • New Features
    • Improved reliability: async circuit breaker + retry behavior, cluster-wide distributed semaphores, endpoint readiness health checks, vLLM-based LLM/embedding/vision clients, two reranker backends, and service registration helpers.
  • Refactor
    • LLM APIs now return structured response payloads (dict); legacy shims preserve older behavior.
  • Chores
    • External-error utilities re-exported; dependency updates for circuit-breaker runtime.
  • Tests
    • Extensive new tests covering resilience, healthchecks, rerankers, vLLM, semaphores, and retry/circuit logic.

Review Change Stack

@coderabbitai

coderabbitai Bot commented May 7, 2026

Copy link
Copy Markdown

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

Adds real retry and circuit-breaker implementations, a Ray distributed semaphore, healthcheck probes, vLLM and reranker HTTP clients, DI registration and shims, embedding shims, package exports, tests, and aiobreaker dependency.

Changes

Inference resilience and clients

Layer / File(s) Summary
Decision log / API contract
REFACTORING_DECISION_LOG.md
Adds Phase 6B notes updating LLM generate/chat contract to return dict and documents stream behavior.
vLLM clients
openrag/services/inference/vllm_client.py
Adds VLLMClient, VLLMEmbedder, VLLMVision with httpx clients, streaming, embedding parsing/caching, captioning, and registration.
Reranker clients
openrag/services/inference/reranker_clients.py
Adds InfinityReranker and OpenAIReranker using httpx, response validation, error mapping, and wrappers.
Retry policy
openrag/services/inference/_retry.py
Implements _is_retryable, _log_retry, and with_retry(...) returning a tenacity-configured retry decorator with exponential-jitter backoff and logging.
Circuit breaker
openrag/services/inference/_circuit_breaker.py
Adds cached aiobreaker CircuitBreaker factory get_breaker, client-error exclusion, logging listener, Prometheus gauge, and with_circuit_breaker(...) decorator factory.
Distributed semaphore
openrag/services/inference/distributed_semaphore.py
Adds Ray detached DistributedSemaphoreActor and DistributedSemaphore async context manager with get-or-create actor and remote acquire/release.
Components utils wiring
openrag/components/utils.py
Removes local Ray semaphore implementation and imports DistributedSemaphore/actor from the new module; semaphores remain eagerly instantiated.
Health & external contracts
openrag/services/inference/healthcheck.py, openrag/utils/external_resource_errors.py
Adds EndpointStatus, HealthResult, health probes, and re-exports external-error helpers from core.utils.external_errors.
Embedding shim & indexer
openrag/components/indexer/embeddings/openai.py, __init__.py
Adds _ShimOpenAIEmbedding delegating to VLLMEmbedder, sync helper threadpool, and updates OpenAIEmbedding typing/behavior.
LLM/reranker shims & pipeline
openrag/components/llm.py, openrag/components/reranker/*, openrag/components/pipeline.py
Adds _LLMShim delegating to VLLMClient, _RerankerShim adapter, registry-based reranker factory wiring, and compatibility aliases.
DI registration & container
openrag/di/*.py, openrag/di/container.py
Adds registration entrypoints for embedders/LLMs/rerankers/VLMs and ServiceContainer factory helpers.
Package exports & deps
openrag/services/inference/__init__.py, pyproject.toml
Re-exports inference primitives/clients and adds aiobreaker>=1.2.0 dependency.
Adapters & stubs
openrag/components/indexer/chunker/chunker.py, openrag/components/retriever.py
Adds stream_chat stubs to LangChain adapters and retriever adapter (no-op).
Tests
openrag/services/inference/test_*.py, openrag/core/utils/test_external_errors.py, openrag/di/test_inference.py
Adds comprehensive tests for retry, circuit-breaker, healthchecks, distributed semaphore, rerankers, vLLM clients/embedder/vision, DI registration, and external-error detection.

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant CircuitBreaker
  participant Retry
  participant HTTPClient
  participant Endpoint
  Caller->>CircuitBreaker: call_async(func)
  CircuitBreaker->>Retry: invoke wrapped call
  Retry->>HTTPClient: httpx request (POST/GET/stream)
  HTTPClient->>Endpoint: network request
  Endpoint-->>HTTPClient: response or error
  HTTPClient-->>Retry: response / raise
  Retry-->>CircuitBreaker: result or exception
  CircuitBreaker-->>Caller: result or CircuitBreakerError (mapped to InferenceConnectionError)
Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

Suggested reviewers

  • paultranvan

"I nibble code, I hop and cheer,
Breaker, retry—resilience near.
Actors guard the cluster's gate,
New clients dance, the tests all wait.
A rabbit twitches: deploy, not fear! 🐇"

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 6.78% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title 'refactor(phase 6): inference adapters' directly and concisely summarizes the main objective of this changeset, which is a Phase 6 refactoring focused on adding new inference adapter clients (vLLM, reranker) and supporting infrastructure (circuit breaker, retry, health checks, distributed semaphore).
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch refactor/phase-6-inference-adapters

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

  • Generate code and open pull requests
  • Plan features and break down work
  • Investigate incidents and troubleshoot customer tickets together
  • Automate recurring tasks and respond to alerts with triggers
  • Summarize progress and report instantly

Built for teams:

  • Shared memory across your entire org—no repeating context
  • Per-thread sandboxes to safely plan and execute work
  • Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get started


Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
openrag/services/inference/_retry.py (1)

1-1: 💤 Low value

Consider adding type hints for consistency.

The companion with_circuit_breaker decorator includes type hints on its parameters, but this function does not. Adding hints would make the API contract more explicit and keep the inference layer's resilience utilities consistent.

📝 Optional type hint addition
-def with_retry(max_attempts=3, base_wait=1.0):
+def with_retry(max_attempts: int = 3, base_wait: float = 1.0):
     """No-op stub — real tenacity implementation replaces this later."""
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@openrag/services/inference/_retry.py` at line 1, The with_retry decorator
function lacks type hints; add explicit typing to its signature (e.g., annotate
max_attempts: int = 3, base_wait: float = 1.0 and the return type as
Callable[..., Callable[..., Any]] or similar) so it matches the style of
with_circuit_breaker and makes the API contract explicit; update any imports
(typing.Callable, typing.Any) if needed and ensure the annotated signature is
applied to the with_retry function and its inner wrapper types.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@openrag/services/inference/_retry.py`:
- Line 1: The with_retry decorator function lacks type hints; add explicit
typing to its signature (e.g., annotate max_attempts: int = 3, base_wait: float
= 1.0 and the return type as Callable[..., Callable[..., Any]] or similar) so it
matches the style of with_circuit_breaker and makes the API contract explicit;
update any imports (typing.Callable, typing.Any) if needed and ensure the
annotated signature is applied to the with_retry function and its inner wrapper
types.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 3cbc3090-fb78-47e8-b904-7910df165569

📥 Commits

Reviewing files that changed from the base of the PR and between 2ed6c89 and 6c39173.

📒 Files selected for processing (2)
  • openrag/services/inference/_circuit_breaker.py
  • openrag/services/inference/_retry.py

@EnjoyBacon7 EnjoyBacon7 changed the title chore(A): fast unlock for downstream tasks 6B, 6C chore(A): fast unlock for downstream tasks 6B, 6C, 6D May 7, 2026
@EnjoyBacon7
EnjoyBacon7 force-pushed the refactor/phase-6-inference-adapters branch from 340f774 to 0013122 Compare May 7, 2026 14:49
@coderabbitai coderabbitai Bot added the feat Add a new feature label May 7, 2026

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 7

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@openrag/services/inference/_circuit_breaker.py`:
- Around line 32-41: get_breaker currently caches CircuitBreaker instances in
_breakers solely by name so the first call's fail_max/timeout_duration silently
win; update get_breaker to either (a) include configuration in the cache key
(e.g., name + fail_max + timeout_duration) or (b) validate the existing
breaker’s configuration against the requested fail_max and timeout_duration and
either raise a clear error or recreate/replace the breaker in _breakers when
they differ. Modify the logic around get_breaker, _breakers and the created
CircuitBreaker (the instance created with fail_max, timeout_duration, name,
exclude=[_is_client_error], listeners=[_LoggingListener()]) to perform this
check or use the expanded key so callers don’t get unexpected silent mismatches.

In `@openrag/services/inference/reranker_clients.py`:
- Around line 58-64: The reranker clients (InfinityReranker and OpenAIReranker)
currently let httpx.HTTPStatusError escape and parse resp.json()["results"]
outside the try block; update the try/except around resp.raise_for_status() to
also catch httpx.HTTPStatusError and rethrow a domain-mapped error (e.g., raise
a new InferenceHTTPStatusError with self._endpoint info) alongside the existing
InferenceConnectionError and InferenceTimeoutError, and move the response JSON
parsing into that same try block so that JSON parsing errors
(ValueError/KeyError) are caught and rewrapped as a domain error (e.g.,
InferenceResponseError) before returning the list comprehension; reference
resp.raise_for_status(), InferenceConnectionError, InferenceTimeoutError,
InfinityReranker, and OpenAIReranker to locate the changes.

In `@openrag/services/inference/test_circuit_breaker.py`:
- Around line 11-16: The _clean_breakers fixture only cleans up after tests;
modify it to also perform the same cleanup before yielding so each test starts
with a reset state: iterate the module-level _breakers dict, call
breaker.close() for each breaker, then clear _breakers prior to the yield
(keeping the existing post-yield cleanup as well) so the fixture ensures
isolation for the first and subsequent tests.

In `@openrag/services/inference/vllm_client.py`:
- Around line 79-88: The chat method builds payload with "stream": False but
places it before **kwargs so callers can override it; update the payload
construction in chat (the dict built with self._model, self._defaults, and
kwargs) so that "stream": False is applied after merging self._defaults and
**kwargs (i.e., pin stream by adding "stream": False last) to ensure the method
always requests non-streaming responses from self._endpoint via
self._client.post and avoid treating SSE as JSON.
- Around line 258-259: caption_images_batch currently uses asyncio.gather(...)
which lets sibling tasks continue running if one raises; change
caption_images_batch to explicitly create tasks for each caption_image (via
asyncio.create_task), use asyncio.wait(tasks,
return_when=asyncio.FIRST_EXCEPTION) to detect the first failure, cancel any
still-pending tasks, then await/gather results to either raise the first
exception or return the list of captions; ensure you cancel pending tasks and
propagate the original exception from caption_image so in-flight HTTP requests
are not orphaned.
- Around line 148-155: The request body sent in vllm_client.py via
self._client.post to f"{self._endpoint}/embeddings" currently nests
truncate_prompt_tokens under "extra_body", which vLLM ignores; change the JSON
payload so "truncate_prompt_tokens": self._max_model_len is a top-level key
alongside "model" and "input" (keep using self._model and texts), removing the
"extra_body" wrapper, and update the assertion in test_vllm_client.py to expect
the top-level field instead of the nested structure.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: bd136a52-6071-4de1-995e-78b265c28995

📥 Commits

Reviewing files that changed from the base of the PR and between 6c39173 and de00cc6.

⛔ Files ignored due to path filters (1)
  • uv.lock is excluded by !**/*.lock
📒 Files selected for processing (16)
  • openrag/components/utils.py
  • openrag/core/utils/test_external_errors.py
  • openrag/services/inference/_circuit_breaker.py
  • openrag/services/inference/_retry.py
  • openrag/services/inference/distributed_semaphore.py
  • openrag/services/inference/healthcheck.py
  • openrag/services/inference/reranker_clients.py
  • openrag/services/inference/test_circuit_breaker.py
  • openrag/services/inference/test_healthcheck.py
  • openrag/services/inference/test_reranker_clients.py
  • openrag/services/inference/test_retry.py
  • openrag/services/inference/test_vllm_client.py
  • openrag/services/inference/vllm_client.py
  • openrag/utils/dependencies.py
  • openrag/utils/external_resource_errors.py
  • pyproject.toml
✅ Files skipped from review due to trivial changes (1)
  • pyproject.toml

Comment thread openrag/services/inference/_circuit_breaker.py Outdated
Comment thread openrag/services/inference/reranker_clients.py Outdated
Comment thread openrag/services/inference/test_circuit_breaker.py
Comment thread openrag/services/inference/vllm_client.py Outdated
Comment thread openrag/services/inference/vllm_client.py Outdated
Comment thread openrag/services/inference/vllm_client.py Outdated
Comment thread openrag/services/inference/vllm_client.py
Ahmath-Gadji and others added 2 commits May 7, 2026 17:16
…nce (phase 6D)

Move DistributedSemaphoreActor and DistributedSemaphore from
components/utils.py to services/inference/distributed_semaphore.py.
Remove module-level config read from @ray.remote decorator.
Old locations re-export from the new canonical home.
@EnjoyBacon7
EnjoyBacon7 force-pushed the refactor/phase-6-inference-adapters branch 3 times, most recently from 4d6f204 to 44d36b7 Compare May 7, 2026 17:26
Ahmath-Gadji and others added 5 commits May 7, 2026 17:29
- Replace no-op stubs with real tenacity-based retry decorator and
  aiobreaker-based circuit breaker in services/inference/
- _retry.py: exponential backoff + jitter, retries on transient httpx
  and OpenRAGError failures (429/502/503/504)
- _circuit_breaker.py: singleton registry, excludes 4xx client errors,
  wraps CircuitBreakerError into InferenceConnectionError, logs state
  transitions
- Re-export utils/external_resource_errors.py from core/utils/
- Add unit tests for both modules and core external_errors
@EnjoyBacon7
EnjoyBacon7 force-pushed the refactor/phase-6-inference-adapters branch from 44d36b7 to b894ca8 Compare May 7, 2026 17:29

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 6

♻️ Duplicate comments (1)
openrag/services/inference/_circuit_breaker.py (1)

48-57: ⚠️ Potential issue | 🟠 Major | ⚖️ Poor tradeoff

Cached breaker ignores later config for the same name.

get_breaker() caches CircuitBreaker instances solely by name, so the first call's fail_max and timeout_duration silently win. Any later caller using the same name with different config values won't get the requested thresholds.

Suggested approach

Consider either:

  1. Include configuration in the cache key (e.g., (name, fail_max, timeout_duration))
  2. Validate the existing breaker's configuration and raise an error if it doesn't match the requested values
 _breakers: dict[str, CircuitBreaker] = {}
+_breaker_config: dict[str, tuple[int, float]] = {}

 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]}, "
+            f"requested={requested}"
+        )
     return _breakers[name]
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@openrag/services/inference/_circuit_breaker.py` around lines 48 - 57,
get_breaker currently caches CircuitBreaker instances only by name so the first
call's fail_max/timeout_duration silently win; change the caching strategy in
get_breaker to either include the configuration in the cache key (e.g., use key
= (name, fail_max, timedelta(seconds=timeout_duration))) so distinct configs
produce distinct breakers, or keep caching by name but validate the existing
breaker's settings (compare fail_max and timeout_duration on the cached
CircuitBreaker) and raise a clear error if they differ; update references to
_breakers and the constructor call in get_breaker accordingly and ensure
timeout_duration comparison accounts for the timedelta conversion used when
constructing CircuitBreaker.
🧹 Nitpick comments (3)
openrag/services/inference/test_distributed_semaphore.py (3)

5-15: ⚡ Quick win

Avoid testing private attributes directly.

Both test_default_params and test_custom_params access private attributes (_name, _namespace, _max_concurrent_ops), which couples the tests to implementation details and makes them brittle to internal refactoring.

Consider one of these alternatives:

  • Expose these as public properties if external callers need them
  • Test the semaphore's behavior rather than its internal state (e.g., verify it acquires/releases correctly with the configured parameters)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@openrag/services/inference/test_distributed_semaphore.py` around lines 5 -
15, Tests are reading private attributes (_name, _namespace,
_max_concurrent_ops) on DistributedSemaphore; change the tests to avoid coupling
to internals by either (A) using public accessors if you add them (e.g., expose
name, namespace, max_concurrent_ops properties on DistributedSemaphore and
update test_default_params and test_custom_params to assert those), or (B) test
behavior instead (create a DistributedSemaphore with given params in
test_default_params/test_custom_params and assert expected behavior such as
acquiring/releasing limits: start N concurrent acquire attempts and verify only
max_concurrent_ops succeed concurrently, and verify namespace/name can be
validated via a public identifier or via the semaphore's public API). Ensure
references to DistributedSemaphore, test_default_params, and test_custom_params
are updated accordingly.

4-18: Consider adding behavioral and integration tests.

The current tests verify constructor parameters and class structure, but don't test the distributed semaphore's actual behavior. Consider adding:

  • Async tests using pytest.mark.asyncio to verify the async context manager protocol
  • Integration tests that verify concurrent acquisition limits (e.g., spawning N tasks, verifying only max_concurrent_ops run simultaneously)
  • Error handling tests (e.g., actor unavailable, timeout scenarios)
  • Ray actor lifecycle tests (initialization, cleanup)

These tests would provide confidence that the distributed coordination works correctly in production scenarios.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@openrag/services/inference/test_distributed_semaphore.py` around lines 4 -
18, Add behavioral and integration tests for the DistributedSemaphore
distributed coordination: write async pytest tests (use pytest.mark.asyncio)
that exercise DistributedSemaphore as an async context manager (enter/exit),
spawn multiple concurrent tasks that call DistributedSemaphore.acquire()/use the
async context to verify only _max_concurrent_ops tasks run simultaneously, and
include timeout/error scenarios (simulate actor unavailability and ensure proper
exception handling/retry) as well as actor lifecycle checks around
DistributedSemaphoreActor initialization and cleanup; reference the
DistributedSemaphore class and DistributedSemaphoreActor to locate where to call
acquire/release and to create/teardown actors for integration tests.

17-18: ⚡ Quick win

Consider testing actor behavior, not just existence.

This test only verifies that DistributedSemaphoreActor has a remote attribute, which confirms it's a Ray remote class but doesn't test any actual functionality. Consider adding a test that instantiates the actor and verifies basic operations (e.g., can acquire/release the underlying semaphore).

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@openrag/services/inference/test_distributed_semaphore.py` around lines 17 -
18, The test currently only checks for the presence of
DistributedSemaphoreActor.remote; replace or extend it to exercise actor
behavior by instantiating the actor via DistributedSemaphoreActor.remote(...)
(after ray.init in the test setup), calling its semaphore methods (e.g., acquire
and release or any public methods like acquire_blocking/release_and_get_count)
using ray.get on returned futures, and asserting the expected state changes (for
example that acquire reduces available permits and release restores them).
Ensure you clean up with ray.shutdown and use short timeouts or simple
concurrency (multiple actor clients) to validate correct semaphore behavior
rather than only checking attribute existence.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@openrag/components/indexer/embeddings/openai.py`:
- Around line 46-48: The embedding_dimension property currently proxies
VLLMEmbedder.dimension but VLLMEmbedder only sets its _dimension after the first
embed; change embedding_dimension to first check whether the delegate actually
has a populated dimension (e.g., hasattr(self._delegate, "_dimension") and
self._delegate._dimension is not None or guard against attribute access raising)
and return that if present, otherwise fall back to the legacy OpenAIEmbedding
behavior (call the original/legacy embedding_dimension or return the stored
legacy default) so callers can read embedding_dimension during setup before any
embed has occurred.
- Around line 50-58: embed_documents and aembed_documents assume texts[0] exists
and only check the first item, causing IndexError on empty lists and allowing
mixed batches (e.g., [str, Document]) to pass Document objects into the HTTP
embedder; fix by guarding empty batches and normalizing every element: in
embed_documents and aembed_documents, first return early for empty texts (return
[]), then map/replace any Document instances with doc.page_content for all items
(not just checking texts[0]) before calling _delegate.embed (and wrap with
_run_sync in embed_documents); update references to Document, _delegate.embed,
_run_sync, embed_documents and aembed_documents when applying the change.

In `@openrag/components/reranker/__init__.py`:
- Around line 23-27: The loop that applies reranker output assumes indices from
ranked map directly into documents (for index, score in ranked ->
documents[index]) which can raise IndexError on bad delegate output; update the
logic in the reranker module to validate each index before accessing documents:
check 0 <= index < len(documents), skip (and optionally process a fallback or
log a warning via the module logger) any out-of-range indices so the rerank
operation degrades gracefully, and ensure output only receives valid Document
objects (preserve metadata injection of "relevance_score" for valid entries).

In `@openrag/core/llm/llm.py`:
- Around line 22-25: stream_chat currently yields the raw dict returned by
chat(), violating its AsyncIterator[str] / SSE contract; update stream_chat to
await self.chat(...), serialize the result to a string (e.g. json.dumps(result))
and emit it as SSE-formatted string(s) (for example prepend "data: " and
terminate with double newline) so consumers receive valid string SSE chunks;
ensure you import json and keep the method signature, or alternatively make
stream_chat abstract if you prefer providers to implement streaming themselves.

In `@openrag/services/inference/vllm_client.py`:
- Around line 223-226: The try block that calls resp.json() and extracts data
can raise JSON decoding errors (ValueError/JSONDecodeError) which are not
currently mapped to EmbeddingResponseError; update the except clause in
openrag.services.inference.vllm_client (the block around resp.json()["data"] and
embeddings = [...]) to also catch ValueError (or json.JSONDecodeError) in
addition to KeyError/IndexError/TypeError and re-raise or convert it to
EmbeddingResponseError so malformed/non-JSON 200 responses are normalized by the
adapter.
- Around line 87-103: The code mutates the caller's metadata by doing
metadata.pop("llm_override", ...) and kwargs.pop("metadata"), which strips the
override across retries; instead read llm_override without mutating the original
dict (e.g., use metadata.get("llm_override") or work on a shallow copy like
metadata_copy = dict(metadata)) and avoid popping metadata out of kwargs; update
the logic in vllm_client.py (the block that references metadata, llm_override,
base_url, model, override_headers) to read values from the copy or via .get(),
and only modify or remove kwargs["metadata"] if you first replace it with a copy
so caller-owned objects are never mutated.

---

Duplicate comments:
In `@openrag/services/inference/_circuit_breaker.py`:
- Around line 48-57: get_breaker currently caches CircuitBreaker instances only
by name so the first call's fail_max/timeout_duration silently win; change the
caching strategy in get_breaker to either include the configuration in the cache
key (e.g., use key = (name, fail_max, timedelta(seconds=timeout_duration))) so
distinct configs produce distinct breakers, or keep caching by name but validate
the existing breaker's settings (compare fail_max and timeout_duration on the
cached CircuitBreaker) and raise a clear error if they differ; update references
to _breakers and the constructor call in get_breaker accordingly and ensure
timeout_duration comparison accounts for the timedelta conversion used when
constructing CircuitBreaker.

---

Nitpick comments:
In `@openrag/services/inference/test_distributed_semaphore.py`:
- Around line 5-15: Tests are reading private attributes (_name, _namespace,
_max_concurrent_ops) on DistributedSemaphore; change the tests to avoid coupling
to internals by either (A) using public accessors if you add them (e.g., expose
name, namespace, max_concurrent_ops properties on DistributedSemaphore and
update test_default_params and test_custom_params to assert those), or (B) test
behavior instead (create a DistributedSemaphore with given params in
test_default_params/test_custom_params and assert expected behavior such as
acquiring/releasing limits: start N concurrent acquire attempts and verify only
max_concurrent_ops succeed concurrently, and verify namespace/name can be
validated via a public identifier or via the semaphore's public API). Ensure
references to DistributedSemaphore, test_default_params, and test_custom_params
are updated accordingly.
- Around line 4-18: Add behavioral and integration tests for the
DistributedSemaphore distributed coordination: write async pytest tests (use
pytest.mark.asyncio) that exercise DistributedSemaphore as an async context
manager (enter/exit), spawn multiple concurrent tasks that call
DistributedSemaphore.acquire()/use the async context to verify only
_max_concurrent_ops tasks run simultaneously, and include timeout/error
scenarios (simulate actor unavailability and ensure proper exception
handling/retry) as well as actor lifecycle checks around
DistributedSemaphoreActor initialization and cleanup; reference the
DistributedSemaphore class and DistributedSemaphoreActor to locate where to call
acquire/release and to create/teardown actors for integration tests.
- Around line 17-18: The test currently only checks for the presence of
DistributedSemaphoreActor.remote; replace or extend it to exercise actor
behavior by instantiating the actor via DistributedSemaphoreActor.remote(...)
(after ray.init in the test setup), calling its semaphore methods (e.g., acquire
and release or any public methods like acquire_blocking/release_and_get_count)
using ray.get on returned futures, and asserting the expected state changes (for
example that acquire reduces available permits and release restores them).
Ensure you clean up with ray.shutdown and use short timeouts or simple
concurrency (multiple actor clients) to validate correct semaphore behavior
rather than only checking attribute existence.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: ee6b1029-8333-4d38-80c9-e907ddd8df18

📥 Commits

Reviewing files that changed from the base of the PR and between de00cc6 and bd2014f.

⛔ Files ignored due to path filters (1)
  • uv.lock is excluded by !**/*.lock
📒 Files selected for processing (35)
  • REFACTORING_DECISION_LOG.md
  • openrag/components/indexer/embeddings/__init__.py
  • openrag/components/indexer/embeddings/openai.py
  • openrag/components/llm.py
  • openrag/components/pipeline.py
  • openrag/components/reranker/__init__.py
  • openrag/components/reranker/base.py
  • openrag/components/reranker/infinity.py
  • openrag/components/reranker/openai.py
  • openrag/components/utils.py
  • openrag/core/llm/llm.py
  • openrag/core/utils/test_external_errors.py
  • openrag/di/container.py
  • openrag/di/embedders.py
  • openrag/di/inference.py
  • openrag/di/llms.py
  • openrag/di/rerankers.py
  • openrag/di/test_inference.py
  • openrag/di/vlms.py
  • openrag/services/inference/__init__.py
  • openrag/services/inference/_circuit_breaker.py
  • openrag/services/inference/_retry.py
  • openrag/services/inference/distributed_semaphore.py
  • openrag/services/inference/healthcheck.py
  • openrag/services/inference/reranker_clients.py
  • openrag/services/inference/test_circuit_breaker.py
  • openrag/services/inference/test_distributed_semaphore.py
  • openrag/services/inference/test_healthcheck.py
  • openrag/services/inference/test_reranker_clients.py
  • openrag/services/inference/test_retry.py
  • openrag/services/inference/test_vllm_client.py
  • openrag/services/inference/vllm_client.py
  • openrag/utils/dependencies.py
  • openrag/utils/external_resource_errors.py
  • pyproject.toml
✅ Files skipped from review due to trivial changes (6)
  • openrag/di/inference.py
  • openrag/di/llms.py
  • REFACTORING_DECISION_LOG.md
  • openrag/services/inference/init.py
  • openrag/core/utils/test_external_errors.py
  • openrag/services/inference/test_reranker_clients.py
🚧 Files skipped from review as they are similar to previous changes (8)
  • openrag/utils/external_resource_errors.py
  • openrag/services/inference/_retry.py
  • openrag/utils/dependencies.py
  • pyproject.toml
  • openrag/services/inference/reranker_clients.py
  • openrag/services/inference/distributed_semaphore.py
  • openrag/components/utils.py
  • openrag/services/inference/test_vllm_client.py

Comment thread openrag/components/indexer/embeddings/openai.py Outdated
Comment thread openrag/components/indexer/embeddings/openai.py Outdated
Comment thread openrag/components/reranker/__init__.py
Comment thread openrag/core/llm/llm.py Outdated
Comment thread openrag/services/inference/vllm_client.py Outdated
Comment thread openrag/services/inference/vllm_client.py Outdated
@EnjoyBacon7
EnjoyBacon7 force-pushed the refactor/phase-6-inference-adapters branch from aa7f134 to 8f0af3e Compare May 11, 2026 14:16

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (1)
openrag/core/llm/llm.py (1)

22-30: ⚖️ Poor tradeoff

Consider declaring stream_chat() as async def for signature consistency.

The abstract method is currently declared as def returning AsyncIterator[str], but the docstring and all implementations use async def. While this works at runtime (Python allows overriding def with async def), the inconsistency can confuse developers.

Since the project uses ruff for linting (not mypy/pyright for type checking), there are no type-checker warnings today. However, for clarity and consistency, consider using:

`@abstractmethod`
async def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> AsyncIterator[str]:
    """Stream chat completion as raw SSE lines."""
    if False:
        yield ""

This makes the abstract signature match implementations and documents the async generator contract explicitly.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@openrag/core/llm/llm.py` around lines 22 - 30, The abstract method
stream_chat is declared as a regular def but implemented everywhere as async
generators; change its signature to "async def stream_chat(self, messages:
list[dict[str, str]], **kwargs) -> AsyncIterator[str]:" (keep the
`@abstractmethod` decorator) and add a no-op guarded yield (e.g., "if False: yield
\"\"") inside the body so the signature explicitly declares an async generator
and matches the implementations.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@openrag/core/retrieval/rrf.py`:
- Line 54: The function declared to return list[T] currently returns
ranked_lists[0] (a Sequence[T]) when len(ranked_lists) == 1, violating the type
signature and causing inconsistent behavior; update the single-list return to
produce a concrete list (e.g., return a defensive copy like
list(ranked_lists[0])) so the function always returns list[T] and matches the
multi-list path; locate the check that returns ranked_lists[0] in the function
(referenced by ranked_lists and the function signature on line 30) and replace
that direct Sequence return with a list conversion/copy.

---

Nitpick comments:
In `@openrag/core/llm/llm.py`:
- Around line 22-30: The abstract method stream_chat is declared as a regular
def but implemented everywhere as async generators; change its signature to
"async def stream_chat(self, messages: list[dict[str, str]], **kwargs) ->
AsyncIterator[str]:" (keep the `@abstractmethod` decorator) and add a no-op
guarded yield (e.g., "if False: yield \"\"") inside the body so the signature
explicitly declares an async generator and matches the implementations.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: d595fa84-05b2-4ccf-87fd-d98b2e7af359

📥 Commits

Reviewing files that changed from the base of the PR and between bd2014f and 8f0af3e.

📒 Files selected for processing (12)
  • openrag/components/indexer/embeddings/__init__.py
  • openrag/components/indexer/embeddings/openai.py
  • openrag/components/llm.py
  • openrag/components/pipeline.py
  • openrag/components/reranker/__init__.py
  • openrag/components/reranker/base.py
  • openrag/core/llm/llm.py
  • openrag/core/retrieval/rrf.py
  • openrag/services/inference/_circuit_breaker.py
  • openrag/services/inference/test_circuit_breaker.py
  • openrag/services/inference/test_vllm_client.py
  • openrag/services/inference/vllm_client.py
🚧 Files skipped from review as they are similar to previous changes (10)
  • openrag/components/pipeline.py
  • openrag/components/indexer/embeddings/init.py
  • openrag/components/reranker/init.py
  • openrag/components/reranker/base.py
  • openrag/components/llm.py
  • openrag/components/indexer/embeddings/openai.py
  • openrag/services/inference/_circuit_breaker.py
  • openrag/services/inference/test_vllm_client.py
  • openrag/services/inference/test_circuit_breaker.py
  • openrag/services/inference/vllm_client.py

Comment thread openrag/core/retrieval/rrf.py Outdated
@Ahmath-Gadji Ahmath-Gadji changed the title chore(A): fast unlock for downstream tasks 6B, 6C, 6D refactor(phase 6): inference adapters May 11, 2026
…s dict

- Extract _parse_response / _parse_rerank_response helpers; unify HTTP
  error and JSON-decode handling in one place per client family
- Add _resolve_overrides to VLLMClient for per-request model/endpoint
  switching via metadata.llm_override
- VLLMVision(VLLMClient, VLM): reuse LLM connection pool and decorators,
  keep VLM in bases for nominal typing
- Separate circuit breakers per reranker backend (reranker_infinity /
  reranker_openai) so one outage doesn't trip the other
- Type LLM.generate() → dict, LLM.chat() → dict to preserve the full
  OpenAI response body; stream_chat stays AsyncIterator[str] (raw SSE)
- Add _ShimOpenAIEmbedding bridging the sync LangChain embed contract
  to the async VLLMEmbedder for legacy vectordb.py callers
- Fix VLLMEmbedder.dimension: raise RuntimeError when called before the
  first embed() instead of silently returning None
- Log Phase 6B decisions (VLLMVision MRO, LLM return type rationale)
@EnjoyBacon7
EnjoyBacon7 force-pushed the refactor/phase-6-inference-adapters branch from ba43af4 to 885a0cd Compare May 11, 2026 15:28

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (2)
openrag/components/indexer/chunker/chunker.py (1)

55-56: ⚡ Quick win

Use explicit unsupported behavior for stream_chat.

Line 56 silently returns None. Please raise NotImplementedError (or intentionally delegate to chat) so accidental calls fail predictably.

Proposed patch
 async def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> Any:
-    pass  # Not implemented since the contextualizer never streams.
+    raise NotImplementedError(
+        "_LangChainLLMAdapter does not support stream_chat in chunker contextualization."
+    )
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@openrag/components/indexer/chunker/chunker.py` around lines 55 - 56, The
stream_chat method currently returns None silently; change it to explicitly
signal unsupported behavior by either raising NotImplementedError("stream_chat
is not supported by the contextualizer") or delegating to the existing chat
method (e.g., return await self.chat(messages, **kwargs)) so accidental callers
fail predictably; update the async def stream_chat(...) implementation in the
chunker.Chunker class to perform one of these two options and ensure the error
message clearly references stream_chat.
openrag/components/retriever.py (1)

69-70: ⚡ Quick win

Raise NotImplementedError to fail fast and comply with the base LLM contract.

Line 70 uses pass, which returns None instead of the expected AsyncIterator[str]. Although the comment states the retriever won't call this method, using pass violates the abstract base class contract and risks a silent TypeError at runtime if streaming is ever attempted. Raise an explicit NotImplementedError instead.

Proposed patch
 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.
+    raise NotImplementedError(
+        "_LangChainLLMAdapter does not support stream_chat; use chat() for this legacy path."
+    )
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@openrag/components/retriever.py` around lines 69 - 70, Replace the no-op
implementation of stream_chat with a fail-fast NotImplementedError: in
openrag/components/retriever.py locate async def stream_chat(self, messages:
list[dict[str, str]], **kwargs) -> Any and change the body to raise
NotImplementedError("stream_chat not implemented for Retriever; use
non-streaming methods") so it adheres to the LLM/async contract (expected
AsyncIterator[str]) and fails explicitly if called.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@openrag/components/indexer/chunker/chunker.py`:
- Around line 55-56: The stream_chat method currently returns None silently;
change it to explicitly signal unsupported behavior by either raising
NotImplementedError("stream_chat is not supported by the contextualizer") or
delegating to the existing chat method (e.g., return await self.chat(messages,
**kwargs)) so accidental callers fail predictably; update the async def
stream_chat(...) implementation in the chunker.Chunker class to perform one of
these two options and ensure the error message clearly references stream_chat.

In `@openrag/components/retriever.py`:
- Around line 69-70: Replace the no-op implementation of stream_chat with a
fail-fast NotImplementedError: in openrag/components/retriever.py locate async
def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> Any and
change the body to raise NotImplementedError("stream_chat not implemented for
Retriever; use non-streaming methods") so it adheres to the LLM/async contract
(expected AsyncIterator[str]) and fails explicitly if called.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 9258596c-9c07-4392-a35a-044a7cbf8d17

📥 Commits

Reviewing files that changed from the base of the PR and between 8f0af3e and 885a0cd.

⛔ Files ignored due to path filters (1)
  • uv.lock is excluded by !**/*.lock
📒 Files selected for processing (19)
  • REFACTORING_DECISION_LOG.md
  • openrag/components/indexer/chunker/chunker.py
  • openrag/components/indexer/embeddings/__init__.py
  • openrag/components/indexer/embeddings/openai.py
  • openrag/components/llm.py
  • openrag/components/pipeline.py
  • openrag/components/reranker/__init__.py
  • openrag/components/reranker/base.py
  • openrag/components/reranker/infinity.py
  • openrag/components/reranker/openai.py
  • openrag/components/reranker/test_rrf_reranking.py
  • openrag/components/retriever.py
  • openrag/core/llm/llm.py
  • openrag/core/retrieval/rrf.py
  • openrag/services/inference/__init__.py
  • openrag/services/inference/_circuit_breaker.py
  • openrag/services/inference/test_circuit_breaker.py
  • openrag/services/inference/test_vllm_client.py
  • openrag/services/inference/vllm_client.py
✅ Files skipped from review due to trivial changes (3)
  • openrag/services/inference/init.py
  • REFACTORING_DECISION_LOG.md
  • openrag/core/retrieval/rrf.py
🚧 Files skipped from review as they are similar to previous changes (13)
  • openrag/components/indexer/embeddings/init.py
  • openrag/components/pipeline.py
  • openrag/services/inference/test_circuit_breaker.py
  • openrag/components/llm.py
  • openrag/components/reranker/openai.py
  • openrag/components/reranker/init.py
  • openrag/components/reranker/base.py
  • openrag/components/reranker/infinity.py
  • openrag/services/inference/_circuit_breaker.py
  • openrag/services/inference/vllm_client.py
  • openrag/core/llm/llm.py
  • openrag/components/indexer/embeddings/openai.py
  • openrag/services/inference/test_vllm_client.py

… circuit breaker

rrf_reranking was wrapping the single input list in list(), breaking the
identity contract. _is_excluded now splits into _is_client_error + an
explicit LLMParsingError check since its 502 status code fell outside
the 4xx exclusion range.
- circuit_breaker: reject reuse with mismatched config (loud at startup)
- circuit_breaker tests: pre-yield cleanup so first test is isolated
- vllm_client: stream_chat maps httpx ConnectError/TimeoutException to
  domain errors; previously raw httpx escaped past the boundary
- vllm_client: _resolve_overrides is now a pure read so retries see the
  original llm_override; outbound payload strips internal metadata
- vllm_client: catch JSON ValueError in embed() so malformed 200 bodies
  surface as EmbeddingResponseError
- embeddings/openai shim: normalize mixed [str, Document] batches and
  guard empty batches before indexing
- reranker shim: skip out-of-range delegate indices instead of raising
- core/llm: stream_chat is @AbstractMethod (default impl violated
  AsyncIterator[str] contract); legacy _LangChainLLMAdapter in
  chunker.py and retriever.py implement the new abstract method so
  ChunkContextualizer / MultiQuery / HyDe stop raising TypeError on
  instantiation
- uv.lock: relock for tenacity dep added in phase 6A
@EnjoyBacon7
EnjoyBacon7 force-pushed the refactor/phase-6-inference-adapters branch from 885a0cd to 28a7100 Compare May 11, 2026 15:38
@Ahmath-Gadji
Ahmath-Gadji merged commit ef3ccc2 into refactor/hexagonal May 11, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feat Add a new feature

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants