LCORE-385: Restructure tool_call responses #350
Conversation
…ool execution events and improve handling of turn completion. Update tests to reflect changes in response format and ensure proper assertions.
WalkthroughThe changes update the streaming query endpoint to emit structured JSON tokens for tool events instead of concatenated strings, and adjust how the final response is captured from the stream. Corresponding unit tests are updated to expect the new token structure, more granular chunking, and a turn completion event. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant API
participant Model/Tools
Client->>API: Send streaming query request
loop For each event in stream
API->>Model/Tools: Process next event
alt Tool execution event
Model/Tools-->>API: Tool event (structured JSON token)
API-->>Client: Stream JSON token (tool_name, arguments, etc.)
else Turn complete event
Model/Tools-->>API: Turn complete event (interleaved content)
API-->>Client: Stream JSON token (final response)
end
end
API-->>Client: Stream ends
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~15 minutes Possibly related PRs
Suggested reviewers
Poem
Note ⚡️ Unit Test Generation is now available in beta!Learn more here, or try it out under "Finishing Touches" below. ✨ Finishing Touches
🧪 Generate unit tests
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:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. 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
🧹 Nitpick comments (2)
src/app/endpoints/streaming_query.py (1)
470-476: Consider more robust JSON parsing for turn_complete events.While the logic correctly captures the turn_complete token, the current implementation manually parses the SSE format. Consider a more robust approach:
- if ( - json.loads(event.replace("data: ", ""))["event"] - == "turn_complete" - ): - complete_response = json.loads(event.replace("data: ", ""))[ - "data" - ]["token"] + # Parse SSE data more safely + if event.startswith("data: "): + try: + data = json.loads(event[6:]) # Skip "data: " prefix + if data.get("event") == "turn_complete": + complete_response = data.get("data", {}).get("token", complete_response) + except json.JSONDecodeError: + logger.warning("Failed to parse event data: %s", event)This approach:
- Safely checks for the "data: " prefix
- Handles JSON parsing errors gracefully
- Uses
.get()to avoid KeyError exceptions- Maintains the default response if parsing fails
tests/unit/app/endpoints/test_streaming_query.py (1)
164-330: Consider adding edge case tests for robustness.While the existing tests cover the happy path, consider adding tests for:
- Missing turn_complete event (should return default "No response from the model")
- Tool responses with missing or malformed structure
- Multiple turn_complete events in a stream
Example test case for missing turn_complete:
@pytest.mark.asyncio async def test_streaming_query_missing_turn_complete(mocker): """Test streaming query when turn_complete event is missing.""" # ... setup code ... # Create response without turn_complete event mock_streaming_response.__aiter__.return_value = [ AgentTurnResponseStreamChunk( event=TurnResponseEvent( payload=AgentTurnResponseStepProgressPayload( event_type="step_progress", step_type="inference", delta=TextDelta(text="Incomplete response", type="text"), step_id="s1", ) ) ), ] # ... rest of test ... # Assert that stored response is "No response from the model"
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
src/app/endpoints/streaming_query.py(6 hunks)tests/unit/app/endpoints/test_streaming_query.py(6 hunks)
🔇 Additional comments (9)
src/app/endpoints/streaming_query.py (6)
216-218: LGTM! Proper content formatting for turn completion.The use of
interleaved_content_as_strensures consistent string formatting of the turn's output message content.
340-343: Correctly implements structured tool call tokens.The token now returns a structured dictionary with
tool_nameandargumentsfields, aligning perfectly with the PR objective of standardizing tool_call responses.
357-360: Proper structured response for memory query tool.The implementation correctly returns a dictionary with
tool_nameandresponsefields, providing clear information about the memory fetch operation.
391-394: Knowledge search response properly structured.The token correctly returns a dictionary with
tool_nameandsummaryfields, maintaining the existing summary extraction logic while conforming to the new structured format.
406-409: Generic tool responses correctly structured.The catch-all handler properly returns a dictionary with
tool_nameandresponsefields for all other tool types.
463-463: Good default response initialization.Setting a meaningful default message helps identify cases where no turn_complete event is received, improving debuggability.
tests/unit/app/endpoints/test_streaming_query.py (3)
184-248: Test correctly models the new streaming behavior.The test properly reflects:
- Text being streamed in chunks during inference
- Turn complete event containing the full assembled response
- Proper event structure with all required fields
297-297: Assertions correctly validate new response handling.The test properly validates that:
- The complete response comes from the turn_complete event
- The chunk count includes the new turn_complete event
Also applies to: 311-311
973-978: Test assertions correctly validate structured tool tokens.The assertions properly verify that tool events return structured JSON objects with the required fields (tool_name with arguments or summary), aligning with the PR objectives.
There was a problem hiding this comment.
LGTM 👍
@TamiTakamiya FYI (thanks for your PR in preparation!)
|
/cc @rawagner |
4 participants
Description
This PR makes "tool_call" responses return a structured JSON with tool_name, arguments, summary, or response as fields (every "tool_call" non inference response will always have tool_name as well as one of the other fields in the response). Please let me know if anything should be changed.
Type of change
Related Tickets & Documents
Checklist before requesting a review
Testing
Summary by CodeRabbit
New Features
Bug Fixes
Tests