refactor(phase 6): inference adapters - #375
Conversation
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds 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. ChangesInference resilience and clients
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)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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.
Built for teams:
One agent for your entire SDLC. Right inside Slack. 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. Comment |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
openrag/services/inference/_retry.py (1)
1-1: 💤 Low valueConsider adding type hints for consistency.
The companion
with_circuit_breakerdecorator 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
📒 Files selected for processing (2)
openrag/services/inference/_circuit_breaker.pyopenrag/services/inference/_retry.py
340f774 to
0013122
Compare
There was a problem hiding this comment.
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
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (16)
openrag/components/utils.pyopenrag/core/utils/test_external_errors.pyopenrag/services/inference/_circuit_breaker.pyopenrag/services/inference/_retry.pyopenrag/services/inference/distributed_semaphore.pyopenrag/services/inference/healthcheck.pyopenrag/services/inference/reranker_clients.pyopenrag/services/inference/test_circuit_breaker.pyopenrag/services/inference/test_healthcheck.pyopenrag/services/inference/test_reranker_clients.pyopenrag/services/inference/test_retry.pyopenrag/services/inference/test_vllm_client.pyopenrag/services/inference/vllm_client.pyopenrag/utils/dependencies.pyopenrag/utils/external_resource_errors.pypyproject.toml
✅ Files skipped from review due to trivial changes (1)
- pyproject.toml
…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.
4d6f204 to
44d36b7
Compare
- 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
44d36b7 to
b894ca8
Compare
There was a problem hiding this comment.
Actionable comments posted: 6
♻️ Duplicate comments (1)
openrag/services/inference/_circuit_breaker.py (1)
48-57:⚠️ Potential issue | 🟠 Major | ⚖️ Poor tradeoffCached breaker ignores later config for the same name.
get_breaker()cachesCircuitBreakerinstances solely byname, so the first call'sfail_maxandtimeout_durationsilently win. Any later caller using the same name with different config values won't get the requested thresholds.Suggested approach
Consider either:
- Include configuration in the cache key (e.g.,
(name, fail_max, timeout_duration))- 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 winAvoid testing private attributes directly.
Both
test_default_paramsandtest_custom_paramsaccess 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.asyncioto verify the async context manager protocol- Integration tests that verify concurrent acquisition limits (e.g., spawning N tasks, verifying only
max_concurrent_opsrun 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 winConsider testing actor behavior, not just existence.
This test only verifies that
DistributedSemaphoreActorhas aremoteattribute, 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
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (35)
REFACTORING_DECISION_LOG.mdopenrag/components/indexer/embeddings/__init__.pyopenrag/components/indexer/embeddings/openai.pyopenrag/components/llm.pyopenrag/components/pipeline.pyopenrag/components/reranker/__init__.pyopenrag/components/reranker/base.pyopenrag/components/reranker/infinity.pyopenrag/components/reranker/openai.pyopenrag/components/utils.pyopenrag/core/llm/llm.pyopenrag/core/utils/test_external_errors.pyopenrag/di/container.pyopenrag/di/embedders.pyopenrag/di/inference.pyopenrag/di/llms.pyopenrag/di/rerankers.pyopenrag/di/test_inference.pyopenrag/di/vlms.pyopenrag/services/inference/__init__.pyopenrag/services/inference/_circuit_breaker.pyopenrag/services/inference/_retry.pyopenrag/services/inference/distributed_semaphore.pyopenrag/services/inference/healthcheck.pyopenrag/services/inference/reranker_clients.pyopenrag/services/inference/test_circuit_breaker.pyopenrag/services/inference/test_distributed_semaphore.pyopenrag/services/inference/test_healthcheck.pyopenrag/services/inference/test_reranker_clients.pyopenrag/services/inference/test_retry.pyopenrag/services/inference/test_vllm_client.pyopenrag/services/inference/vllm_client.pyopenrag/utils/dependencies.pyopenrag/utils/external_resource_errors.pypyproject.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
aa7f134 to
8f0af3e
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
openrag/core/llm/llm.py (1)
22-30: ⚖️ Poor tradeoffConsider declaring
stream_chat()asasync deffor signature consistency.The abstract method is currently declared as
defreturningAsyncIterator[str], but the docstring and all implementations useasync def. While this works at runtime (Python allows overridingdefwithasync 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
📒 Files selected for processing (12)
openrag/components/indexer/embeddings/__init__.pyopenrag/components/indexer/embeddings/openai.pyopenrag/components/llm.pyopenrag/components/pipeline.pyopenrag/components/reranker/__init__.pyopenrag/components/reranker/base.pyopenrag/core/llm/llm.pyopenrag/core/retrieval/rrf.pyopenrag/services/inference/_circuit_breaker.pyopenrag/services/inference/test_circuit_breaker.pyopenrag/services/inference/test_vllm_client.pyopenrag/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
…ference package __init__
…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)
ba43af4 to
885a0cd
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (2)
openrag/components/indexer/chunker/chunker.py (1)
55-56: ⚡ Quick winUse explicit unsupported behavior for
stream_chat.Line 56 silently returns
None. Please raiseNotImplementedError(or intentionally delegate tochat) 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 winRaise
NotImplementedErrorto fail fast and comply with the baseLLMcontract.Line 70 uses
pass, which returnsNoneinstead of the expectedAsyncIterator[str]. Although the comment states the retriever won't call this method, usingpassviolates the abstract base class contract and risks a silentTypeErrorat runtime if streaming is ever attempted. Raise an explicitNotImplementedErrorinstead.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
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (19)
REFACTORING_DECISION_LOG.mdopenrag/components/indexer/chunker/chunker.pyopenrag/components/indexer/embeddings/__init__.pyopenrag/components/indexer/embeddings/openai.pyopenrag/components/llm.pyopenrag/components/pipeline.pyopenrag/components/reranker/__init__.pyopenrag/components/reranker/base.pyopenrag/components/reranker/infinity.pyopenrag/components/reranker/openai.pyopenrag/components/reranker/test_rrf_reranking.pyopenrag/components/retriever.pyopenrag/core/llm/llm.pyopenrag/core/retrieval/rrf.pyopenrag/services/inference/__init__.pyopenrag/services/inference/_circuit_breaker.pyopenrag/services/inference/test_circuit_breaker.pyopenrag/services/inference/test_vllm_client.pyopenrag/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
885a0cd to
28a7100
Compare
Summary by CodeRabbit