Repository navigation
feat: add Cerebras integration to the Python SDK - #7342
Conversation
github-actions
Bot
added
dependencies
![baz-reviewer[bot]](https://avatars.githubusercontent.com/in/933528?s=60&v=4){kind=small}
| aggregated_response = { | ||
| "choices": [{"index": 0, "message": {"role": "", "content": ""}}], | ||
| "created": first_chunk.created, | ||
| "id": first_chunk.id, | ||
| "model": first_chunk.model, |
There was a problem hiding this comment.
Duplicate chunk aggregation
aggregate() is a near copy of opik/integrations/openai/chat_completion_chunks_aggregator.py:22-63, so OpenAI and Cerebras chunk stitching will drift if this logic changes; should we extract it into a reusable helper?
Want Baz to fix this for you? Activate Fixer
| if delta.content: | ||
| text_chunks.append(delta.content) |
There was a problem hiding this comment.
Drops streamed reasoning output
delta.reasoning is dropped when the reducer aggregates streamed choices, so reasoning-enabled completions log a synthetic choices[0].message without the reasoning text and the span output is incomplete — should we preserve it too, as Cerebras’ official reasoning docs say it’s delivered in delta.reasoning?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/src/opik/integrations/cerebras/chat_completion_chunks_aggregator.py around
lines 48-49 inside the `aggregate` function’s loop, the code currently appends only
`delta.content` into `text_chunks` and never copies `delta.reasoning` from the streamed
chunks. Refactor the aggregation to also collect `delta.reasoning` tokens (e.g., a
separate buffer) and set them onto the synthetic
`aggregated_response['choices'][0]['message']` in the appropriate field so downstream
span/log output includes reasoning-enabled text. Ensure this works when reasoning is
present across multiple streamed chunks by concatenating all reasoning tokens in order,
similar to how `content` is concatenated.
| def Stream__iter__decorator(dunder_iter_func: Callable) -> Callable: | ||
| @functools.wraps(dunder_iter_func) | ||
| def wrapper(self: cerebras.Stream) -> Iterator[StreamItem]: | ||
| try: | ||
| accumulated_items: List[StreamItem] = [] |
There was a problem hiding this comment.
Duplicate stream patching
patch_sync_stream/patch_async_stream duplicate the opik/integrations/openai/stream_patchers.py:30‑157 logic almost verbatim for cerebras.Stream, so any fix to generations_aggregator or finally_callback has to be mirrored in both places and can drift — should we parameterize the shared patcher instead?
Want Baz to fix this for you? Activate Fixer
| LOGGER.debug( | ||
| "Exception raised from cerebras.Stream.", | ||
| str(exception), | ||
| exc_info=True, | ||
| ) |
There was a problem hiding this comment.
Logging call raises formatting error
In both the sync and async stream wrappers, LOGGER.debug("Exception raised from cerebras.Stream.", str(exception), exc_info=True) passes str(exception) to a message with no %s, so logging raises a TypeError before the debug line is emitted and can interrupt exception handling while the span is still open — should we change it to LOGGER.debug("Exception raised from cerebras.Stream: %s", exception, exc_info=True)?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/src/opik/integrations/cerebras/stream_patchers.py around lines 43-47 in
patch_sync_stream’s Stream__iter__decorator exception handler, the LOGGER.debug call
passes str(exception) as a positional arg to a message string with no formatting
placeholders, which can raise TypeError during logging and prevent cleanup. Refactor
that logging line to use standard logging formatting with a “%s” placeholder (e.g.,
append “: %s” to the message and pass the exception as the corresponding argument)
while keeping exc_info=True. Do the same for the analogous LOGGER.debug call in
patch_async_stream around lines 96-100 inside AsyncStream__aiter__decorator.
| import cerebras.cloud.sdk as cerebras | ||
|
|
||
| from opik.types import LLMProvider | ||
|
|
||
| from . import cerebras_chat_completions_decorator, chat_completion_chunks_aggregator |
There was a problem hiding this comment.
Optional import crashes package load
sdks/python/src/opik/integrations/cerebras/__init__.py imports cerebras.cloud.sdk and the decorator helpers at module load, so import opik.integrations.cerebras raises ModuleNotFoundError in environments without cerebras-cloud-sdk before track_cerebras() is reached — should we move those imports behind the public function boundary or guard them like the other optional integrations, as .agents/skills/python-sdk/SKILL.md requires?
Want Baz to fix this for you? Activate Fixer You can also update your AI coding guidelines based on this comment by apply pr to [branch name]
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/src/opik/integrations/cerebras/opik_tracker.py around lines 1-70,
`track_cerebras()` relies on `import cerebras.cloud.sdk as cerebras` and module-scope
imports of the Cerebras decorator modules, which execute during package import and can
raise ModuleNotFoundError when the optional `cerebras-cloud-sdk` isn’t installed.
Refactor by removing the top-level Cerebras SDK import and any other Cerebras-runtime
imports from module scope; instead, perform them inside `track_cerebras()` (and
`_extract_metadata_from_client` if needed) with a clear error message if Cerebras
isn’t installed. Keep typing aliases (e.g., `TypeVar`) behind `typing.TYPE_CHECKING`
or adjust them so they don’t require the Cerebras SDK at import time.
|
Same as #7341 (Groq) — these three points are intentional consistency with the existing OpenAI/Anthropic integrations:
|
Add track_cerebras() to wrap Cerebras / AsyncCerebras clients so calls to chat.completions.create() (including stream=True) are logged as Opik LLM spans, following the existing OpenAI/Anthropic integration pattern. Cerebras returns OpenAI-shaped payloads, so usage is parsed with the OpenAI converter while the span provider is recorded as 'cerebras'. Adds mock-based unit tests and wires the cerebras suite into the integration runner.
729ed0d to
b89405f
Compare
|
Rebased onto latest Note for anyone comparing SHAs: a plain rebase wasn't possible ( The only conflict was in Also re-checked the decorator API against current |
|
Thanks for the contribution and apologies about the delay! We're going to proceed with this review and potentially add this Cerebras integration if once all is ready in the PR. When you have a chance, could you take a look at the automated review comments from Baz? For each comment, please either address it with a code change or reply with a brief explanation if you believe the current implementation is the right approach. Thanks! |
| env: | ||
| CEREBRAS_API_KEY: ${{ secrets.CEREBRAS_API_KEY }} |
There was a problem hiding this comment.
PR tests inherit Cerebras credential
CEREBRAS_API_KEY is exposed at workflow scope, so PR-controlled code run by pip install . or pytest can read and transmit the real secret even though the tests use a fake key and monkeypatch _post — should we remove the workflow-wide env entry or scope credentials to the specific test step that needs Cerebras access?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
`.github/workflows/lib-cerebras-tests.yml` around lines 9-10, remove the workflow-wide
`CEREBRAS_API_KEY` environment variable because it exposes the secret to package
installation and PR-controlled test code, while these tests use fake credentials and
mocked requests. If a future test genuinely requires the key, scope it only to that
specific step rather than defining it at workflow scope.
| def _extract_metadata_from_client(client: CerebrasClient) -> Dict[str, Any]: | ||
| return {"base_url": str(client.base_url)} |
There was a problem hiding this comment.
Secrets in Cerebras URL reach Opik telemetry
_extract_metadata_from_client stores str(client.base_url) verbatim, so caller-supplied userinfo, query, and fragment components reach metadata["base_url"] and the span/trace payloads — should we record only the validated host/origin or strip those components first?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
`sdks/python/src/opik/integrations/cerebras/opik_tracker.py` around lines 69-70, update
`_extract_metadata_from_client` so `metadata["base_url"]` cannot contain credentials,
query parameters, or fragments from the caller-provided Cerebras `base_url`. Parse the
URL and record only its validated scheme and host/origin (including a non-sensitive port
when needed), with safe handling for malformed or missing values before the metadata is
merged into spans and traces.
| cerebras.Stream.__iter__ = Stream__iter__decorator(original_stream_iter_method) | ||
|
|
||
| stream.opik_tracked_instance = True |
There was a problem hiding this comment.
patch_sync_stream replaces cerebras.Stream.__iter__ at the class level with per-call finally_callback and accumulated_items state, so concurrent track_cerebras() calls share state and can finalize the wrong span/trace or make untracked streams accumulate chunks; patch_async_stream has the same issue with cerebras.AsyncStream.__aiter__ — should we wrap each tracked stream instance or use a thread-safe per-stream dispatcher?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
`sdks/python/src/opik/integrations/cerebras/stream_patchers.py` around lines 68-70 and
121-125, refactor `patch_sync_stream` and `patch_async_stream` so they no longer replace
`cerebras.Stream.__iter__` or `cerebras.AsyncStream.__aiter__` at the class level with a
closure holding call-specific state. Instead, attach the iterator wrapper only to the
tracked stream instance (or use a stable per-instance/thread-safe dispatcher that reads
finalization state from that stream), ensuring: (1) concurrent `track_cerebras()` calls
finalize the correct span, trace, and callback without racing; (2) untracked streams
never accumulate chunks or otherwise change SDK behavior; and (3) successive tracked
streams don't stack wrappers around each other — while preserving the original SDK
methods and existing finalization/error handling.
| except Exception as exception: | ||
| LOGGER.debug( | ||
| "Exception raised from cerebras.AsyncStream.", | ||
| str(exception), | ||
| exc_info=True, | ||
| ) | ||
| error_info = error_info_collector.collect(exception) | ||
| raise exception | ||
| finally: | ||
| if hasattr(self, "opik_tracked_instance"): | ||
| delattr(self, "opik_tracked_instance") | ||
| output = ( |
There was a problem hiding this comment.
Interrupted streams recorded as successful
The sync and async stream wrappers catch only Exception, so cancellations and generator termination leave error_info=None; their finally blocks pass partial aggregation to _after_call, which records interrupted streams as successful output — should we capture these terminations, pass output=None, and re-raise the original interruption unchanged?
Want Baz to fix this for you? Activate Fixer You can also update your AI coding guidelines based on this comment by apply pr to [branch name]
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
`sdks/python/src/opik/integrations/cerebras/stream_patchers.py` around lines 42-64 and
95-117, update the sync and async stream wrapper exception handling so cancellation,
`GeneratorExit`, and other relevant `BaseException` interruptions are collected into
`error_info` before finalization. Ensure these paths re-raise the original termination
unchanged, while the `finally` logic passes `output=None` instead of aggregating partial
items; preserve normal `Exception` handling and successful aggregation for genuinely
completed streams.
|
Thanks for the contribution @NishchayMahor! |
…rror paths (#8238) The two LOGGER.debug calls in the Cerebras stream wrappers pass str(exception) positionally to a message with no %s. logging defers the interpolation to the handler, so with debug logging on it raises inside Handler.emit, handleError prints a "--- Logging error ---" block to stderr, and the record — including the exc_info traceback — is dropped. Same fix as #7850, which covered the twelve call sites that existed when it merged. The Cerebras integration landed afterwards in #7342, so these two were never in its scope. An AST pass over sdks/python/src/opik for LOGGER calls whose message carries fewer %-placeholders than positional args now reports none.
Summary
Adds a Cerebras integration to the Opik Python SDK, so calls made with the native
cerebras-cloud-sdkclient are automatically logged as Opik LLM spans:Cerebras (very fast inference for Llama and other open models) has a distinct native SDK, so the OpenAI wrapper (
track_openai) can't instrument it — this adds first-class support, matching the existing standalone integrations.Details
track_cerebras()follows the existing OpenAI/Anthropic pattern:chat.completions.create()on bothCerebrasandAsyncCerebras, includingstream=True(sync + async streaming via aggregated chunks)."cerebras"; Cerebras returns OpenAI-shaped usage, so token usage is parsed with the OpenAI converter (the comment notes this denotes payload format, not the span provider).error_infoon failures.Files:
src/opik/integrations/cerebras/—track_cerebras, decorator, chunk aggregator, stream patchers.tests/library_integration/cerebras/— mock-based unit tests +requirements.txt.lib-cerebras-tests.yml+ registration inlib-integration-tests-runner.yml.Testing
pytest tests/library_integration/cerebras/→ 4 passed (sync happy-path with full trace-tree assertion, async, and error handling), using thefake_backendfixture and a mocked transport so no Cerebras API key is required.ruffandmypyclean.Happy to add a documentation/cookbook page as a follow-up if preferred.