[SDK][Python] Add Chroma retrieval integration#5426
Conversation
github-actions
bot
added
python
📋 PR Linter Failed❌ Invalid Title Format. Your PR title must include a ticket/issue number and may optionally include component tags (
Example: ❌ Incomplete Issues Section. You must reference at least one GitHub issue ( |
1 similar comment
📋 PR Linter Failed❌ Invalid Title Format. Your PR title must include a ticket/issue number and may optionally include component tags (
Example: ❌ Incomplete Issues Section. You must reference at least one GitHub issue ( |
github-actions
bot
added
documentation
📋 PR Linter Failed❌ Invalid Title Format. Your PR title must include a ticket/issue number and may optionally include component tags (
Example: ❌ Incomplete Issues Section. You must reference at least one GitHub issue ( |
1 similar comment
📋 PR Linter Failed❌ Invalid Title Format. Your PR title must include a ticket/issue number and may optionally include component tags (
Example: ❌ Incomplete Issues Section. You must reference at least one GitHub issue ( |
![baz-reviewer[bot]](https://avatars.githubusercontent.com/in/933528?s=60&v=4){kind=small}
| from opik.integrations._retrieval_tracker import ( | ||
| RetrievalTrackingConfig, | ||
| as_tuple, | ||
| patch_retrieval_client, | ||
| ) |
There was a problem hiding this comment.
sdks/python/AGENTS.md says new code should prefer module-style imports; this change directly imports RetrievalTrackingConfig/as_tuple/patch_retrieval_client, so the file no longer follows that guideline—can we import _retrieval_tracker as a module and reference retrieval_tracker.RetrievalTrackingConfig/patch_retrieval_client instead so the style matches the documented convention?
Finding type: AI Coding Guidelines
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents:
In sdks/python/src/opik/integrations/chroma/opik_tracker.py around lines 5 to 9 (and
usages in lines ~12 to ~24) the file currently does `from
opik.integrations._retrieval_tracker import RetrievalTrackingConfig, as_tuple,
patch_retrieval_client`, which violates the module-style import guideline. Change the
import to `import opik.integrations._retrieval_tracker as retrieval_tracker` and update
all references to use module-qualified names: `RetrievalTrackingConfig` ->
`retrieval_tracker.RetrievalTrackingConfig`, `as_tuple` -> `retrieval_tracker.as_tuple`,
and `patch_retrieval_client` -> `retrieval_tracker.patch_retrieval_client`. Ensure the
`track_chroma` function is updated accordingly and run tests/lint to confirm style
compliance.
| def track_chroma(chroma_client_or_collection: Any, project_name: Optional[str] = None) -> Any: | ||
| """Adds Opik tracking wrappers to a Chroma client or collection object.""" | ||
| config = RetrievalTrackingConfig( | ||
| provider="chroma", | ||
| operation_paths=as_tuple( | ||
| [ | ||
| "query", | ||
| "get", | ||
| ] | ||
| ), | ||
| project_name=project_name, | ||
| ) | ||
| return patch_retrieval_client(chroma_client_or_collection, config) |
There was a problem hiding this comment.
track_chroma advertises wrapping a Chroma client or collection, but only looks for top-level query/get methods; chromadb.PersistentClient exposes those on its collections, not the client, so passing a client leaves nothing patched and no spans are emitted—can we either walk into the collection paths or clarify that only collection objects are supported?
Finding type: Logical Bugs
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents:
In sdks/python/src/opik/integrations/chroma/opik_tracker.py around lines 12 to 24, the
track_chroma function claims to wrap a Chroma client or collection but only patches
top-level query/get methods, so passing a chromadb.PersistentClient results in no
instrumentation because those methods live on collections. Change the function to detect
when a client object (vs a collection) is passed: if it's a client, iterate and patch
any existing collections returned by the client (e.g., via list_collections or similar),
and wrap the client's collection factory methods (e.g.,
get_collection/create_collection) so they return patched collections; if it's already a
collection, keep the current behavior. Ensure the docstring/type hints reflect this
behavior and add minimal tests or comments documenting the supported client->collection
wrapping strategy.
| return Collection() | ||
|
|
||
|
|
||
| def test_track_chroma__patches_methods_with_retrieval_metadata(monkeypatch: Any) -> None: |
There was a problem hiding this comment.
test_track_chroma__patches_methods_with_retrieval_metadata lacks the required test_WHAT__CASE_DESCRIPTION__EXPECTED_RESULT suffix (happy path tests must end with __happyflow or similar explicit expected result per .agents/skills/python-sdk/testing.md), so it violates the SDK test naming standard—can we rename it to include the expected-result segment (e.g., __then_tracks_retrieval_metadata or __happyflow)?
Finding type: AI Coding Guidelines
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents:
In sdks/python/tests/unit/integrations/test_chroma_tracker.py around line 42, the test
function test_track_chroma__patches_methods_with_retrieval_metadata does not follow the
SDK test naming standard requiring an expected-result suffix. Rename the function
definition to include an explicit expected-result suffix, for example change its name to
test_track_chroma__patches_methods_with_retrieval_metadata__then_tracks_retrieval_metadata
(or append __happyflow). Update only the function name at the definition (and any direct
references if present) so pytest continues to discover the test with the new compliant
name.
| assert "opik.operation" in call["metadata"] | ||
|
|
||
|
|
||
| def test_track_chroma__idempotent(monkeypatch: Any) -> None: |
There was a problem hiding this comment.
test_track_chroma__idempotent also skips the EXPECTED_RESULT suffix required by test_WHAT__CASE_DESCRIPTION__EXPECTED_RESULT (per .agents/skills/python-sdk/testing.md), so the naming rule is violated—can we rename it to include an explicit outcome segment such as __then_is_idempotent?
Finding type: AI Coding Guidelines
Want Baz to fix this for you? Activate Fixer
1 participant
Details
Adds first-pass Python retrieval integration wrapper for Chroma.
This PR includes:
track_chroma__init__.py)Change checklist
Issues
Testing
cd sdks/python && PYTHONPATH=src pytest tests/unit/integrations/test_chroma_tracker.py2 passedDocumentation
docs/tracing/integrations/chroma.mdxdocs.yml