feat(wren-ai-service): Add trace ID propagation across service responses #1364
Conversation
paopa
added
module/ai-service
|
Warning Rate limit exceeded@paopa has exceeded the limit for the number of commits or files that can be reviewed per hour. Please wait 26 minutes and 44 seconds before requesting another review. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. 📒 Files selected for processing (2)
WalkthroughThe changes add a new optional Changes
Sequence Diagram(s)sequenceDiagram
participant C as Client
participant R as Router/Endpoint
participant D as trace_metadata Decorator
participant S as Service
participant M as Response Model
C->>R: Send request
R->>D: Invoke decorated function
D->>D: Retrieve trace_id via langfuse_context
D->>S: Call service with trace_id passed in kwargs
S->>M: Build response including trace_id
M-->>S: Return response object
S-->>D: Return result with trace_id
D-->>R: Pass result with trace_id
R-->>C: Respond with trace_id included
Suggested reviewers
Poem
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🔭 Outside diff range comments (9)
wren-ai-service/src/web/v1/services/sql_question.py (1)
131-137: 🛠️ Refactor suggestionMissing trace_id propagation in error handling case
The error handling in
get_sql_question_resultdoesn't include trace_id when returning a "not found" error response. For consistency, trace_id should be included here as well.return SqlQuestionResultResponse( status="failed", error=SqlQuestionResultResponse.SqlQuestionError( code="OTHERS", message=f"{sql_question_result_request.query_id} is not found", ), + trace_id=None, )wren-ai-service/src/web/v1/services/sql_expansion.py (2)
262-269: 🛠️ Refactor suggestionMissing trace_id propagation in stop_sql_expansion method
The
stop_sql_expansionmethod doesn't include trace_id when setting the status to "stopped". For consistency with the rest of the implementation, trace_id should be included here as well.def stop_sql_expansion( self, stop_sql_expansion_request: StopSqlExpansionRequest, ): self._sql_expansion_results[ stop_sql_expansion_request.query_id - ] = SqlExpansionResultResponse(status="stopped") + ] = SqlExpansionResultResponse(status="stopped", trace_id=None)
270-288: 🛠️ Refactor suggestionMissing trace_id propagation in error handling case
The error handling in
get_sql_expansion_resultdoesn't include trace_id when returning a "not found" error response. For consistency with the rest of the implementation, trace_id should be included here as well.return SqlExpansionResultResponse( status="failed", error=AskError( code="OTHERS", message=f"{sql_expansion_result_request.query_id} is not found", ), + trace_id=None, )wren-ai-service/src/web/v1/services/ask_details.py (1)
181-197: 🛠️ Refactor suggestionMissing trace_id propagation in error handling case
The error handling in
get_ask_details_resultdoesn't include trace_id when returning a "not found" error response. For consistency with the rest of the implementation, trace_id should be included here as well.return AskDetailsResultResponse( status="failed", error=AskDetailsResultResponse.AskDetailsError( code="OTHERS", message=f"{ask_details_result_request.query_id} is not found", ), + trace_id=None, )wren-ai-service/src/web/v1/services/relationship_recommendation.py (1)
98-110: 🛠️ Refactor suggestionMissing trace_id in the getitem error path.
The
__getitem__method doesn't include trace_id when creating a Resource for a not found case, which can break the trace chain when retrieving non-existent resources.Consider adding trace_id to the Resource creation in the not found case:
if response is None: message = f"Relationship Recommendation Resource with ID '{id}' not found." logger.exception(message) return self.Resource( id=id, status="failed", error=self.Resource.Error(code="RESOURCE_NOT_FOUND", message=message), + trace_id=None, # Consider retrieving from current context if possible )wren-ai-service/src/web/v1/services/semantics_description.py (1)
134-146: 🛠️ Refactor suggestionMissing trace_id in the getitem error path.
Similar to other services, the
__getitem__method doesn't include trace_id when creating a Resource for a not found case.Consider adding trace_id to the Resource creation in the not found case:
if response is None: message = f"Semantics Description Resource with ID '{id}' not found." logger.exception(message) return self.Resource( id=id, status="failed", error=self.Resource.Error(code="RESOURCE_NOT_FOUND", message=message), + trace_id=None, # Consider retrieving from current context if possible )wren-ai-service/src/web/v1/services/sql_pairs.py (2)
99-120:⚠️ Potential issueMissing trace_id handling in the delete method.
The
deletemethod doesn't extract or propagate the trace_id, unlike theindexmethod. This creates inconsistency in trace ID propagation.Update the delete method to handle trace_id similar to the index method:
@observe(name="Delete SQL Pairs") @trace_metadata async def delete( self, request: DeleteRequest, **kwargs, ): logger.info(f"Request {request.id}: SQL Pairs Deletion process is running...") + trace_id = kwargs.get("trace_id") try: sql_pairs = [SqlPair(id=id) for id in request.sql_pair_ids] await self._pipelines["sql_pairs"].clean( sql_pairs=sql_pairs, project_id=request.project_id ) - self._cache[request.id] = self.Resource(id=request.id, status="finished") + self._cache[request.id] = self.Resource( + id=request.id, + status="finished", + trace_id=trace_id + ) except Exception as e: self._handle_exception( request.id, f"Failed to delete SQL pairs: {e}", + trace_id=trace_id, )
122-134: 🛠️ Refactor suggestionMissing trace_id in the getitem error path.
The
__getitem__method doesn't include trace_id when creating a Resource for a not found case.Consider adding trace_id to the Resource creation in the not found case:
if response is None: message = f"SQL Pairs Resource with ID '{id}' not found." logger.exception(message) return self.Resource( id=id, status="failed", error=self.Resource.Error(code="OTHERS", message=message), + trace_id=None, # Consider retrieving from current context if possible )wren-ai-service/src/web/v1/services/question_recommendation.py (1)
215-227: 🛠️ Refactor suggestionMissing trace_id in the getitem error path.
The
__getitem__method doesn't include trace_id when creating a Resource for a not found case.Consider adding trace_id to the Resource creation in the not found case:
if response is None: message = f"Question Recommendation Resource with ID '{id}' not found." logger.exception(message) return self.Resource( id=id, status="failed", error=self.Resource.Error(code="RESOURCE_NOT_FOUND", message=message), + trace_id=None, # Consider retrieving from current context if possible )
🧹 Nitpick comments (12)
wren-ai-service/src/web/v1/routers/sql_pairs.py (1)
149-149: Add default value for consistency with other response models.While the implementation is correct, consider adding a default value of
Noneto maintain consistency with the other similar response models in the codebase.- trace_id: Optional[str] + trace_id: Optional[str] = Nonewren-ai-service/src/web/v1/routers/semantics_description.py (1)
126-132: Consider adding a comment on the trace_id field's purposeTo improve code readability and understanding, it would be helpful to add a brief comment explaining the purpose of the trace_id field and its role in request tracing.
class GetResponse(BaseModel): id: str status: Literal["generating", "finished", "failed"] response: Optional[list[dict]] error: Optional[dict] + # Optional trace_id for request tracing and debugging across services trace_id: Optional[str] = Nonewren-ai-service/src/web/v1/services/sql_answer.py (1)
141-147: Missing trace_id propagation in error case.The error response in
get_sql_answer_resultmethod doesn't include the trace_id, which could break the tracing chain. Consider adding a trace_id parameter to the method and including it in the error response.def get_sql_answer_result( self, sql_answer_result_request: SqlAnswerResultRequest, + trace_id: Optional[str] = None, ) -> SqlAnswerResultResponse: if ( result := self._sql_answer_results.get(sql_answer_result_request.query_id) ) is None: logger.exception( f"sql answer pipeline - OTHERS: {sql_answer_result_request.query_id} is not found" ) return SqlAnswerResultResponse( status="failed", error=SqlAnswerResultResponse.SqlAnswerError( code="OTHERS", message=f"{sql_answer_result_request.query_id} is not found", ), + trace_id=trace_id, )wren-ai-service/src/web/v1/services/chart.py (2)
191-193: Missing trace_id in stop_chart method.The
stop_chartmethod doesn't propagate the trace_id when setting status to "stopped". Consider adding a trace_id parameter to maintain continuity of tracing.def stop_chart( self, stop_chart_request: StopChartRequest, + trace_id: Optional[str] = None, ): self._chart_results[stop_chart_request.query_id] = ChartResultResponse( status="stopped", + trace_id=trace_id, )
203-209: Missing trace_id propagation in error case.Similar to other services, the error response in the
get_chart_resultmethod doesn't include the trace_id.def get_chart_result( self, chart_result_request: ChartResultRequest, + trace_id: Optional[str] = None, ) -> ChartResultResponse: if (result := self._chart_results.get(chart_result_request.query_id)) is None: logger.exception( f"chart pipeline - OTHERS: {chart_result_request.query_id} is not found" ) return ChartResultResponse( status="failed", error=ChartError( code="OTHERS", message=f"{chart_result_request.query_id} is not found", ), + trace_id=trace_id, )wren-ai-service/src/web/v1/services/chart_adjustment.py (2)
207-211: Missing trace_id in stop_chart_adjustment method.Similar to the
stop_chartmethod, thestop_chart_adjustmentmethod doesn't include trace_id when setting status to "stopped".def stop_chart_adjustment( self, stop_chart_adjustment_request: StopChartAdjustmentRequest, + trace_id: Optional[str] = None, ): self._chart_adjustment_results[ stop_chart_adjustment_request.query_id ] = ChartAdjustmentResultResponse( status="stopped", + trace_id=trace_id, )
225-231: Missing trace_id propagation in error case.The error response in the
get_chart_adjustment_resultmethod doesn't include the trace_id, which could break the tracing chain.def get_chart_adjustment_result( self, chart_adjustment_result_request: ChartAdjustmentResultRequest, + trace_id: Optional[str] = None, ) -> ChartAdjustmentResultResponse: if ( result := self._chart_adjustment_results.get( chart_adjustment_result_request.query_id ) ) is None: logger.exception( f"chart adjustment pipeline - OTHERS: {chart_adjustment_result_request.query_id} is not found" ) return ChartAdjustmentResultResponse( status="failed", error=ChartAdjustmentError( code="OTHERS", message=f"{chart_adjustment_result_request.query_id} is not found", ), + trace_id=trace_id, )wren-ai-service/src/web/v1/services/ask.py (4)
539-541: Missing trace_id in stop_ask method.The
stop_askmethod doesn't include trace_id when setting status to "stopped".def stop_ask( self, stop_ask_request: StopAskRequest, + trace_id: Optional[str] = None, ): self._ask_results[stop_ask_request.query_id] = AskResultResponse( status="stopped", + trace_id=trace_id, )
730-734: Missing trace_id in stop_ask_feedback method.Similar to other stop methods,
stop_ask_feedbackdoesn't propagate the trace_id.def stop_ask_feedback( self, stop_ask_feedback_request: StopAskFeedbackRequest, + trace_id: Optional[str] = None, ): self._ask_feedback_results[stop_ask_feedback_request.query_id] = ( AskFeedbackResultResponse( status="stopped", + trace_id=trace_id, ) )
551-558: Missing trace_id propagation in get result error responses.Both
get_ask_resultandget_ask_feedback_resultmethods don't include the trace_id in their error responses, which could break the tracing chain.Consider adding a trace_id parameter to both methods and including it in the error responses. For example:
def get_ask_result( self, ask_result_request: AskResultRequest, + trace_id: Optional[str] = None, ) -> AskResultResponse: if (result := self._ask_results.get(ask_result_request.query_id)) is None: logger.exception( f"ask pipeline - OTHERS: {ask_result_request.query_id} is not found" ) return AskResultResponse( status="failed", type="TEXT_TO_SQL", error=AskError( code="OTHERS", message=f"{ask_result_request.query_id} is not found", ), + trace_id=trace_id, )And similarly for the
get_ask_feedback_resultmethod.Also applies to: 748-754
52-52: End-to-end trace ID propagation is well implemented.The implementation of trace ID propagation across all service responses is well done and follows a consistent pattern. This significantly enhances observability by allowing requests to be traced across different components of the system, making debugging and monitoring much easier.
The pattern of extracting the trace ID from kwargs (passed by the trace_metadata decorator) and then including it in all response objects is consistent across all services. This ensures that the trace ID is maintained throughout the entire request lifecycle.
Also applies to: 73-73, 79-79, 92-92, 100-100, 109-109, 121-121, 159-159, 196-197, 591-592
wren-ai-service/src/web/v1/services/semantics_description.py (1)
116-117: Status update and trace_id setting could be combined.The trace_id is set in a separate statement after updating the status, which works but could be more concise.
Consider combining the status and trace_id updates:
- self[request.id].status = "finished" - self[request.id].trace_id = trace_id + resource = self[request.id] + resource.status = "finished" + resource.trace_id = trace_idor:
- self[request.id].status = "finished" - self[request.id].trace_id = trace_id + self[request.id] = self.Resource( + **self[request.id].model_dump(), + status="finished", + trace_id=trace_id + )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (16)
wren-ai-service/src/utils.py(1 hunks)wren-ai-service/src/web/v1/routers/question_recommendation.py(2 hunks)wren-ai-service/src/web/v1/routers/relationship_recommendation.py(2 hunks)wren-ai-service/src/web/v1/routers/semantics_description.py(2 hunks)wren-ai-service/src/web/v1/routers/sql_pairs.py(2 hunks)wren-ai-service/src/web/v1/services/ask.py(23 hunks)wren-ai-service/src/web/v1/services/ask_details.py(5 hunks)wren-ai-service/src/web/v1/services/chart.py(6 hunks)wren-ai-service/src/web/v1/services/chart_adjustment.py(7 hunks)wren-ai-service/src/web/v1/services/question_recommendation.py(5 hunks)wren-ai-service/src/web/v1/services/relationship_recommendation.py(3 hunks)wren-ai-service/src/web/v1/services/semantics_description.py(4 hunks)wren-ai-service/src/web/v1/services/sql_answer.py(5 hunks)wren-ai-service/src/web/v1/services/sql_expansion.py(7 hunks)wren-ai-service/src/web/v1/services/sql_pairs.py(4 hunks)wren-ai-service/src/web/v1/services/sql_question.py(5 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (4)
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: pytest
- GitHub Check: pytest
- GitHub Check: Analyze (go)
🔇 Additional comments (51)
wren-ai-service/src/utils.py (1)
126-128: Great implementation of trace ID propagation.The addition of the trace ID retrieval and passing it as a keyword argument to the decorated function ensures proper propagation of tracing context throughout the function call chain. This implementation aligns perfectly with the PR objective of enhancing observability across service responses.
I particularly like how this solution maintains backward compatibility while adding the new trace functionality, as the trace_id is passed as a keyword argument.
wren-ai-service/src/web/v1/routers/question_recommendation.py (2)
125-125: Appropriate field addition for trace propagation.The optional trace_id field with a default value of None allows for backward compatibility while enabling the new tracing functionality.
151-151: Consistent implementation of trace ID inclusion in response.Properly propagates the trace_id from the resource to the API response, enhancing the observability of requests.
wren-ai-service/src/web/v1/routers/relationship_recommendation.py (2)
110-110: Appropriate field addition for trace propagation.The optional trace_id field with a default value of None allows for backward compatibility while enabling the new tracing functionality.
127-127: Consistent implementation of trace ID inclusion in response.Properly propagates the trace_id from the resource to the API response, enhancing the observability of requests.
wren-ai-service/src/web/v1/routers/sql_pairs.py (1)
161-161: Consistent implementation of trace ID inclusion in response.Properly propagates the trace_id from the resource to the API response, enhancing the observability of requests.
wren-ai-service/src/web/v1/services/sql_question.py (3)
48-48: Good addition of trace_id field to response modelAdding trace_id to SqlQuestionResultResponse allows for improved request tracking and debugging capabilities across services.
69-69: LGTM: Proper extraction of trace_id from kwargsCorrectly extracting the trace_id from kwargs to be used in response objects throughout the request lifecycle.
83-83: Consistent trace_id propagation across all response statesGood implementation of trace_id propagation across all possible service states: generating, succeeded, and failed. This ensures end-to-end traceability regardless of the request outcome.
Also applies to: 96-96, 112-112
wren-ai-service/src/web/v1/services/sql_expansion.py (3)
74-74: Good addition of trace_id field to response modelAdding trace_id to SqlExpansionResultResponse enhances the observability of the system by enabling tracing across service boundaries.
103-103: LGTM: Proper extraction of trace_id from kwargsCorrectly extracting the trace_id from kwargs to be used in response objects throughout the request lifecycle.
119-119: Comprehensive trace_id propagation across all response statesExcellent implementation of trace_id propagation across all possible service states: understanding, searching, generating, finished, and failed. This ensures the trace_id is consistently maintained throughout the entire request flow.
Also applies to: 125-125, 148-148, 156-156, 219-219, 239-239, 255-255
wren-ai-service/src/web/v1/services/ask_details.py (3)
62-62: Good addition of trace_id field to response modelAdding trace_id to AskDetailsResultResponse enables consistent request tracing across different services.
92-92: LGTM: Proper extraction of trace_id from kwargsCorrectly extracting the trace_id from kwargs to be used in response objects throughout the request lifecycle.
108-108: Thorough trace_id propagation across all response statesGood implementation of trace_id propagation across all possible service states: understanding, searching, generating, finished, and failed. This ensures that trace_id is consistently available throughout the entire processing flow.
Also applies to: 113-113, 118-118, 157-157, 173-173
wren-ai-service/src/web/v1/routers/semantics_description.py (3)
131-131: Good addition of trace_id field to the GetResponse modelAdding trace_id to the GetResponse model enables consistent request tracing across different services and endpoints.
168-168: LGTM: Proper propagation of trace_id from resource to responseCorrectly propagating the trace_id from the resource to the GetResponse ensures end-to-end traceability for semantic description requests.
1-169:✅ Verification successful
Verify that SemanticsDescription.Resource can store trace_id
The code assumes that
resource.trace_idexists in line 168, but it's not clear from this file whether theSemanticsDescription.Resourceclass has been updated to include this field. Please verify that thetrace_idfield has been added to theResourceclass elsewhere in the codebase.
🏁 Script executed:
#!/bin/bash # Check if SemanticsDescription.Resource class includes the trace_id field # by searching for its definition in the codebase # Search for the definition of the Resource class in the relevant file(s) rg -A 20 "class Resource" --glob "src/web/v1/services/**/*.py"Length of output: 213
Let's run a new check without the restrictive glob to search for the
Resourceclass definition (and itstrace_idattribute) across the entire repository:Please execute this script and share the output so we can verify whether the
trace_idattribute is indeed present in theSemanticsDescription.Resourceclass definition.
🏁 Script executed:
#!/bin/bash echo "Searching the entire repository for the 'Resource' class definition along with nearby context (to check for 'trace_id'):" rg -t py -A 30 "class\s+Resource" .Length of output: 11799
Trace_ID Field Verified in SemanticsDescription.Resource
- The search across the repository confirms that
SemanticsDescription.Resource(in wren-ai-service/src/web/v1/services/semantics_description.py) defines atrace_idfield asOptional[str] = None.- This matches the assumption in the router code at line 168.
No further actions are necessary.
wren-ai-service/src/web/v1/services/sql_answer.py (3)
52-52: Trace ID field addition looks good.The addition of the optional
trace_idfield to theSqlAnswerResultResponseclass is a good enhancement for observability.
73-73: Properly extracting trace_id from kwargs.This correctly extracts the trace ID from kwargs that was passed by the trace_metadata decorator.
89-90: Propagation of trace_id is consistent.The trace ID is properly propagated to all response objects when setting different statuses.
Also applies to: 99-100, 124-125
wren-ai-service/src/web/v1/services/chart.py (3)
79-79: Trace ID field addition follows consistent pattern.The addition of the optional
trace_idfield to theChartResultResponseclass is consistent with the tracing implementation across services.
109-109: Properly extracting trace_id from kwargs.This correctly extracts the trace ID from kwargs passed by the trace_metadata decorator.
121-125: Comprehensive trace_id propagation in all response paths.The trace ID is properly propagated across all response states including initial state, generation, failure, and successful completion. This provides complete request traceability.
Also applies to: 136-139, 153-159, 163-167, 174-181
wren-ai-service/src/web/v1/services/chart_adjustment.py (3)
92-92: Trace ID field addition is consistent with other models.The addition of the optional
trace_idfield to theChartAdjustmentResultResponseclass follows the same pattern used in other response models.
121-121: Properly extracting trace_id from kwargs.This correctly extracts the trace ID from kwargs passed by the trace_metadata decorator.
133-136: Comprehensive trace_id propagation in all response paths.The trace ID is consistently included in all response states, which is excellent for maintaining trace context throughout the request lifecycle.
Also applies to: 145-148, 163-171, 175-181, 188-197
wren-ai-service/src/web/v1/services/ask.py (4)
100-100: Trace ID fields added to both response models.The addition of the optional
trace_idfield to bothAskResultResponseandAskFeedbackResultResponseclasses is consistent with the overall tracing approach.Also applies to: 159-159
196-197: Properly extracting trace_id from kwargs in both methods.Both
askandask_feedbackmethods correctly extract the trace ID from kwargs passed by the trace_metadata decorator.Also applies to: 591-592
222-224: Comprehensive trace_id propagation in the ask method.The trace ID is consistently propagated across all paths in the
askmethod, including all various states and error conditions. This provides excellent traceability.Also applies to: 272-273, 294-295, 304-305, 312-313, 339-340, 356-357, 383-384, 394-395, 453-454, 492-493, 510-511, 527-528
608-609: Consistent trace_id propagation in the ask_feedback method.The trace ID is properly propagated across all response paths in the
ask_feedbackmethod, maintaining consistency with the tracing implementation.Also applies to: 624-625, 656-657, 691-692, 703-704, 719-720
wren-ai-service/src/web/v1/services/relationship_recommendation.py (5)
32-32: Good implementation of the optional trace_id field.This correctly adds the trace ID tracking capability to the Resource model while maintaining backward compatibility with the optional typing.
45-57: Properly updated exception handling to propagate trace_id.The exception handling has been updated to accept and propagate the trace ID, which is essential for maintaining traceability during error conditions.
64-64: Good extraction of trace_id from kwargs.Extracting the trace_id from kwargs at the beginning of the method ensures it's available throughout the entire function scope.
76-81: Trace ID correctly propagated to successful responses.The trace_id is properly included in the Resource instance when creating a successful response, which completes the traceability chain.
82-94: Error handling properly includes trace_id in both exception cases.Both error handling branches consistently include the trace_id parameter, ensuring traceability is maintained in exception paths.
wren-ai-service/src/web/v1/services/semantics_description.py (4)
27-27: Good addition of the trace_id field.The trace_id is correctly added as an optional field to maintain backward compatibility.
38-50: Exception handling correctly updated to include trace_id.The _handle_exception method has been properly updated to accept and propagate the trace_id parameter.
106-107: Good extraction of trace_id from kwargs.Extracting the trace_id at the beginning of the method ensures it's available throughout the function scope.
118-130: Error handling consistently includes trace_id.Both error cases properly pass the trace_id parameter to the _handle_exception method.
wren-ai-service/src/web/v1/services/sql_pairs.py (5)
25-25: Good addition of the trace_id field.The trace_id is correctly added as an optional field to maintain backward compatibility.
36-48: Exception handling correctly updated to include trace_id.The _handle_exception method has been properly updated to accept and propagate the trace_id parameter.
64-65: Good extraction of trace_id from kwargs.Extracting the trace_id at the beginning of the method ensures it's available throughout the function scope.
78-83: Resource creation correctly includes trace_id.The trace_id is properly included when creating a new Resource instance for successful operations.
84-89: Error handling correctly includes trace_id.The exception handling case properly passes the trace_id to the _handle_exception method.
wren-ai-service/src/web/v1/services/question_recommendation.py (6)
37-37: Good addition of the trace_id field.The trace_id is correctly added as an optional field to maintain backward compatibility.
50-62: Exception handling correctly updated to include trace_id.The _handle_exception method has been properly updated to accept and propagate the trace_id parameter.
157-158: Good extraction of trace_id from kwargs.Extracting the trace_id at the beginning of the method ensures it's available throughout the function scope.
171-173: Trace ID correctly set after recommendation.The resource is retrieved and the trace_id is properly set.
199-211: Error handling consistently includes trace_id.Both exception handlers properly pass the trace_id to the _handle_exception method.
64-211:❓ Verification inconclusive
Verify trace_id is consistently propagated when multiple service requests are involved.
There are multiple entry points where trace_id could be lost if not explicitly passed in the
_validate_questionand_recommendmethods that interact with pipelines.Let's check if trace_id is properly propagated in the internal helper methods:
🏁 Script executed:
#!/bin/bash # Check if trace_id is propagated in the helper methods # First, check if the utils.trace_metadata handles trace_id propagation correctly rg "def trace_metadata" -A 20 src/utils.py # Check if _validate_question and _recommend functions in question_recommendation.py pass trace_id to pipeline calls rg -A 5 "_validate_question" src/web/v1/services/question_recommendation.py rg -A 5 "_recommend\(" src/web/v1/services/question_recommendation.py # Check if other services pass trace_id to helper methods and pipeline calls rg -A 5 "await self\._pipelines\[" src/web/v1/services/Length of output: 538
Action Required: Validate Consistent Trace ID Propagation Across Pipelines
The current implementation extracts
trace_idin the publicrecommendmethod and assigns it to the cached resource. However, the internal helper methods (_validate_questionand_recommend) do not explicitly propagate thetrace_idwhen initiating pipeline runs. This could lead to scenarios in which, during multiple asynchronous service requests, the trace identifier is not consistently available across all internal pipeline invocations.
- Verify that the
@trace_metadatadecorator (or any global context mechanism) reliably injects thetrace_idinto every pipeline call.- If it does not, consider refactoring these helper methods to accept and forward the
trace_idexplicitly to ensure end-to-end traceability.
paopa
force-pushed
the
chore/trace-id-for-all-endpoints
branch
2 times, most recently
from
08ac9ad to
edea1db
Compare
paopa
force-pushed
the
chore/trace-id-for-all-endpoints
branch
from
edea1db to
931f1f4
Compare
Changes
trace_idfield to response models across multiple services to enable end-to-end tracingtrace_metadatadecorator to pass trace ID through function callsImpact
This change improves observability by allowing requests to be traced across different services and components. The trace ID is now consistently propagated through the entire request lifecycle, making it easier to debug and monitor request flows.
Summary by CodeRabbit