feat(llm): add default framework with OpenAI-compatible client - #1797
Conversation
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
ca11b55 to
344c8f7
Compare
9287b4b to
c574bb4
Compare
533aeb3 to
5e9d170
Compare
Greptile SummaryThis PR introduces a native OpenAI-compatible HTTP client (
|
| Filename | Overview |
|---|---|
| nemoguardrails/llm/clients/_errors.py | New error classification module; "is not supported" keyword in _UNSUPPORTED_PARAMS_KEYWORDS is too broad and will misclassify feature/region/plan unavailability errors as parameter errors |
| nemoguardrails/llm/clients/base.py | New BaseClient with retry logic, SSE streaming, and header building; retry/backoff and cleanup logic look correct |
| nemoguardrails/llm/clients/openai_chat_model.py | Chat model adapter with streaming tool call accumulation; previously-flagged issues (tc["id"] KeyError, tc_delta["index"] KeyError) are fixed; reasoning-model heuristic is tightened |
| nemoguardrails/llm/clients/openai_compatible.py | Thin OpenAI-compatible client; stream_options injection is now opt-outable via include_usage_in_stream=False; deterministic aclose chain looks correct |
| nemoguardrails/llm/default_framework.py | New DefaultFramework; client-pool keying uses json.dumps (stable for unhashable values, previously flagged and fixed); reset() behavior documented as destructive to provider registrations |
| nemoguardrails/llm/frameworks.py | Framework registry with async-safe reset; _areset_frameworks uses finally block for guaranteed cleanup; per-framework errors are caught and logged; _reset_frameworks limitation in running loop is documented via test |
| nemoguardrails/llm/clients/_sse.py | SSE decoder correctly handles multi-line events, id clearing, retry fields, and trailing event flush |
| nemoguardrails/exceptions.py | Adds structured LLM exception hierarchy with status code, error metadata, and provider context; clean design |
| tests/llm/test_frameworks.py | Comprehensive new tests for async reset, error isolation, missing-reset-method, and default-framework behavior; explicitly tests the asyncio.run limitation in running loops |
| tests/conftest.py | Adds langchain_framework fixture that switches default framework and resets after; used correctly in sync test contexts only |
Sequence Diagram
sequenceDiagram
participant Caller
participant OpenAIChatModel
participant OpenAICompatibleClient
participant BaseClient
participant Provider as OpenAI-compatible Provider
Caller->>OpenAIChatModel: generate_async(prompt, **kwargs)
OpenAIChatModel->>OpenAIChatModel: _to_messages() + _prepare_params()
OpenAIChatModel->>OpenAICompatibleClient: chat_completion(model, messages, **params)
OpenAICompatibleClient->>OpenAICompatibleClient: _build_payload()
OpenAICompatibleClient->>BaseClient: _apost("/chat/completions", payload)
loop Retry (up to max_retries)
BaseClient->>Provider: POST /chat/completions
alt 2xx
Provider-->>BaseClient: JSON response
else retryable (408/429/5xx)
BaseClient->>BaseClient: sleep (backoff/Retry-After)
else 4xx error
BaseClient->>BaseClient: raise_for_status() → LLMClientError subclass
end
end
BaseClient-->>OpenAICompatibleClient: response dict
OpenAICompatibleClient-->>OpenAIChatModel: response dict
OpenAIChatModel->>OpenAIChatModel: _parse_response()
OpenAIChatModel-->>Caller: LLMResponse
Caller->>OpenAIChatModel: stream_async(prompt, **kwargs)
OpenAIChatModel->>OpenAICompatibleClient: stream_chat_completion(model, messages)
OpenAICompatibleClient->>BaseClient: _apost_stream (SSE)
loop SSE chunks
Provider-->>BaseClient: data: {...}
BaseClient->>BaseClient: SSEDecoder.decode()
BaseClient-->>OpenAICompatibleClient: parsed chunk
OpenAICompatibleClient-->>OpenAIChatModel: chunk dict
OpenAIChatModel->>OpenAIChatModel: accumulate tool_calls + _parse_chunk()
OpenAIChatModel-->>Caller: LLMResponseChunk
end
Note over OpenAIChatModel,BaseClient: aclose() chain on exit/abort
Prompt To Fix All With AI
This is a comment left during a code review.
Path: nemoguardrails/llm/clients/_errors.py
Line: 44-50
Comment:
**`"is not supported"` keyword too broad — same class of problem as the removed `"not allowed"`**
Any 400 response whose message contains "is not supported" — including "Streaming is not supported for this endpoint", "Image generation is not supported in your plan", or "This model is not supported in your region" — will be raised as `LLMUnsupportedParamsError`. Callers that catch `LLMUnsupportedParamsError` to retry with different parameters will silently swallow these capability/region/plan errors and get the same rejection on retry, masking the real cause. The previous iteration removed `"not allowed"` for this exact reason; `"is not supported"` should get the same treatment or be replaced with tighter phrases like `"parameter is not supported"` or `"unsupported parameter"` (which is already in the list).
How can I resolve this? If you propose a fix, please make it concise.Reviews (11): Last reviewed commit: "fix(llm/default_framework): use stable s..." | Re-trigger Greptile
tgasser-nv
left a comment
There was a problem hiding this comment.
A high-level comment, this PR is way too large. It could easily have been stacked into 4 (or more) PRs.
- New client implementation.
- Client fixture JSON and code to re-generate them later.
- Moving tests (without edits).
- New / modified tests.
Non-blocking feedback / comments:
- What is the customer-benefit of this change? Lightweight dependencies compared with Langchain? Better performance? More specific customizations.
- Can you measure performance difference between Langchain and Default LLM clients to make sure there are no regressions?
Introduces structured exception types for LLM client failures: LLMClientError base with model_name/provider_name/base_url enrichment, plus LLMAuthenticationError, LLMRateLimitError (with retry_after_seconds), LLMBadRequestError (with LLMContextWindowError and LLMUnsupportedParamsError subtypes), LLMServerError, and LLMResponseValidationError. These replace ad-hoc LangChain exceptions and are consumed by the new HTTP transport and OpenAI chat model in subsequent commits.
Introduces SSEDecoder used by streaming chat completions. Implements the HTML5 Server-Sent Events parsing spec (multi-line data fields, comments, event boundaries on blank lines) and detects mid-stream error payloads so upstream provider errors surface as exceptions instead of silent truncation. Standalone leaf utility with direct tests; used by the HTTP transport introduced in the next commit.
Introduces the transport layer underneath the LLMModel protocol. BaseClient owns the httpx.AsyncClient, retry algorithm (exponential backoff with jitter, Retry-After and x-should-retry header honored), connection pool, and error-status handling. OpenAICompatibleClient is a thin provider-aware wrapper exposing chat_completion and stream_chat_completion against the /v1/chat/completions path. Errors raised by the transport are classified from the HTTP status and response body (OpenAI, vLLM, and generic detail formats supported), enriched with provider_name and base_url, and have API keys redacted from the message. All retry/error/payload behavior is covered in the transport test file, which mocks httpx to avoid coupling to the protocol layer.
Introduces the protocol layer on top of OpenAICompatibleClient. OpenAIChatModel exposes generate_async/stream_async, serializes ChatMessage lists (and tool-call arguments) into the wire format, parses chat-completion responses into LLMResponse and streaming chunks into LLMResponseChunk, accumulates tool-call fragments across stream chunks, strips temperature for reasoning models, extracts provider metadata and response headers, and enriches transport errors with the model name. Tests mock the client (client.chat_completion / stream_chat_completion return raw dicts) and assert parsing, validation, serialization, and streaming accumulation. test_client_config covers construction knobs (timeout, connect_timeout, http_client injection and lifecycle, custom_headers, custom_query) on both the client and the model.
Lands end-to-end coverage for the new client/model stack. record_fixtures.py captures real OpenAI and NIM responses (generate/stream x text/tool_call/ reasoning) once, producing the JSON fixtures under fixtures/. The live test file runs those fixtures through a mocked httpx transport for deterministic CI coverage, and has an opt-in suite that hits the real OpenAI and NIM endpoints when API keys are present. Separated from the protocol-layer commit so reviewers can skim the model code without scrolling through ~1200 lines of recorded JSON.
DefaultFramework wires OpenAIChatModel into the main engine as an LLMFramework implementation. Constructs one OpenAICompatibleClient per distinct (base_url, api_key) pair and reuses it across models to share connection pools; the pool key hashes the api_key so secrets never appear in cache keys or logs. Adds pipeline-level tests for _stream_llm_call that exercise accumulation of delta content, reasoning, tool calls, usage, request ids, and provider metadata into the corresponding context vars, with stale state cleared when chunks omit those fields. The framework is registered in the next commit as part of the cutover.
Registers DefaultFramework, flips the default from langchain to default,
and migrates existing tests to the new protocol.
- frameworks.py: NEMOGUARDRAILS_LLM_FRAMEWORK defaults to default, and
get_framework("default") lazily constructs DefaultFramework.
- tests/conftest.py: adds an opt-in langchain_framework fixture for the
tests that still need LangChain-specific behavior.
- tests/integrations/langchain/: LangChain-only tests move here, with an
autouse fixture applying langchain_framework, plus the LangChain copy
of test_configs that exercise custom_llm/custom_chat_model hooks.
- tests/test_configs/**: rewritten to use LLMModel directly instead of
BaseChatModel/LLM subclasses.
- tests/test_llm_params_e2e.py, test_supported_llm_providers.py,
test_task_specific_model.py, tests/llm/test_frameworks.py: updated
for the new default and for chat-only models.
- examples/bots/abc/config.yml, test_task_specific_model.py: switch
gpt-3.5-turbo-instruct (completion) to a chat model, since the new
stack is chat-only.
📝 WalkthroughWalkthroughThis pull request introduces a new default OpenAI-compatible LLM framework alongside the existing LangChain integration. It adds standardized HTTP client infrastructure with retry logic, error handling, SSE parsing, and support for OpenAI-compatible chat completions. A comprehensive test suite with fixtures validates the implementation. Changes
Sequence Diagram(s)sequenceDiagram
participant User as User Code
participant Framework as DefaultFramework
participant Client as OpenAICompatibleClient
participant Base as BaseClient
participant HTTP as OpenAI API
participant Model as OpenAIChatModel
User->>Framework: create_model("gpt-4", "openai")
Framework->>Framework: resolve base_url & api_key
Framework->>Framework: _get_or_create_client
Framework->>Client: new OpenAICompatibleClient(...)
Framework->>Model: new OpenAIChatModel(client, model)
Framework-->>User: OpenAIChatModel
User->>Model: generate_async(prompt, **kwargs)
Model->>Model: convert prompt to messages
Model->>Client: chat_completion(model, messages, **kwargs)
Client->>Client: _build_payload
Client->>Base: _apost("/chat/completions", payload)
Base->>HTTP: POST /chat/completions
HTTP-->>Base: 200 OK + JSON response
Base->>Base: parse response.json()
Base-->>Client: dict response
Client-->>Model: response dict
Model->>Model: parse response fields
Model->>Model: validate structure
Model-->>User: LLMResponse(content, finish_reason, usage, ...)
sequenceDiagram
participant Client as BaseClient
participant HTTP as OpenAI API
participant Decoder as SSEDecoder
participant Model as OpenAIChatModel
Client->>HTTP: POST /chat/completions (stream=true)
HTTP-->>Client: 200 OK + streaming body
loop For each streamed chunk
Client->>Decoder: decode(line)
Decoder-->>Client: ServerSentEvent
Client->>Client: parse event.data JSON
Client->>Client: check for [DONE]
alt Error in stream
Client->>Client: _check_sse_error
Client-->>Client: raise LLMClientError variant
end
end
Client-->>Model: async generator of chunks
Model->>Model: accumulate deltas
Model->>Model: aggregate tool_calls
Model-->>User: async iterator LLMResponseChunk
sequenceDiagram
participant Client as BaseClient
participant HTTP as OpenAI API
Client->>HTTP: POST request
HTTP-->>Client: 5xx or 429 response
alt Retryable Status
Client->>Client: parse retry-after header
Client->>Client: exponential backoff + jitter
Client->>HTTP: retry POST request
else Non-Retryable Status
Client->>Client: raise_for_status(status_code)
Client->>Client: map to LLMClientError subclass
Client-->>Client: raise exception
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 15
♻️ Duplicate comments (4)
nemoguardrails/llm/clients/_errors.py (1)
103-115:⚠️ Potential issue | 🟠 MajorRedact
bodyandresponse_headersbefore attaching them to the exception.Only
error_messageis sanitized here.kwargs["body"]andkwargs["response_headers"]still carry provider data verbatim, so a provider that echoes tokens or auth headers will still leak secrets throughexc.body/exc.response_headerseven though the message is redacted.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemoguardrails/llm/clients/_errors.py` around lines 103 - 115, The function _build_error_fields currently only redacts error_message; update it to sanitize the parsed_body and headers before attaching them to kwargs by applying _redact_secrets (or converting non-string bodies/headers to strings and then redacting) so that kwargs["body"] and kwargs["response_headers"] do not contain raw secrets; do this in _build_error_fields (before creating kwargs) using the existing _redact_secrets helper and then attach the redacted values (instead of raw parsed_body / dict(headers)) while keeping ctx.as_kwargs() unchanged.nemoguardrails/llm/clients/base.py (2)
158-162:⚠️ Potential issue | 🟠 MajorWrap successful-response JSON parsing in a typed validation error.
A
200response with HTML/plain text currently bubbles up as a rawJSONDecodeError, which bypasses the new LLM client exception hierarchy and loses the provider/model context.Suggested fix
-from nemoguardrails.exceptions import LLMConnectionError, LLMTimeoutError +from nemoguardrails.exceptions import ( + LLMConnectionError, + LLMResponseValidationError, + LLMTimeoutError, +) ... if response.status_code >= 400: raise_for_status(response.status_code, response.text, response.headers, ctx) - data = response.json() + try: + data = response.json() + except json.JSONDecodeError as err: + raise LLMResponseValidationError( + f"Provider returned non-JSON response: {err}", + response_data=None, + **ctx.as_kwargs(), + ) from err data["_response_headers"] = dict(response.headers) return data🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemoguardrails/llm/clients/base.py` around lines 158 - 162, The code calling response.json() in the response handling path should catch JSONDecodeError and re-raise a typed validation/response error (instead of letting JSONDecodeError bubble) so provider and model context from ctx aren't lost; update the block around response.json() in the handler (the same place that calls raise_for_status and sets data["_response_headers"]) to wrap the JSON parsing in a try/except, on JSONDecodeError construct and raise a specific LLM client validation error (e.g., ResponseValidationError or a suitable existing client exception) that includes the original exception as the cause, the raw response.text, response.headers (like data["_response_headers"]), the status_code and the ctx/provider/model info. Ensure you still populate/attach response headers and preserve the original exception (using from) so downstream callers can inspect both the typed error and the original parsing error.
87-92:⚠️ Potential issue | 🟠 MajorMerge auth headers case-insensitively before sending the request.
HTTP header names are case-insensitive. With the current
dict.update(),custom_headers={"authorization": "..."}can coexist withAuthorization, so the request may carry ambiguous auth state instead of a deterministic override or rejection.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemoguardrails/llm/clients/base.py` around lines 87 - 92, The _build_headers method currently uses dict.update which can leave duplicate auth headers with different casing; change it to merge headers case-insensitively so custom headers deterministically override built headers (especially Authorization). In _build_headers, build the initial headers dict (including Authorization from self._api_key if present), then iterate over self._custom_headers and for each key compare key.lower() against existing header keys lowercased; if a match exists, replace the existing header value with the custom one (rather than adding a new key), otherwise insert the custom header preserving its original casing; ensure this logic applies to "Authorization"/"authorization" so the custom header wins.nemoguardrails/exceptions.py (1)
199-216:⚠️ Potential issue | 🟠 MajorMake response-validation failures a
LLMServerErrorsubtype.
LLMResponseValidationErroris raised for malformed provider responses, and it already carries a synthetic502. SinceLLMClientErrorexplicitly tells callers to branch on exception class, keeping this outsideLLMServerErrormeans the same server-side failure won't be caught by server/transient handlers.Proposed fix
-class LLMResponseValidationError(LLMClientError): +class LLMResponseValidationError(LLMServerError):🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemoguardrails/exceptions.py` around lines 199 - 216, LLMResponseValidationError is currently subclassing LLMClientError but should subclass LLMServerError so server-side/transient handlers treat response-validation failures as server errors; modify the class declaration for LLMResponseValidationError to inherit from LLMServerError (keep the __init__ signature and the synthetic status_code=502, response_data, model_name, provider_name, base_url handling unchanged) and ensure any necessary import for LLMServerError is present.
🧹 Nitpick comments (9)
nemoguardrails/llm/clients/constants.py (1)
20-20: Consider lowering default connection limits to safer baseline values.
max_connections=1000per client is aggressive and can cause avoidable FD/memory pressure when multiple pooled clients exist.Proposed safer default
-DEFAULT_CONNECTION_LIMITS = httpx.Limits(max_connections=1000, max_keepalive_connections=100) +DEFAULT_CONNECTION_LIMITS = httpx.Limits(max_connections=100, max_keepalive_connections=20)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemoguardrails/llm/clients/constants.py` at line 20, DEFAULT_CONNECTION_LIMITS is set to an overly high default (httpx.Limits with max_connections=1000), which risks file-descriptor and memory pressure; change the constant DEFAULT_CONNECTION_LIMITS to use much lower, safer defaults (for example lowering max_connections and max_keepalive_connections to a reasonable baseline such as ~100 and ~20 respectively) by updating the assignment in constants.py (the DEFAULT_CONNECTION_LIMITS definition that constructs httpx.Limits) and adjust any tests or documentation that assume the old values.tests/conftest.py (1)
38-44: Initialize the fixture from a clean framework state.This fixture only resets after the test. Pre-resetting before
set_default_framework("langchain")avoids inheriting state from previously run tests.♻️ Suggested update
`@pytest.fixture` def langchain_framework(): from nemoguardrails.llm.frameworks import _reset_frameworks, set_default_framework + _reset_frameworks() set_default_framework("langchain") yield _reset_frameworks()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/conftest.py` around lines 38 - 44, The fixture langchain_framework should reset frameworks before setting the default to avoid inheriting prior test state; call _reset_frameworks() before set_default_framework("langchain"), keep the existing yield, and retain the post-test _reset_frameworks() cleanup so the sequence is: _reset_frameworks(), set_default_framework("langchain"), yield, _reset_frameworks().tests/llm/clients/fixtures/nim_stream_tool_calls.json (1)
334-342: Consider adding fragmentedtool_calls.function.argumentschunks for stronger stream coverage.This fixture currently provides the full arguments JSON in one delta, so it doesn’t exercise incremental argument accumulation in streaming. Splitting arguments across 2+ deltas would better validate the aggregation path.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/llm/clients/fixtures/nim_stream_tool_calls.json` around lines 334 - 342, The fixture's tool call currently supplies the full JSON in one delta; update the "tool_calls" entry for id "call_1fab8199a2fd4c23acf6fac5" so the nested "function.arguments" value is split across multiple streamed deltas (e.g., two or more fragments that concatenate to form "{\"city\": \"Paris\"}") to exercise incremental argument accumulation logic; ensure the sequence order and indices remain valid so the test harness will aggregate the fragments into the original JSON string.tests/llm/clients/record_fixtures.py (1)
70-77: Renameaiterparameter to avoid shadowing Python builtin.
aiteris a Python builtin (since 3.10). Consider renaming tostreamorasync_iter.♻️ Proposed fix
-async def _try_record_stream(name, aiter): +async def _try_record_stream(name, stream): try: chunks = [] - async for chunk in aiter: + async for chunk in stream: chunks.append(chunk) save(name, chunks) except Exception as e: print(f" FAILED {name}: {e}")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/llm/clients/record_fixtures.py` around lines 70 - 77, The parameter name aiter in the function _try_record_stream shadows the Python builtin aiter; rename the parameter (for example to stream or async_iter) and update all uses inside _try_record_stream (the async for loop and any references like chunks append/save) to the new name to avoid builtin shadowing while preserving the existing behavior.nemoguardrails/llm/clients/openai_compatible.py (1)
24-35: Provider detection relies on URL heuristics.The substring-based provider detection may misclassify URLs (e.g., a proxy URL containing "openai" in the path). Consider documenting this behavior or allowing explicit provider override via constructor parameter.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemoguardrails/llm/clients/openai_compatible.py` around lines 24 - 35, The provider_name property uses fragile substring heuristics on self._base_url and can misclassify proxied or path-containing URLs; update the class (add an optional constructor parameter like provider or explicit_provider to the __init__ of the OpenAI-compatible client class) to accept an explicit provider override and have provider_name return that override when set, keep the heuristic as fallback, and add a short docstring comment on the class/__init__ explaining that heuristics are fallback-only and users should prefer the explicit provider parameter.nemoguardrails/llm/clients/_sse.py (2)
20-37: Parameteridshadows Python builtin (acceptable for SSE field).The static analysis correctly flags that
idshadows the builtin. However, since this is a standard SSE field name and the class is narrowly scoped, this is acceptable. You could rename toevent_idfor clarity, but it's not required.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemoguardrails/llm/clients/_sse.py` around lines 20 - 37, The constructor parameter name id in ServerSentEvent shadows the Python builtin; rename the parameter to event_id and update internal storage (change self._id to self._event_id) and any usages of the attribute to the new name in ServerSentEvent.__init__ and elsewhere to avoid shadowing and improve clarity; alternatively, if you prefer to keep the SSE name, add a lint suppression comment (e.g., noqa for W0622) next to the parameter and document the reason.
76-92: Persistent_last_event_idcauses events on consecutive blank lines after id is set.Per the SSE spec,
_last_event_idis intentionally not reset after dispatch (correct). However, this means once anidis set, any subsequent blank line will dispatch an event (since line 78's condition checksnot self._last_event_id). This may dispatch "empty" events with only an id after the first event with that id.If this is intentional behavior for reconnection handling, consider adding a brief comment explaining why
_last_event_idis excluded from the reset on lines 88-90.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemoguardrails/llm/clients/_sse.py` around lines 76 - 92, The decode method currently preserves _last_event_id while resetting _event, _data, and _retry which can make blank lines after an id dispatch an event; add a brief explanatory comment immediately above the block that resets state in decode (the lines that set self._event = "", self._data = [], self._retry = None) clarifying that per the SSE spec last-event-id must be retained across dispatches for reconnection semantics, so _last_event_id is intentionally not cleared while other per-event fields are reset; reference decode, _last_event_id, _event, _data, _retry and ServerSentEvent in the comment so future readers understand the rationale.tests/llm/clients/test_sse.py (1)
195-198: Use specific exception type instead of bareException.The test should assert
json.JSONDecodeError(orValueErrorfor older Python) instead of catching anyException. This makes the test more precise and avoids masking unexpected failures.♻️ Proposed fix
+import json + def test_json_raises_on_invalid(self): sse = ServerSentEvent(data="not json") - with pytest.raises(Exception): + with pytest.raises(json.JSONDecodeError): sse.json()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/llm/clients/test_sse.py` around lines 195 - 198, Update the test_json_raises_on_invalid test to assert the specific JSON decoding exception instead of a bare Exception: call ServerSentEvent(data="not json").json() inside a pytest.raises that expects json.JSONDecodeError (falling back to ValueError on older Python if needed). Locate the test function test_json_raises_on_invalid in tests/llm/clients/test_sse.py and replace the generic pytest.raises(Exception) with pytest.raises(json.JSONDecodeError) (or pytest.raises(ValueError) behind a version check or try/except import) so the test precisely verifies ServerSentEvent.json() error behavior.tests/llm/clients/fixtures/openai_multiturn_tool_roundtrip.json (1)
51-78: Consider redacting sensitive metadata from fixtures.The
_response_headerssections contain production identifiers (openai-organization,openai-project) and session cookies (set-cookie) that could be sensitive. While these are test fixtures, consider redacting or replacing these values with placeholder strings to avoid accidentally committing real organizational identifiers.Also applies to: 122-148
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/llm/clients/fixtures/openai_multiturn_tool_roundtrip.json` around lines 51 - 78, The fixture's _response_headers contains real identifiers and cookies; update the JSON in openai_multiturn_tool_roundtrip.json by replacing the values for keys like "openai-organization", "openai-project", and "set-cookie" (and any other org/project/cookie-like headers elsewhere in the file such as the other _response_headers block) with non-sensitive placeholders (e.g., "REDACTED_ORG", "REDACTED_PROJECT", "REDACTED_COOKIE") so tests keep structure but no production identifiers are committed.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@nemoguardrails/integrations/langchain/llm_adapter.py`:
- Around line 265-266: The reset method currently returns immediately and must
instead clear framework-owned provider state so _reset_frameworks() can perform
cleanup; modify async def reset(self) in the LLM adapter to asynchronously clear
any internal provider references and state (e.g., set cached client/connection
attributes to None, cancel/await any pending tasks, and call/await any
provider-specific cleanup methods such as client.close() or provider.reset() if
available), ensure all awaited cleanup completes before returning, and keep this
behavior invoked by _reset_frameworks().
In `@nemoguardrails/llm/clients/base.py`:
- Around line 104-113: The _calculate_retry_delay function only parses numeric
delta-seconds from headers["retry-after"]; update it to also accept HTTP-date
form per RFC7231 by attempting to parse the value as an RFC‑1123/HTTP-date
(e.g., use email.utils.parsedate_to_datetime or equivalent) when float(...)
fails, compute the seconds as (parsed_date - now).total_seconds(), validate 0 <
delay <= MAX_RETRY_AFTER, return that delay if valid, and keep existing debug
logs for out-of-range or parse failures; reference the _calculate_retry_delay
function, the headers["retry-after"] value, MAX_RETRY_AFTER constant, and
log.debug calls when implementing this behavior.
In `@nemoguardrails/llm/clients/openai_chat_model.py`:
- Around line 242-279: _parsed_chunk currently only preserves _response_headers
into provider_metadata causing loss of other top-level provider fields (e.g.,
system_fingerprint) present in non-streaming paths; update the _parse_chunk
method to collect the same non-standard top-level fields as _parse_response and
include them in provider_metadata (merge existing _response_headers with any
other keys from the incoming data that are not the known standard keys like
"choices", "id", "model", "usage", "delta", "finish_reason", etc.), and then
pass that merged provider_metadata into the returned LLMResponseChunk so
streaming chunks and non-streaming responses preserve the same provider-specific
metadata.
- Around line 123-135: The code only sets "id", "type", and "function_name" when
a new index is first created in tool_call_acc, but later tc_delta chunks may
include those fields; update the logic in the loop that processes raw_tool_calls
so that for each tc_delta you always overwrite or set tool_call_acc[idx]["id"],
["type"], and ["function_name"] if the incoming tc_delta provides non-empty
values (while continuing to append argument fragments into
["arguments_buffer"]). Refer to the variables/tool structures tool_call_acc,
tc_delta, raw_tool_calls and ensure you do not lose existing arguments_buffer
content when updating metadata.
In `@nemoguardrails/llm/default_framework.py`:
- Around line 69-77: The pool key currently embeds the raw api_key in the tuple
assigned to variable "key", which retains secrets in self._clients; instead
compute a stable non-reversible fingerprint (e.g., SHA-256 hex digest) of
api_key and use that fingerprint (or empty string when no api_key) in the tuple;
add the necessary import (hashlib) and replace the api_key usage in the key
construction in default_framework.py so the rest of the tuple items (base_url,
timeout, connect_timeout, max_retries, default_headers/default_query tuples)
remain unchanged.
- Around line 132-136: The reset() loop currently swallows exceptions from await
client.close(), which lets failures go unnoticed while _clients and _providers
are still cleared; update reset() (the loop iterating over
self._clients.values() and the subsequent clearing of self._clients/_providers)
to collect any exceptions from client.close() instead of silently passing,
attempt to close all clients, and if any close fails raise or return an
aggregated exception (or re-raise the first) so callers know the reset did not
fully release transports — ensure you only clear self._clients and
self._providers after successful closes (or after recording failures and still
raising) so leaked transports are not hidden.
In `@nemoguardrails/llm/frameworks.py`:
- Around line 64-70: The current _reset_frameworks implementation calls
asyncio.run(fw.reset()) which fails silently if invoked from an already-running
event loop; add an async-safe entry point and update callers: implement an async
function reset_frameworks_async that awaits fw.reset() for each framework
(iterating over list(_frameworks.values())), and change or add a synchronous
wrapper reset_frameworks that checks for a running loop (using
asyncio.get_running_loop() or asyncio.get_event_loop()/loops) and calls
asyncio.run(reset_frameworks_async()) when safe; update or document that
_reset_frameworks (or the existing public reset helper) must be used only from
sync contexts if you prefer keeping the sync-only API, and ensure
_frameworks.clear() and resetting _default_framework remain executed after the
awaited resets.
In `@tests/llm/clients/_helpers.py`:
- Around line 36-38: The helpers currently create a real OpenAICompatibleClient
that owns an httpx.AsyncClient and then replace client._client with test
doubles, leaking the original pool and desynchronizing client._owns_client; fix
by ensuring you either construct OpenAICompatibleClient with the test double
injected (avoid creating a real httpx.AsyncClient in make_client) or, if you
must replace client._client in mock_httpx_post() / stream_client(), first check
client._owns_client and if true await client._client.aclose() before assigning
the replacement, then set client._client to the mock and update
client._owns_client to reflect ownership of the new object so cleanup paths
remain correct.
In `@tests/llm/clients/fixtures/openai_error_401.json`:
- Around line 12-30: The fixture contains sensitive session-like headers in the
response_headers block (notably the "set-cookie" header and similar runtime IDs
like "x-request-id"); update the openai_error_401.json fixture to remove or
replace these values with non-sensitive placeholders (e.g. set "set-cookie":
"[REDACTED]" or delete the key) and sanitize any other runtime headers in
response_headers so tests keep structure but no real tokens or cookies are
committed.
In `@tests/llm/clients/fixtures/openai_generate_finish_length.json`:
- Around line 47-50: Replace real tenant/project/request identifiers in the
fixture headers with stable redacted placeholders: update the keys
"openai-organization", "openai-project", and any "x-request-id" (and similar
request/tenant headers such as "openai-version" if necessary) in
tests/llm/clients/fixtures/openai_generate_finish_length.json to use neutral
placeholders like "REDACTED_ORG", "REDACTED_PROJECT", and "REDACTED_REQUEST_ID"
(or similar) so no internal identifiers are committed; ensure the placeholders
are consistent with other fixtures referenced in the test suite.
In `@tests/llm/clients/fixtures/openai_generate_refusal.json`:
- Around line 36-59: The fixture
tests/llm/clients/fixtures/openai_generate_refusal.json contains
tenant-identifying header values; update the _response_headers object to replace
real identifiers with synthetic placeholders (e.g. replace values for
"openai-organization", "openai-project", "x-request-id", "cf-ray" and any other
provider-specific ids) so no internal org/project/request IDs remain, keeping
header keys intact but using neutral values like "org_test", "proj_test",
"req_test", "ray_test" to preserve structure.
In `@tests/llm/clients/fixtures/openai_generate_tool_call.json`:
- Around line 54-77: The committed fixture exposes tenant/session-specific
headers under _response_headers (e.g., openai-organization, openai-project,
x-request-id, set-cookie); sanitize those entries by replacing real values with
neutral placeholders (for example "REDACTED_ORG", "REDACTED_PROJECT",
"REQ_ID_REDACTED", "REDACTED_COOKIE") or remove the sensitive keys entirely so
tests still run; update the fixture file's _response_headers block to use these
placeholders for openai-organization, openai-project, x-request-id, set-cookie
(and any other environment-specific headers) and add a short comment explaining
that values are intentionally redacted.
In `@tests/llm/clients/test_client_config.py`:
- Around line 32-113: Several sync tests create httpx.AsyncClient or
OpenAICompatibleClient instances without closing them; convert those tests that
instantiate AsyncClient/OpenAICompatibleClient (e.g., TestTimeout.test_defaults,
TestTimeout.test_custom, TestTimeout.test_http_client_timeout_inferred,
TestConnectionPool.test_limits, TestCustomHeaders.* ,
TestCustomQuery.test_stored, TestHttpClientInjection.test_uses_injected_client,
TestHttpClientInjection.test_invalid_type_raises) to async pytest tests (use
`@pytest.mark.asyncio`) and ensure you close resources in a finally block by
calling await custom.aclose() for injected httpx.AsyncClient and await
client.close() for owned OpenAICompatibleClient; keep the existing tests that
already use async/await unchanged. Ensure assertions remain the same but run
after proper cleanup handling.
In `@tests/llm/clients/test_openai_compatible.py`:
- Around line 48-63: The tests in TestProviderName (test_openai, test_nim,
test_local, test_azure) construct OpenAICompatibleClient instances without
closing the internally owned httpx.AsyncClient, leaking transports; update each
test to either (a) create and pass a managed http_client to
OpenAICompatibleClient or (b) make the test async and await c.aclose() after
assertions (e.g., await c.aclose()), ensuring the internal AsyncClient is
properly closed; apply the same change to the analogous tests at the other noted
locations.
In `@tests/test_task_specific_model.py`:
- Around line 90-92: Add a non-empty assertion for the filtered list before
iterating so the test fails if there are no non-task-specific calls;
specifically, after computing other_calls (from res.log.llm_calls filtered by
task_specific_tasks) add an assertion like "assert other_calls" or "assert
len(other_calls) > 0" with a clear message, then proceed to iterate and assert
each call.llm_model_name == "gpt-4o".
---
Duplicate comments:
In `@nemoguardrails/exceptions.py`:
- Around line 199-216: LLMResponseValidationError is currently subclassing
LLMClientError but should subclass LLMServerError so server-side/transient
handlers treat response-validation failures as server errors; modify the class
declaration for LLMResponseValidationError to inherit from LLMServerError (keep
the __init__ signature and the synthetic status_code=502, response_data,
model_name, provider_name, base_url handling unchanged) and ensure any necessary
import for LLMServerError is present.
In `@nemoguardrails/llm/clients/_errors.py`:
- Around line 103-115: The function _build_error_fields currently only redacts
error_message; update it to sanitize the parsed_body and headers before
attaching them to kwargs by applying _redact_secrets (or converting non-string
bodies/headers to strings and then redacting) so that kwargs["body"] and
kwargs["response_headers"] do not contain raw secrets; do this in
_build_error_fields (before creating kwargs) using the existing _redact_secrets
helper and then attach the redacted values (instead of raw parsed_body /
dict(headers)) while keeping ctx.as_kwargs() unchanged.
In `@nemoguardrails/llm/clients/base.py`:
- Around line 158-162: The code calling response.json() in the response handling
path should catch JSONDecodeError and re-raise a typed validation/response error
(instead of letting JSONDecodeError bubble) so provider and model context from
ctx aren't lost; update the block around response.json() in the handler (the
same place that calls raise_for_status and sets data["_response_headers"]) to
wrap the JSON parsing in a try/except, on JSONDecodeError construct and raise a
specific LLM client validation error (e.g., ResponseValidationError or a
suitable existing client exception) that includes the original exception as the
cause, the raw response.text, response.headers (like data["_response_headers"]),
the status_code and the ctx/provider/model info. Ensure you still
populate/attach response headers and preserve the original exception (using
from) so downstream callers can inspect both the typed error and the original
parsing error.
- Around line 87-92: The _build_headers method currently uses dict.update which
can leave duplicate auth headers with different casing; change it to merge
headers case-insensitively so custom headers deterministically override built
headers (especially Authorization). In _build_headers, build the initial headers
dict (including Authorization from self._api_key if present), then iterate over
self._custom_headers and for each key compare key.lower() against existing
header keys lowercased; if a match exists, replace the existing header value
with the custom one (rather than adding a new key), otherwise insert the custom
header preserving its original casing; ensure this logic applies to
"Authorization"/"authorization" so the custom header wins.
---
Nitpick comments:
In `@nemoguardrails/llm/clients/_sse.py`:
- Around line 20-37: The constructor parameter name id in ServerSentEvent
shadows the Python builtin; rename the parameter to event_id and update internal
storage (change self._id to self._event_id) and any usages of the attribute to
the new name in ServerSentEvent.__init__ and elsewhere to avoid shadowing and
improve clarity; alternatively, if you prefer to keep the SSE name, add a lint
suppression comment (e.g., noqa for W0622) next to the parameter and document
the reason.
- Around line 76-92: The decode method currently preserves _last_event_id while
resetting _event, _data, and _retry which can make blank lines after an id
dispatch an event; add a brief explanatory comment immediately above the block
that resets state in decode (the lines that set self._event = "", self._data =
[], self._retry = None) clarifying that per the SSE spec last-event-id must be
retained across dispatches for reconnection semantics, so _last_event_id is
intentionally not cleared while other per-event fields are reset; reference
decode, _last_event_id, _event, _data, _retry and ServerSentEvent in the comment
so future readers understand the rationale.
In `@nemoguardrails/llm/clients/constants.py`:
- Line 20: DEFAULT_CONNECTION_LIMITS is set to an overly high default
(httpx.Limits with max_connections=1000), which risks file-descriptor and memory
pressure; change the constant DEFAULT_CONNECTION_LIMITS to use much lower, safer
defaults (for example lowering max_connections and max_keepalive_connections to
a reasonable baseline such as ~100 and ~20 respectively) by updating the
assignment in constants.py (the DEFAULT_CONNECTION_LIMITS definition that
constructs httpx.Limits) and adjust any tests or documentation that assume the
old values.
In `@nemoguardrails/llm/clients/openai_compatible.py`:
- Around line 24-35: The provider_name property uses fragile substring
heuristics on self._base_url and can misclassify proxied or path-containing
URLs; update the class (add an optional constructor parameter like provider or
explicit_provider to the __init__ of the OpenAI-compatible client class) to
accept an explicit provider override and have provider_name return that override
when set, keep the heuristic as fallback, and add a short docstring comment on
the class/__init__ explaining that heuristics are fallback-only and users should
prefer the explicit provider parameter.
In `@tests/conftest.py`:
- Around line 38-44: The fixture langchain_framework should reset frameworks
before setting the default to avoid inheriting prior test state; call
_reset_frameworks() before set_default_framework("langchain"), keep the existing
yield, and retain the post-test _reset_frameworks() cleanup so the sequence is:
_reset_frameworks(), set_default_framework("langchain"), yield,
_reset_frameworks().
In `@tests/llm/clients/fixtures/nim_stream_tool_calls.json`:
- Around line 334-342: The fixture's tool call currently supplies the full JSON
in one delta; update the "tool_calls" entry for id
"call_1fab8199a2fd4c23acf6fac5" so the nested "function.arguments" value is
split across multiple streamed deltas (e.g., two or more fragments that
concatenate to form "{\"city\": \"Paris\"}") to exercise incremental argument
accumulation logic; ensure the sequence order and indices remain valid so the
test harness will aggregate the fragments into the original JSON string.
In `@tests/llm/clients/fixtures/openai_multiturn_tool_roundtrip.json`:
- Around line 51-78: The fixture's _response_headers contains real identifiers
and cookies; update the JSON in openai_multiturn_tool_roundtrip.json by
replacing the values for keys like "openai-organization", "openai-project", and
"set-cookie" (and any other org/project/cookie-like headers elsewhere in the
file such as the other _response_headers block) with non-sensitive placeholders
(e.g., "REDACTED_ORG", "REDACTED_PROJECT", "REDACTED_COOKIE") so tests keep
structure but no production identifiers are committed.
In `@tests/llm/clients/record_fixtures.py`:
- Around line 70-77: The parameter name aiter in the function _try_record_stream
shadows the Python builtin aiter; rename the parameter (for example to stream or
async_iter) and update all uses inside _try_record_stream (the async for loop
and any references like chunks append/save) to the new name to avoid builtin
shadowing while preserving the existing behavior.
In `@tests/llm/clients/test_sse.py`:
- Around line 195-198: Update the test_json_raises_on_invalid test to assert the
specific JSON decoding exception instead of a bare Exception: call
ServerSentEvent(data="not json").json() inside a pytest.raises that expects
json.JSONDecodeError (falling back to ValueError on older Python if needed).
Locate the test function test_json_raises_on_invalid in
tests/llm/clients/test_sse.py and replace the generic pytest.raises(Exception)
with pytest.raises(json.JSONDecodeError) (or pytest.raises(ValueError) behind a
version check or try/except import) so the test precisely verifies
ServerSentEvent.json() error behavior.
🪄 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: Enterprise
Run ID: 80ca8e45-0301-48a1-9ee5-52242a6376da
📒 Files selected for processing (52)
examples/bots/abc/config.ymlnemoguardrails/exceptions.pynemoguardrails/integrations/langchain/llm_adapter.pynemoguardrails/llm/clients/__init__.pynemoguardrails/llm/clients/_errors.pynemoguardrails/llm/clients/_sse.pynemoguardrails/llm/clients/base.pynemoguardrails/llm/clients/constants.pynemoguardrails/llm/clients/openai_chat_model.pynemoguardrails/llm/clients/openai_compatible.pynemoguardrails/llm/default_framework.pynemoguardrails/llm/frameworks.pynemoguardrails/types.pytests/conftest.pytests/integrations/langchain/conftest.pytests/llm/clients/__init__.pytests/llm/clients/_helpers.pytests/llm/clients/fixtures/nim_generate_reasoning.jsontests/llm/clients/fixtures/nim_generate_text.jsontests/llm/clients/fixtures/nim_generate_tool_call.jsontests/llm/clients/fixtures/nim_multiturn_tool_roundtrip.jsontests/llm/clients/fixtures/nim_stream_reasoning.jsontests/llm/clients/fixtures/nim_stream_text.jsontests/llm/clients/fixtures/nim_stream_tool_calls.jsontests/llm/clients/fixtures/openai_error_400_context_length.jsontests/llm/clients/fixtures/openai_error_401.jsontests/llm/clients/fixtures/openai_generate_finish_length.jsontests/llm/clients/fixtures/openai_generate_multimodal.jsontests/llm/clients/fixtures/openai_generate_refusal.jsontests/llm/clients/fixtures/openai_generate_text.jsontests/llm/clients/fixtures/openai_generate_tool_call.jsontests/llm/clients/fixtures/openai_multiturn_tool_roundtrip.jsontests/llm/clients/fixtures/openai_stream_multimodal.jsontests/llm/clients/fixtures/openai_stream_text.jsontests/llm/clients/fixtures/openai_stream_tool_calls.jsontests/llm/clients/record_fixtures.pytests/llm/clients/test_client_config.pytests/llm/clients/test_openai_chat_model.pytests/llm/clients/test_openai_compatible.pytests/llm/clients/test_openai_compatible_live.pytests/llm/clients/test_sse.pytests/llm/clients/test_stream_llm_call.pytests/llm/test_frameworks.pytests/test_configs/with_custom_chat_model/config.pytests/test_configs/with_custom_chat_model/custom_chat_model.pytests/test_configs/with_custom_llm/config.pytests/test_configs/with_custom_llm/custom_llm.pytests/test_configs/with_custom_llm_prompt_action_v2_x/actions.pytests/test_llm_params_e2e.pytests/test_supported_llm_providers.pytests/test_task_specific_model.pytests/test_types.py
actionable error for stream_usagte
b456289 to
46c4d0b
Compare
Strict layering: OpenAICompatibleClient pure transport, no provider concept. OpenAIChatModel owns engine identity, enriches client errors at method boundaries. - Drop URL-substring provider_name heuristic from OpenAICompatibleClient - Resolve provider_name in OpenAIChatModel: explicit override > _KNOWN_PROVIDER_URLS lookup > "openai" default - Wrap generate_async/stream_async in try/except LLMClientError; _enrich(exc) stamps provider_name, model_name, base_url - DefaultFramework pass provider_name verbatim, no canonicalization Tests split by layer: TestClientErrorMetadata locks transport-only contract (no provider_name on client-direct errors); TestErrorEnrichment covers model-layer enrichment via mocked client + real-client integration.
Strict transport/model layering follow-up to b8470c2a3. BaseClient._error_context no longer mines payload["model"] for model_name on error context. Transport stamps base_url only; model_name and provider_name added by model layer's _enrich at OpenAIChatModel boundary. Removes last upper-layer leak from BaseClient. Direct client calls raise with base_url only; via OpenAIChatModel get full enrichment. TestClientErrorMetadata assertions flip from model_name == "gpt-4o" to model_name is None to lock new contract.
|
Thanks @tgasser-nv for the review :+1 I know this was a painful one given the size.
Agreed. I went with self-contained commits hoping that would make the review easier/more managble and also my life easier. Excluding the ~4700 fixture related json files, considering AI code reviewers and your experience I totally agree that splitting it would have been the right call :+1 .
we're dropping langchain as a hard dependency on the default path.
I've done it already, no regression at either layer and we'll open the PR once this is merged :+1 |
|
@tgasser-nv three of your items are already implemented as follow-up branches, will open as PRs after this lands:
|
tgasser-nv
left a comment
There was a problem hiding this comment.
Thanks for making the updates, just a few cleanups needed before merging
Emits a UserWarning when api_key is set and base_url is http:// to a non-local host. Catches the typo case where users mean https:// but write http://. Local hosts (localhost, 127.0.0.1, ::1, *.local) are suppressed since they're the common dev case (Ollama, vLLM, llama.cpp).
Signed-off-by: Pouyan <13303554+Pouyanpi@users.noreply.github.com>
…llow unhashable query values
tuple(sorted((default_query or {}).items())) crashed with TypeError when query
values were lists or dicts (legitimate inputs since signature is Dict[str, Any]
and HTTP query params are commonly multi-valued). Switching to
json.dumps(..., sort_keys=True, default=str) accepts arbitrary nested values,
preserves hashability, and stays deterministic across processes.
Same change applied to default_headers for symmetry.
eb2e130 to
a5ce8e7
Compare
|
#1828 #1829 #1830 |
Description
Adds a native, HTTP client for OpenAI compatible endpoints (OpenAI, NVIDIA NIM, vLLM, local Ollama), a chat-model adapter implementing the
LLMModelprotocol, andDefaultFrameworkthat wires this as the default LLM framework in place of LangChain.End state:
DefaultFrameworkis now the default; LangChain stays reachable viaNEMOGUARDRAILS_LLM_FRAMEWORK=langchainHow to review
The 7 commits are self-contained and build on each other. Review in order, top-down
Note:
+4317 additions belongs to a6a3fd5 were fixture files are created (tests/llm/clients/fixtures/*.json)
Testing
Client subset:
Opt-in live tests (real API calls):
Summary by CodeRabbit
New Features
Bug Fixes & Improvements
Configuration
gpt-3.5-turbo-instructtogpt-4.