fix(tracing): Add association property #3524
Conversation
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughAdds a Gemini-based chatbot sample with tool-enabled function calling and observability, and introduces an associations API in the Traceloop SDK to attach CUSTOMER_ID, USER_ID, and SESSION_ID into OpenTelemetry context and span attributes. Changes
Sequence Diagram(s)sequenceDiagram
actor User
participant Chatbot
participant "Gemini API"
participant Tools
User->>Chatbot: process_message(session_id, user_message, history)
Chatbot->>Gemini API: generate_content(history + user_message, tools)
alt Gemini requests a tool
Gemini API-->>Chatbot: function_call (name, args)
Chatbot->>Tools: execute_function(name, args)
Tools-->>Chatbot: result
Chatbot->>Chatbot: append function_call & result to history
Chatbot->>Gemini API: generate_content(history, tools)
Gemini API-->>Chatbot: (repeat until text)
end
alt Gemini returns assistant text
Gemini API-->>Chatbot: assistant_text
Chatbot->>Chatbot: append assistant_text to history
Chatbot-->>User: return (assistant_text, updated_history)
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~30 minutes
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: defaults Review profile: CHILL Plan: Pro Disabled knowledge base sources:
📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (5)
Comment |
![ellipsis-dev[bot]](https://avatars.githubusercontent.com/in/64358?s=60&v=4){kind=small}
There was a problem hiding this comment.
Important
Looks good to me! 👍
Reviewed everything up to 8c2987c in 1 minute and 52 seconds. Click for details.
- Reviewed
538lines of code in7files - Skipped
0files when reviewing. - Skipped posting
5draft comments. View those below. - Modify your settings and rules to customize what types of comments Ellipsis leaves. And don't forget to react with 👍 or 👎 to teach Ellipsis.
1. packages/sample-app/sample_app/chats/gemini_chatbot.py:129
- Draft comment:
Using associations.set to tag the conversation is good. Consider adding input validation for the conversation_id to avoid empty or malformed values. - Reason this comment was not posted:
Decided after close inspection that this draft comment was likely wrong and/or not actionable: usefulness confidence = 20% vs. threshold = 50% This is a new file being added, so all the code is technically new. The comment is suggesting defensive programming by adding validation. However, looking at the actual usage inmain(), theconversation_idis always generated as a UUID string, so it will never be empty or malformed in practice. The comment is speculative - it's suggesting to guard against a problem that doesn't exist in the current code. The function could theoretically be called from elsewhere with bad input, but that's speculative. According to the rules, I should not keep speculative comments like "If X, then Y is an issue" - I should only comment if it's definitely an issue. There's no evidence of actual misuse here. The comment could be valid if this function is part of a public API that could be called from other places with arbitrary input. I might be missing context about whether this function is intended to be called from other modules or if it's only used internally within this file. The validation could prevent runtime errors if someone calls this function incorrectly. Even if the function could be called from elsewhere, the comment is still speculative because there's no evidence in the current code that it will be called with invalid input. The only caller in this file (main()) always passes a valid UUID. The rules explicitly state not to make speculative comments. This is a "consider adding" suggestion for defensive programming, not a clear bug or issue. This comment should be deleted. It's a speculative suggestion for defensive programming without evidence of an actual issue. The only usage in the code always passes a valid UUID string, and there's no indication that invalid input will be passed. This violates the rule against speculative comments.
2. packages/traceloop-sdk/tests/test_associations.py:60
- Draft comment:
The tests correctly set and assert associations. Note that span order is assumed; consider verifying spans based on attributes rather than fixed order for increased robustness. - Reason this comment was not posted:
Confidence changes required:50%<= threshold50%None
3. packages/traceloop-sdk/traceloop/sdk/associations/associations.py:44
- Draft comment:
Ensure get_value returns a dict before updating associations to prevent potential type errors. Consider adding a type check or fallback logic. - Reason this comment was not posted:
Comment looked like it was already resolved.
4. packages/traceloop-sdk/traceloop/sdk/client/client.py:72
- Draft comment:
Initialization of associations in the client constructor is straightforward and ensures associations are available in subsequent operations. - Reason this comment was not posted:
Confidence changes required:0%<= threshold50%None
5. packages/traceloop-sdk/traceloop/sdk/tracing/tracing.py:352
- Draft comment:
Associations are correctly applied to spans in default_span_processor_on_start. For extra safety, consider prefixing association keys to avoid any collision with other span attributes. - Reason this comment was not posted:
Decided after close inspection that this draft comment was likely wrong and/or not actionable: usefulness confidence = 20% vs. threshold = 50% This is a code quality suggestion about the new associations feature. The comment points out that the keys from the associations dictionary are being used directly as span attribute names without any namespace/prefix. This could indeed cause collisions if the associations dictionary contains keys that match existing span attribute names. Looking at the pattern in the codebase, other similar features do use prefixes (association_properties, prompt_template_variables). However, I need to consider: 1) Is this definitely about changed code? Yes, lines 352-357 are new additions. 2) Is this actionable? Yes, it's a clear suggestion to add prefixing. 3) Is this speculative or definite? It's somewhat speculative - it says "consider" and talks about avoiding "any collision" which is a potential issue, not a definite one. The comment doesn't prove there IS a collision, just that there COULD be one. The comment is somewhat speculative - it uses "consider" and talks about avoiding potential collisions rather than pointing to a definite issue. Without knowing what keys are actually in the associations dictionary or what the intended use case is, we can't be certain this is a problem. Perhaps the associations feature is intentionally designed to allow setting arbitrary span attributes without prefixes. While the comment is somewhat speculative, it's a reasonable code quality concern based on the established patterns in the codebase. However, the rules state "Do NOT make speculative comments, such as 'If X, then Y is an issue', Only comment if it is definitely an issue." The word "consider" and the phrase "to avoid any collision" make this a speculative suggestion rather than a definite issue. Without more context about the intended design of the associations feature, we cannot be certain this is incorrect. This comment should be deleted. While it raises a reasonable code quality concern, it is speculative in nature ("consider...to avoid any collision") rather than pointing to a definite issue. The rules explicitly state not to make speculative comments. Without evidence that collisions will actually occur or that this design is incorrect, we should assume the author implemented it correctly.
Workflow ID: wflow_ZvxkHXwf5pprlvcm
You can customize 
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
packages/traceloop-sdk/traceloop/sdk/__init__.py (1)
34-34: Re-export import flagged as unused by Flake8.The import is intentionally a re-export for public API access, but Flake8 reports it as unused (F401). To silence the linter and make the intent explicit, consider adding an
__all__declaration or using an explicit re-export pattern.🔎 Proposed fix
from traceloop.sdk.client.client import Client -from traceloop.sdk.associations.associations import AssociationProperty +from traceloop.sdk.associations.associations import AssociationProperty # noqa: F401Alternatively, define
__all__at module level to explicitly declare public exports.packages/traceloop-sdk/traceloop/sdk/associations/__init__.py (1)
1-7: LGTM! Clean package re-export structure.The file correctly exposes the public API for the associations module. The
__all__list is functional but could be sorted alphabetically per Ruff's RUF022 recommendation.🔎 Optional: Sort __all__ alphabetically
-__all__ = ["Associations", "AssociationProperty", "Association"] +__all__ = ["Association", "AssociationProperty", "Associations"]packages/sample-app/sample_app/chats/gemini_chatbot.py (1)
221-223: Consider narrowing the exception type.Catching bare
Exception(Ruff BLE001) can mask unexpected errors. For a sample app this may be acceptable, but consider catching more specific exceptions likegoogle.genai.errors.APIErroror similar.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Disabled knowledge base sources:
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (7)
packages/sample-app/sample_app/chats/gemini_chatbot.py(1 hunks)packages/traceloop-sdk/tests/test_associations.py(1 hunks)packages/traceloop-sdk/traceloop/sdk/__init__.py(1 hunks)packages/traceloop-sdk/traceloop/sdk/associations/__init__.py(1 hunks)packages/traceloop-sdk/traceloop/sdk/associations/associations.py(1 hunks)packages/traceloop-sdk/traceloop/sdk/client/client.py(3 hunks)packages/traceloop-sdk/traceloop/sdk/tracing/tracing.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.py: Store API keys only in environment variables/secure vaults; never hardcode secrets in code
Use Flake8 for code linting and adhere to its rules
Files:
packages/traceloop-sdk/traceloop/sdk/associations/__init__.pypackages/traceloop-sdk/traceloop/sdk/client/client.pypackages/traceloop-sdk/tests/test_associations.pypackages/traceloop-sdk/traceloop/sdk/__init__.pypackages/traceloop-sdk/traceloop/sdk/associations/associations.pypackages/sample-app/sample_app/chats/gemini_chatbot.pypackages/traceloop-sdk/traceloop/sdk/tracing/tracing.py
🧠 Learnings (1)
📚 Learning: 2025-08-04T15:23:44.799Z
Learnt from: nina-kollman
Repo: traceloop/openllmetry PR: 3219
File: packages/traceloop-sdk/tests/datasets/mock_objects.py:105-106
Timestamp: 2025-08-04T15:23:44.799Z
Learning: In the traceloop-sdk project, directly accessing private attributes like _client is acceptable in test mock objects and utilities, even though it breaks encapsulation principles, because it's isolated to test code and not production code.
Applied to files:
packages/traceloop-sdk/traceloop/sdk/__init__.py
🧬 Code graph analysis (4)
packages/traceloop-sdk/traceloop/sdk/associations/__init__.py (1)
packages/traceloop-sdk/traceloop/sdk/associations/associations.py (2)
Associations(22-54)AssociationProperty(9-15)
packages/traceloop-sdk/traceloop/sdk/client/client.py (1)
packages/traceloop-sdk/traceloop/sdk/associations/associations.py (1)
Associations(22-54)
packages/traceloop-sdk/tests/test_associations.py (3)
packages/traceloop-sdk/traceloop/sdk/__init__.py (2)
Traceloop(37-268)init(49-199)packages/traceloop-sdk/traceloop/sdk/associations/associations.py (1)
AssociationProperty(9-15)packages/traceloop-sdk/traceloop/sdk/utils/in_memory_span_exporter.py (4)
export(45-51)InMemorySpanExporter(22-61)clear(35-38)get_finished_spans(40-43)
packages/traceloop-sdk/traceloop/sdk/__init__.py (1)
packages/traceloop-sdk/traceloop/sdk/associations/associations.py (1)
AssociationProperty(9-15)
🪛 Flake8 (7.3.0)
packages/traceloop-sdk/traceloop/sdk/__init__.py
[error] 34-34: 'traceloop.sdk.associations.associations.AssociationProperty' imported but unused
(F401)
🪛 Ruff (0.14.8)
packages/traceloop-sdk/traceloop/sdk/associations/__init__.py
7-7: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
packages/sample-app/sample_app/chats/gemini_chatbot.py
221-221: Do not catch blind exception: Exception
(BLE001)
🔇 Additional comments (9)
packages/traceloop-sdk/traceloop/sdk/tracing/tracing.py (1)
352-356: LGTM! Association enrichment follows existing patterns.The code correctly reads associations from context and applies them as span attributes, consistent with the handling of other context values (workflow_name, agent_name, entity_path). The stringification with
str(value)ensures attribute values are safe for OpenTelemetry.packages/traceloop-sdk/traceloop/sdk/client/client.py (1)
9-9: LGTM! Clean integration of Associations into Client.The
Associationsattribute follows the existing pattern established by other client features (user_feedback, datasets, etc.), providing a consistent API surface for users to access viaclient.associations.set(...).Also applies to: 30-30, 72-72
packages/traceloop-sdk/tests/test_associations.py (4)
48-73: Test coverage for single association is well-structured.The test properly validates that associations set before workflow execution propagate to both task and workflow spans.
76-111: Good coverage for multiple associations.The test validates that all three associations propagate correctly to both span types.
114-146: Tests associations set within workflow context.This validates the important use case of setting associations mid-workflow.
149-171: Comprehensive test for all enum values.Good validation that all
AssociationPropertyenum members work correctly.packages/traceloop-sdk/traceloop/sdk/associations/associations.py (1)
22-53: Well-structured Associations class with clear API.The implementation correctly:
- Merges new associations with existing context values
- Attaches updated associations to OpenTelemetry context
- Sets attributes on the current span for immediate visibility
The docstring provides helpful usage examples.
packages/sample-app/sample_app/chats/gemini_chatbot.py (2)
53-108: Tool declarations are well-structured.The function declarations follow the Gemini API pattern correctly with proper schema definitions.
124-187: Workflow implementation demonstrates associations API effectively.The
process_messagefunction correctly uses the new associations API to attach conversation IDs to spans, providing good observability for multi-turn conversations.
There was a problem hiding this comment.
Caution
Changes requested ❌
Reviewed 1af680b in 2 minutes and 4 seconds. Click for details.
- Reviewed
103lines of code in3files - Skipped
0files when reviewing. - Skipped posting
2draft comments. View those below. - Modify your settings and rules to customize what types of comments Ellipsis leaves. And don't forget to react with 👍 or 👎 to teach Ellipsis.
1. packages/sample-app/sample_app/chats/gemini_chatbot.py:128
- Draft comment:
Inline comment is outdated—update 'conversation_id' to 'session_id' to match the code change. - Reason this comment was not posted:
Comment was on unchanged code.
2. packages/sample-app/sample_app/chats/gemini_chatbot.py:125
- Draft comment:
Update the function docstring to reference 'session_id' instead of the outdated 'chat_id'. - Reason this comment was not posted:
Comment was on unchanged code.
Workflow ID: wflow_oIcBnmvFwCTrKM8n
You can customize
| spans = exporter.get_finished_spans() | ||
| workflow_span = spans[0] | ||
|
|
||
| assert workflow_span.attributes["session_id"] == "conv-1" |
There was a problem hiding this comment.
The test sets two SESSION_ID associations ("conv-1" then "session-4"), but the implementation overwrites duplicates. Adjust the assertions to expect only the final value.
| assert workflow_span.attributes["session_id"] == "conv-1" |
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (2)
packages/sample-app/sample_app/chats/gemini_chatbot.py (2)
11-11: Handle potential None from Traceloop.init().This issue was previously raised and remains unresolved. Line 129 will raise
AttributeErroriftraceloopisNone.
150-176: Add bounds checking for response access.This issue was previously raised and remains unresolved. Direct indexing of
response.candidates[0].content.parts[0]can raiseIndexErrorfor empty or malformed responses.
🧹 Nitpick comments (3)
packages/traceloop-sdk/traceloop/sdk/associations/associations.py (1)
24-53: Consider adding input validation for association values.The method doesn't validate that association values are non-empty. While OpenTelemetry context is typically thread-local (mitigating concurrency concerns), adding a simple check for empty strings could improve robustness.
🔎 Optional: Add validation for empty values
@staticmethod def set(associations: Sequence[Association]) -> None: """ Set associations that will be added directly to all spans in the current context. Args: associations: A sequence of (property, value) tuples Example: # Single association traceloop.associations.set([(AssociationProperty.SESSION_ID, "conv-123")]) # Multiple associations traceloop.associations.set([ (AssociationProperty.USER_ID, "user-456"), (AssociationProperty.SESSION_ID, "session-789") ]) """ + # Validate inputs + for prop, value in associations: + if not value or not value.strip(): + raise ValueError(f"Association value for {prop.value} cannot be empty") + # Store all associations in context current_associations: dict[str, str] = get_value(ASSOCIATIONS_KEY) or {} # type: ignore for prop, value in associations: current_associations[prop.value] = value attach(set_value(ASSOCIATIONS_KEY, current_associations)) # Also set directly on the current span span = trace.get_current_span() if span and span.is_recording(): for prop, value in associations: span.set_attribute(prop.value, value)packages/sample-app/sample_app/chats/gemini_chatbot.py (2)
138-187: Add error handling and loop safeguards to prevent infinite loops.The
while Trueloop lacks error handling for API failures and has no maximum iteration limit. If the Gemini API consistently returns malformed responses or the model gets stuck in a tool-calling loop, this could run indefinitely.🔎 Proposed safeguards
# Keep trying until we get a final response (handle tool calls) + max_iterations = 10 + iteration = 0 - while True: + while iteration < max_iterations: + iteration += 1 + # Generate content with tools - response = client.models.generate_content( - model="gemini-2.0-flash-exp", - contents=conversation_history, - config=types.GenerateContentConfig( - tools=[weather_tool, time_tool, knowledge_tool], - temperature=0.7, - ) - ) + try: + response = client.models.generate_content( + model="gemini-2.0-flash-exp", + contents=conversation_history, + config=types.GenerateContentConfig( + tools=[weather_tool, time_tool, knowledge_tool], + temperature=0.7, + ) + ) + except Exception as e: + return f"Error generating response: {e}", conversation_history # Check if the model wants to use a tool - if response.candidates[0].content.parts[0].function_call: + if (response.candidates and + response.candidates[0].content.parts and + response.candidates[0].content.parts[0].function_call): function_call = response.candidates[0].content.parts[0].function_call function_name = function_call.name function_args = dict(function_call.args) print(f"[Tool Call]: {function_name}({function_args})") # Execute the function - function_result = execute_function(function_name, function_args) + try: + function_result = execute_function(function_name, function_args) + except Exception as e: + function_result = f"Error executing function: {e}" + print(f"[Tool Result]: {function_result}") # Add the model's function call to history conversation_history.append({ "role": "model", "parts": [{"function_call": function_call}] }) # Add the function result to history conversation_history.append({ "role": "user", "parts": [{ "function_response": types.FunctionResponse( name=function_name, response={"result": function_result} ) }] }) else: # Got a text response, we're done with this turn assistant_message = response.text # Add assistant response to conversation history conversation_history.append({ "role": "model", "parts": [{"text": assistant_message}] }) return assistant_message, conversation_history + + # Max iterations reached + return "Maximum conversation turns reached. Please try rephrasing your question.", conversation_history
216-223: Consider more specific exception handling.The broad
Exceptioncatch (flagged by Ruff BLE001) is acceptable for a sample application to maintain conversation flow, but consider catching specific exceptions (e.g.,genai.APIError,ValueError) separately for better error diagnostics.🔎 Optional: More specific exception handling
# Process the message try: assistant_message, conversation_history = process_message( chat_id, user_message, conversation_history ) print(f"\nAssistant: {assistant_message}") + except (KeyError, ValueError, AttributeError) as e: + print(f"\nConfiguration Error: {e}") + print("Please check your setup and try again.") + break except Exception as e: print(f"\nError: {e}") print("Please try again.")
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Disabled knowledge base sources:
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (3)
packages/sample-app/sample_app/chats/gemini_chatbot.py(1 hunks)packages/traceloop-sdk/tests/test_associations.py(1 hunks)packages/traceloop-sdk/traceloop/sdk/associations/associations.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- packages/traceloop-sdk/tests/test_associations.py
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.py: Store API keys only in environment variables/secure vaults; never hardcode secrets in code
Use Flake8 for code linting and adhere to its rules
Files:
packages/sample-app/sample_app/chats/gemini_chatbot.pypackages/traceloop-sdk/traceloop/sdk/associations/associations.py
🪛 Ruff (0.14.8)
packages/sample-app/sample_app/chats/gemini_chatbot.py
221-221: Do not catch blind exception: Exception
(BLE001)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (5)
- GitHub Check: Test Packages (3.11)
- GitHub Check: Test Packages (3.12)
- GitHub Check: Test Packages (3.10)
- GitHub Check: Build Packages (3.11)
- GitHub Check: Lint
🔇 Additional comments (8)
packages/traceloop-sdk/traceloop/sdk/associations/associations.py (3)
1-6: LGTM!Imports are appropriate and the constant follows Python naming conventions.
9-14: LGTM!The
str, Enuminheritance pattern enables direct string comparison and JSON serialization, which is appropriate for span attributes.
17-18: LGTM!Clean type alias definition using modern Python syntax.
packages/sample-app/sample_app/chats/gemini_chatbot.py (5)
1-8: LGTM!All necessary imports are present and properly organized.
17-49: LGTM!The tool functions are appropriately simplified for demonstration purposes. The implementations use mock data and basic logic suitable for a sample application.
52-108: LGTM!Tool declarations are correctly structured with appropriate schemas, descriptions, and required parameters matching the function implementations.
111-121: LGTM!The function routing logic is clear and handles unknown functions gracefully by returning an error message.
226-227: LGTM!Standard entry point pattern correctly implemented.
There was a problem hiding this comment.
Important
Looks good to me! 👍
Reviewed 302fa80 in 25 seconds. Click for details.
- Reviewed
12lines of code in1files - Skipped
0files when reviewing. - Skipped posting
1draft comments. View those below. - Modify your settings and rules to customize what types of comments Ellipsis leaves. And don't forget to react with 👍 or 👎 to teach Ellipsis.
1. packages/traceloop-sdk/tests/test_associations.py:165
- Draft comment:
The test now checks that duplicate SESSION_ID values are overwritten (last wins). Please ensure that this behavior is documented in the associations API. - Reason this comment was not posted:
Comment did not seem useful. Confidence is useful =0%<= threshold50%This comment is asking the PR author to ensure documentation, which is not allowed by the rules. It doesn't provide a specific code suggestion or ask for a specific test to be written.
Workflow ID: wflow_zkbIX8xi7GqSaRsv
You can customize
feat(instrumentation): ...orfix(instrumentation): ....Important
Added a Gemini-based chatbot and an associations API for enhanced tracing capabilities, with tests for association propagation.
gemini_chatbot.pyfor a Gemini-based chatbot with weather, time, and knowledge-base query support.AssociationsandAssociationPropertyinassociations.pyfor managing trace associations.tracing.pyto handle associations in spans usingAssociations.Clientinclient.pyto includeassociationsattribute.test_associations.pyto validate association propagation in spans.__init__.pyto importAssociationProperty.This description was created by
for 302fa80. You can customize this summary. It will automatically update as commits are pushed.
Summary by CodeRabbit
New Features
Tests
✏️ Tip: You can customize this high-level summary in your review settings.