feat: add embedding tracer support for Bedrock, LiteLLM, OpenAI#631
Open
viniciusdsmello wants to merge 17 commits intomainfrom
Open
feat: add embedding tracer support for Bedrock, LiteLLM, OpenAI#631viniciusdsmello wants to merge 17 commits intomainfrom
viniciusdsmello wants to merge 17 commits intomainfrom
Conversation
Used by superpowers workflows to host isolated git worktrees during implementation, never meant to be tracked. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…EN-10480) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…480) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…N-10480) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…OPEN-10480) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…gression (OPEN-10480) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds the same file-level pragma already used by test_portkey_integration.py to suppress reportUnknown* and reportMissingParameterType — these come from openlayer.lib.integrations being in pyright's ignore list, which causes imports from there to be typed as Unknown. Per-line pyright ignores added on direct imports of botocore.response and openai, which are not present in the lint job's environment. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Code review findings addressed: - Move per-call imports of _openai_embedding_common to module-level (was in hot path of every embedding call). - Extract build_embedding_step_kwargs into _openai_embedding_common so that sync and async OpenAI handlers each become ~10 lines instead of ~50, and LiteLLM reuses the same kwargs assembly. - Drop LiteLLM's local _parse_embedding_response and _get_embedding_model_parameters; both now delegate to the shared helpers (LiteLLM-specific timeout/api_base/api_version/cost/metadata are layered on top of the common kwargs). - Type Bedrock _parse_embedding_output return as Tuple[Union[List[float], List[List[float]]], int, int] instead of bare tuple. Net: -34 lines across the 5 touched source files. Tests unchanged, all 77 embedding tests + 448 lib tests still green. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds native embedding tracing support across the Python SDK so that embedding API calls (Titan via Bedrock,
litellm.embedding,OpenAI.embeddings.create) generate proper traces in Openlayer with correct model, tokens, dimensions, and output.StepType.EMBEDDING+add_embedding_step_to_tracehelper (src/openlayer/lib/tracing/)."embed"inmodelIdand routes to a dedicated handler with parsers for Titan v1/v2 and Cohere v3 (single + batch). Existing chat path is untouched and locked in by a backfilled regression test.litellm.embeddingalongside the existinglitellm.completionpatch. Reusesdetect_provider_from_response,extract_usage_from_response, andextract_litellm_metadata.client.embeddings.createfor both sync (trace_openai) and async (trace_async_openai) clients via a small shared helper module (_openai_embedding_common.py).Linear
OPEN-10480
Verification
ruff checkclean on all touched files.pyrightclean on all touched source files.Test plan
amazon.titan-embed-text-v2:0) — confirm trace appears with model name and prompt tokens populated.litellm.embedding(model="text-embedding-3-small", input="x").step_type=embeddingis rendered correctly (out of scope for this PR but required for end-to-end value).Out of scope
🤖 Generated with Claude Code