Add work-around for tool_calls with granite models.#211
Add work-around for tool_calls with granite models.#211tisnik merged 1 commit intolightspeed-core:mainfrom
Conversation
WalkthroughA new Changes
Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (4)
⏰ 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). (2)
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. 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 (
|
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
tests/unit/utils/test_types.py (1)
40-48: Fix docstring and improve mock dataThe docstring is incorrect and the mock data could be more realistic.
Apply this diff to fix the issues:
- def test_get_tool_calls_from_completion_message_when_message_has_tool_calls(self): - """Test that get_tool_calls returns an empty array when CompletionMessage has no tool_calls.""" + def test_get_tool_calls_from_completion_message_when_message_has_tool_calls(self): + """Test that get_tool_calls returns the tool_calls when CompletionMessage has tool_calls.""" tool_parser = GraniteToolParser.get_parser("granite-3.3-8b-instruct") completion_message = Mock() - tool_calls = [{"tool_name": "tool-1"}, {"tool_name": "tool-2"}] + tool_calls = [Mock(tool_name="tool-1"), Mock(tool_name="tool-2")] completion_message.tool_calls = tool_calls assert ( tool_parser.get_tool_calls(completion_message) == tool_calls ), f"get_tool_calls should return {tool_calls}"
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
src/app/endpoints/streaming_query.py(2 hunks)src/utils/types.py(2 hunks)tests/unit/utils/test_types.py(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (2)
src/app/endpoints/streaming_query.py (1)
src/utils/types.py (2)
GraniteToolParser(23-37)get_parser(33-37)
tests/unit/utils/test_types.py (1)
src/utils/types.py (3)
GraniteToolParser(23-37)get_parser(33-37)get_tool_calls(26-30)
⏰ 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). (2)
- GitHub Check: unit_tests (3.12)
- GitHub Check: build-pr
🔇 Additional comments (8)
src/utils/types.py (2)
3-7: LGTM: Imports are appropriateThe new imports are necessary for the
GraniteToolParserimplementation and follow proper typing conventions.
22-37: LGTM: Well-implemented workaround with good defensive programmingThe implementation correctly:
- References the GitHub issue for context
- Handles null/empty inputs gracefully
- Uses case-insensitive model ID matching
- Provides a clean factory pattern via the static method
The defensive checks (
if output_message and output_message.tool_callsandif model_id and model_id.lower().startswith("granite")) properly handle edge cases.src/app/endpoints/streaming_query.py (2)
27-27: LGTM: Clean import additionThe import is properly placed and follows the existing import structure.
65-65: LGTM: Proper integration of the granite tool parserThe integration correctly passes the result of
GraniteToolParser.get_parser(model_id)to the AsyncAgent constructor. Since the method returnsNonefor non-granite models, this maintains backward compatibility while enabling the workaround for granite models.tests/unit/utils/test_types.py (4)
1-6: LGTM: Proper test setupThe test module is well-structured with appropriate imports and clear organization.
9-24: LGTM: Good coverage of the get_parser static methodThe tests properly verify that the parser is only returned for model IDs that start with "granite", covering both negative and positive cases.
26-29: LGTM: Proper null handling testThe test correctly verifies that
get_tool_callsreturns an empty list when passedNone.
31-38: LGTM: Empty tool_calls testThe test correctly verifies the behavior when a completion message has an empty tool_calls list.
3a00de6 to
61651bc
Compare
tisnik
left a comment
There was a problem hiding this comment.
nit, will be probably repeated by Rabbit too ;)
src/utils/types.py
Outdated
| @@ -1,5 +1,11 @@ | |||
| """Common types for the project.""" | |||
|
|
|||
| from typing import List, Optional | |||
There was a problem hiding this comment.
Ah, well, @tisnik this is interesting.
I originally had list as I know you prefer this over List however one of the CI checks failed saying the method signature of the get_tool_calls(..) didn't match that of the super class (that uses List)...
I've changed to list.. let's see what CI says 😄
61651bc to
d661a29
Compare
d661a29 to
0596fbf
Compare
Description
During the integration of Ansible Lightspeed with vanilla
llama-stackinvocation of MCP Tool Calls was troublesum.After much investigation and research by @TamiTakamiya an issue in
llama-stackwith IBMgranitemodels was identified.The fix in this PR was required in our prototype
llama-stackwork to optimise MCP Tool Calls withgranitemodels.We are contributing the same work-around to
lightspeed-stack.Type of change
Related Tickets & Documents
Checklist before requesting a review
Testing
Summary by CodeRabbit