Refine toolbox and extend tool handling modes#374
Conversation
|
Caution Review failedThe pull request is closed. """ WalkthroughThis change introduces extensive updates across the codebase to enhance tool invocation, error handling, and result aggregation in the Draive framework. The tool handling system is refactored to support multiple handling modes ("auto", "extend", "direct", "spawn"), with corresponding updates to type annotations, protocols, and implementation logic. The Estimated code review effort5 (~150–180 minutes) Possibly related PRs
📜 Recent review detailsConfiguration used: CodeRabbit UI ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (25)
✨ Finishing Touches
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. 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: 4
🔭 Outside diff range comments (3)
src/draive/generation/model/default.py (1)
27-36: Missing generic type parameter for decoderThe
decoderparameter should use the generic type to ensure type safety consistency with the protocol definition.- decoder: ModelGeneratorDecoder | None, + decoder: ModelGeneratorDecoder[Generated] | None,src/draive/conversation/completion/default.py (1)
167-330: Complex but well-structured streaming implementation.The refactored streaming logic correctly handles concurrent tool execution and incremental result streaming. The use of
asyncio.waitwithFIRST_COMPLETEDis an excellent choice for processing tool responses as they complete.Consider adding more inline documentation to explain the complex async flow, particularly around the concurrent tool handling logic (lines 262-292).
src/draive/tools/toolbox.py (1)
176-176: Missing required parameter inavailable()callsAccording to the
Toolprotocol, theavailable()method requires atools_turnparameter, but it's being called without any arguments here.Apply this fix:
- elif self.suggest_call.available: # use suggested tool if able + elif self.suggest_call.available(tools_turn): # use suggested tool if able- tuple(tool.specification for tool in self.tools.values() if tool.available), + tuple(tool.specification for tool in self.tools.values() if tool.available(tools_turn)),Also applies to: 187-187
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (22)
pyproject.toml(1 hunks)src/draive/__init__.py(2 hunks)src/draive/anthropic/lmm_generation.py(1 hunks)src/draive/choice/default.py(2 hunks)src/draive/conversation/completion/default.py(3 hunks)src/draive/generation/model/default.py(2 hunks)src/draive/generation/model/state.py(1 hunks)src/draive/generation/model/types.py(2 hunks)src/draive/generation/text/default.py(3 hunks)src/draive/generation/text/state.py(1 hunks)src/draive/generation/text/types.py(0 hunks)src/draive/lmm/types.py(3 hunks)src/draive/mcp/client.py(1 hunks)src/draive/mistral/lmm_generation.py(2 hunks)src/draive/ollama/lmm_generation.py(2 hunks)src/draive/openai/lmm_generation.py(2 hunks)src/draive/stages/stage.py(2 hunks)src/draive/tools/__init__.py(2 hunks)src/draive/tools/function.py(16 hunks)src/draive/tools/toolbox.py(10 hunks)src/draive/tools/types.py(4 hunks)src/draive/vllm/lmm_generation.py(2 hunks)
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit Inference Engine (CLAUDE.md)
**/*.py: Use absolute imports fromdraivepackage
Follow Ruff import ordering (standard library, third party, local)
Use Python 3.12+ type features (type unions with|, generic syntax)
Use base and abstract types likeSequenceorIterableinstead of concrete types
Use custom exceptions for specific errors
Use NumPy docstring convention for all functions, classes, and methods
Skip module-level docstrings unless explicitly requested
Include sections: Parameters, Returns, Raises, Notes (if needed) in docstrings
ALWAYS useSequence[T]instead oflist[T],Mapping[K,V]instead ofdict[K,V], andSet[T]instead ofset[T]for collections in State, Config and DataModel classes
Immutable updates to State, Config, and DataModel classes should be performed through the.updated()method
Files:
src/draive/mcp/client.pysrc/draive/tools/__init__.pysrc/draive/generation/model/state.pysrc/draive/vllm/lmm_generation.pysrc/draive/__init__.pysrc/draive/ollama/lmm_generation.pysrc/draive/openai/lmm_generation.pysrc/draive/anthropic/lmm_generation.pysrc/draive/generation/model/types.pysrc/draive/generation/text/state.pysrc/draive/choice/default.pysrc/draive/conversation/completion/default.pysrc/draive/generation/model/default.pysrc/draive/lmm/types.pysrc/draive/generation/text/default.pysrc/draive/tools/function.pysrc/draive/mistral/lmm_generation.pysrc/draive/stages/stage.pysrc/draive/tools/types.pysrc/draive/tools/toolbox.py
**/__init__.py
📄 CodeRabbit Inference Engine (CLAUDE.md)
Put exported symbols into
__init__.py
Files:
src/draive/tools/__init__.pysrc/draive/__init__.py
🧠 Learnings (7)
pyproject.toml (1)
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: draive builds on top of haiway and exports its symbols
src/draive/tools/__init__.py (3)
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/*.py : Use absolute imports from draive package
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/init.py : Put exported symbols into __init__.py
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/*.py : Use custom exceptions for specific errors
src/draive/__init__.py (4)
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/*.py : Use absolute imports from draive package
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/init.py : Put exported symbols into __init__.py
Learnt from: KaQuMiQ
PR: #338
File: src/draive/lmm/init.py:1-2
Timestamp: 2025-06-16T10:28:07.434Z
Learning: The draive project requires Python 3.12+ as specified in pyproject.toml with "requires-python = ">=3.12"" and uses Python 3.12+ specific features like PEP 695 type aliases and generic syntax extensively throughout the codebase.
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: draive builds on top of haiway and exports its symbols
src/draive/generation/text/state.py (1)
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/*.py : Use base and abstract types like Sequence or Iterable instead of concrete types
src/draive/conversation/completion/default.py (1)
Learnt from: KaQuMiQ
PR: #338
File: src/draive/bedrock/lmm_generation.py:102-107
Timestamp: 2025-06-16T09:44:08.688Z
Learning: LMMInstruction is a type alias for str in the draive codebase, so LMMInstruction and str are equivalent types and can be used interchangeably.
src/draive/generation/text/default.py (1)
Learnt from: KaQuMiQ
PR: #338
File: src/draive/bedrock/lmm_generation.py:102-107
Timestamp: 2025-06-16T09:44:08.688Z
Learning: LMMInstruction is a type alias for str in the draive codebase, so LMMInstruction and str are equivalent types and can be used interchangeably.
src/draive/tools/types.py (1)
Learnt from: KaQuMiQ
PR: #327
File: src/draive/helpers/instruction_preparation.py:28-34
Timestamp: 2025-05-28T17:41:57.460Z
Learning: The draive project uses and requires Python 3.12+, so PEP-695 generic syntax with square brackets (e.g., def func[T: Type]()) is valid and should be used instead of the older TypeVar approach.
🧬 Code Graph Analysis (8)
src/draive/tools/__init__.py (1)
src/draive/tools/types.py (1)
ToolError(24-31)
src/draive/generation/model/state.py (1)
src/draive/generation/model/types.py (1)
ModelGeneratorDecoder(17-21)
src/draive/__init__.py (1)
src/draive/tools/types.py (1)
ToolError(24-31)
src/draive/ollama/lmm_generation.py (1)
src/draive/lmm/types.py (10)
LMMToolRequests(289-306)of(85-108)of(144-153)of(164-173)of(195-211)of(222-231)of(267-281)of(291-302)of(316-327)of(382-392)
src/draive/openai/lmm_generation.py (2)
src/draive/lmm/types.py (10)
LMMToolRequests(289-306)of(85-108)of(144-153)of(164-173)of(195-211)of(222-231)of(267-281)of(291-302)of(316-327)of(382-392)src/draive/ollama/lmm_generation.py (3)
lmm_completion(41-51)lmm_completion(54-64)lmm_completion(66-175)
src/draive/anthropic/lmm_generation.py (2)
src/draive/lmm/types.py (11)
LMMToolRequests(289-306)of(85-108)of(144-153)of(164-173)of(195-211)of(222-231)of(267-281)of(291-302)of(316-327)of(382-392)LMMToolRequest(265-286)src/draive/tools/function.py (5)
call(186-253)tool(289-308)tool(312-361)tool(365-410)tool(413-443)
src/draive/tools/function.py (4)
src/draive/utils/processing.py (4)
Processing(103-209)event(164-168)event(172-180)event(183-205)src/draive/tools/types.py (7)
ToolAvailabilityChecking(35-40)ToolResultFormatting(44-48)ToolErrorFormatting(52-56)handling(77-77)meta(74-74)call(84-89)name(62-62)src/draive/multimodal/content.py (2)
meta(151-164)MultimodalContent(24-236)src/draive/lmm/types.py (10)
of(85-108)of(144-153)of(164-173)of(195-211)of(222-231)of(267-281)of(291-302)of(316-327)of(382-392)LMMToolError(132-139)
src/draive/tools/toolbox.py (3)
src/draive/tools/types.py (4)
meta(74-74)Tool(60-89)call(84-89)handling(77-77)src/draive/multimodal/content.py (2)
meta(151-164)MultimodalContent(24-236)src/draive/lmm/types.py (11)
of(85-108)of(144-153)of(164-173)of(195-211)of(222-231)of(267-281)of(291-302)of(316-327)of(382-392)LMMToolResponse(193-217)LMMToolError(132-139)
💤 Files with no reviewable changes (1)
- src/draive/generation/text/types.py
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit Inference Engine (CLAUDE.md)
**/*.py: Use absolute imports fromdraivepackage
Follow Ruff import ordering (standard library, third party, local)
Use Python 3.12+ type features (type unions with|, generic syntax)
Use base and abstract types likeSequenceorIterableinstead of concrete types
Use custom exceptions for specific errors
Use NumPy docstring convention for all functions, classes, and methods
Skip module-level docstrings unless explicitly requested
Include sections: Parameters, Returns, Raises, Notes (if needed) in docstrings
ALWAYS useSequence[T]instead oflist[T],Mapping[K,V]instead ofdict[K,V], andSet[T]instead ofset[T]for collections in State, Config and DataModel classes
Immutable updates to State, Config, and DataModel classes should be performed through the.updated()method
Files:
src/draive/mcp/client.pysrc/draive/tools/__init__.pysrc/draive/generation/model/state.pysrc/draive/vllm/lmm_generation.pysrc/draive/__init__.pysrc/draive/ollama/lmm_generation.pysrc/draive/openai/lmm_generation.pysrc/draive/anthropic/lmm_generation.pysrc/draive/generation/model/types.pysrc/draive/generation/text/state.pysrc/draive/choice/default.pysrc/draive/conversation/completion/default.pysrc/draive/generation/model/default.pysrc/draive/lmm/types.pysrc/draive/generation/text/default.pysrc/draive/tools/function.pysrc/draive/mistral/lmm_generation.pysrc/draive/stages/stage.pysrc/draive/tools/types.pysrc/draive/tools/toolbox.py
**/__init__.py
📄 CodeRabbit Inference Engine (CLAUDE.md)
Put exported symbols into
__init__.py
Files:
src/draive/tools/__init__.pysrc/draive/__init__.py
🧠 Learnings (7)
pyproject.toml (1)
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: draive builds on top of haiway and exports its symbols
src/draive/tools/__init__.py (3)
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/*.py : Use absolute imports from draive package
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/init.py : Put exported symbols into __init__.py
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/*.py : Use custom exceptions for specific errors
src/draive/__init__.py (4)
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/*.py : Use absolute imports from draive package
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/init.py : Put exported symbols into __init__.py
Learnt from: KaQuMiQ
PR: #338
File: src/draive/lmm/init.py:1-2
Timestamp: 2025-06-16T10:28:07.434Z
Learning: The draive project requires Python 3.12+ as specified in pyproject.toml with "requires-python = ">=3.12"" and uses Python 3.12+ specific features like PEP 695 type aliases and generic syntax extensively throughout the codebase.
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: draive builds on top of haiway and exports its symbols
src/draive/generation/text/state.py (1)
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/*.py : Use base and abstract types like Sequence or Iterable instead of concrete types
src/draive/conversation/completion/default.py (1)
Learnt from: KaQuMiQ
PR: #338
File: src/draive/bedrock/lmm_generation.py:102-107
Timestamp: 2025-06-16T09:44:08.688Z
Learning: LMMInstruction is a type alias for str in the draive codebase, so LMMInstruction and str are equivalent types and can be used interchangeably.
src/draive/generation/text/default.py (1)
Learnt from: KaQuMiQ
PR: #338
File: src/draive/bedrock/lmm_generation.py:102-107
Timestamp: 2025-06-16T09:44:08.688Z
Learning: LMMInstruction is a type alias for str in the draive codebase, so LMMInstruction and str are equivalent types and can be used interchangeably.
src/draive/tools/types.py (1)
Learnt from: KaQuMiQ
PR: #327
File: src/draive/helpers/instruction_preparation.py:28-34
Timestamp: 2025-05-28T17:41:57.460Z
Learning: The draive project uses and requires Python 3.12+, so PEP-695 generic syntax with square brackets (e.g., def func[T: Type]()) is valid and should be used instead of the older TypeVar approach.
🧬 Code Graph Analysis (8)
src/draive/tools/__init__.py (1)
src/draive/tools/types.py (1)
ToolError(24-31)
src/draive/generation/model/state.py (1)
src/draive/generation/model/types.py (1)
ModelGeneratorDecoder(17-21)
src/draive/__init__.py (1)
src/draive/tools/types.py (1)
ToolError(24-31)
src/draive/ollama/lmm_generation.py (1)
src/draive/lmm/types.py (10)
LMMToolRequests(289-306)of(85-108)of(144-153)of(164-173)of(195-211)of(222-231)of(267-281)of(291-302)of(316-327)of(382-392)
src/draive/openai/lmm_generation.py (2)
src/draive/lmm/types.py (10)
LMMToolRequests(289-306)of(85-108)of(144-153)of(164-173)of(195-211)of(222-231)of(267-281)of(291-302)of(316-327)of(382-392)src/draive/ollama/lmm_generation.py (3)
lmm_completion(41-51)lmm_completion(54-64)lmm_completion(66-175)
src/draive/anthropic/lmm_generation.py (2)
src/draive/lmm/types.py (11)
LMMToolRequests(289-306)of(85-108)of(144-153)of(164-173)of(195-211)of(222-231)of(267-281)of(291-302)of(316-327)of(382-392)LMMToolRequest(265-286)src/draive/tools/function.py (5)
call(186-253)tool(289-308)tool(312-361)tool(365-410)tool(413-443)
src/draive/tools/function.py (4)
src/draive/utils/processing.py (4)
Processing(103-209)event(164-168)event(172-180)event(183-205)src/draive/tools/types.py (7)
ToolAvailabilityChecking(35-40)ToolResultFormatting(44-48)ToolErrorFormatting(52-56)handling(77-77)meta(74-74)call(84-89)name(62-62)src/draive/multimodal/content.py (2)
meta(151-164)MultimodalContent(24-236)src/draive/lmm/types.py (10)
of(85-108)of(144-153)of(164-173)of(195-211)of(222-231)of(267-281)of(291-302)of(316-327)of(382-392)LMMToolError(132-139)
src/draive/tools/toolbox.py (3)
src/draive/tools/types.py (4)
meta(74-74)Tool(60-89)call(84-89)handling(77-77)src/draive/multimodal/content.py (2)
meta(151-164)MultimodalContent(24-236)src/draive/lmm/types.py (11)
of(85-108)of(144-153)of(164-173)of(195-211)of(222-231)of(267-281)of(291-302)of(316-327)of(382-392)LMMToolResponse(193-217)LMMToolError(132-139)
🔇 Additional comments (34)
pyproject.toml (1)
24-24: LGTM - Standard dependency update.The haiway version bump from 0.27.1 to 0.27.2 aligns with the broader tool handling refactoring across the codebase.
src/draive/__init__.py (2)
213-213: LGTM - Consistent exception rename.The import correctly updates from
ToolExceptiontoToolError, aligning with the broader tool error handling refactor.
376-376: LGTM - Export list updated consistently.The
__all__export correctly reflects theToolException→ToolErrorrename.src/draive/tools/__init__.py (2)
7-7: LGTM - Import updated consistently.The import correctly updates from
ToolExceptiontoToolError, maintaining consistency with the tool error handling refactor.
18-18: LGTM - Export list synchronized.The
__all__export correctly reflects the exception rename.src/draive/mcp/client.py (2)
603-603: LGTM - Parameter name updated consistently.The parameter rename from
availability_checktoavailabilityaligns with the broader FunctionTool API refactor.
605-605: LGTM - Error formatting parameter renamed consistently.The parameter rename from
format_failuretoformat_errormaintains consistency with the updated tool error handling terminology.src/draive/vllm/lmm_generation.py (1)
289-301: LGTM - Factory method pattern adoption.The change from direct constructor to
LMMToolRequests.of()factory method improves API consistency across LMM generation modules while maintaining the same functionality.src/draive/ollama/lmm_generation.py (1)
147-159: LGTM! Consistent factory method usage.The refactoring correctly uses
LMMToolRequests.of()with the list ofLMMToolRequestobjects as the first positional argument and content as a keyword argument, aligning with the factory method signature.src/draive/anthropic/lmm_generation.py (1)
264-275: LGTM! Proper factory method implementation.The refactoring correctly uses
LMMToolRequests.of()with appropriate positional and keyword arguments. The explicit null check in the content expression and themeta=META_EMPTYparameter are properly handled.src/draive/openai/lmm_generation.py (1)
292-304: LGTM! Consistent factory method refactoring.The change properly uses
LMMToolRequests.of()with the correct argument structure, maintaining consistency with other LMM provider implementations.src/draive/mistral/lmm_generation.py (1)
269-281: LGTM! Factory method correctly implemented.The refactoring follows the established pattern of using
LMMToolRequests.of()with requests as the first positional argument and content as a keyword argument, maintaining consistency across LMM providers.src/draive/generation/model/state.py (1)
29-29: LGTM! Improved type safety with generic parameter.The type annotation correctly adds the
Generatedtype parameter toModelGeneratorDecoder, ensuring type consistency with the protocol definition and improving type safety throughout the generation pipeline.src/draive/generation/model/types.py (2)
17-21: Type safety improvement with generic decoderThe introduction of generic typing for
ModelGeneratorDecoderimproves type safety by ensuring the decoder returns the expectedDataModeltype rather than a generic dictionary.
36-36: Consistent generic typing for decoder parameterThe decoder parameter correctly uses the generic type, ensuring type consistency between the expected output type and the decoder's return type.
src/draive/generation/text/state.py (2)
2-2: Clean import removalCorrectly removed unused imports after eliminating streaming overloads.
18-35: API simplification by removing streaming parameterThe removal of the
streamparameter and associated overloads simplifies the API while maintaining the ability to return either strings or async iterables. This aligns with the broader refactoring to handle streaming at a different layer.src/draive/generation/model/default.py (2)
90-116: Well-structured tool response handlingThe refactored tool handling logic properly distinguishes between different response types (completion, extension, regular) and accumulates results appropriately. The
tools_turncounter provides clear tracking of interaction rounds.
118-137: Robust decoding with clear fallback chainThe decoding logic provides a clear hierarchy of decoding attempts with appropriate debug logging and error context. The error handling preserves the exception type and provides useful debugging information.
src/draive/choice/default.py (1)
85-123: Consistent tool handling and proper choice extractionThe implementation follows the established pattern for tool handling with proper accumulation of results. The choice extraction logic correctly parses the accumulated content for the selection.
src/draive/lmm/types.py (2)
190-190: Enhanced tool response handling typesThe expansion of handling types supports the new tool interaction patterns with clearer semantics for different response behaviors.
300-305: Good API improvement with non-optional contentMaking the
contentattribute non-optional with a default empty value eliminates the need for None checks and simplifies the API usage.src/draive/conversation/completion/default.py (1)
111-165: LGTM! The turn-based tool handling is well-implemented.The refactoring from
repetition_leveltotools_turnand the addition of extension accumulation logic are correctly implemented. The function properly handles direct completions, extensions, and regular tool responses.src/draive/tools/types.py (1)
21-90: Well-designed type system improvements.The changes enhance the tool system's flexibility and consistency:
- Expanded
ToolHandlingmodes ("auto", "extend", "direct", "spawn") provide richer interaction patterns- Renaming to
ToolErrorimproves naming consistency- Unified use of
MultimodalContentthroughout- Addition of
tools_turnparameter enables sophisticated availability controlsrc/draive/generation/text/default.py (1)
23-102: Clean refactoring with streaming removal.The simplified non-streaming implementation is cleaner and follows the same turn-based tool handling pattern as other modules. The accumulation of extensions and proper handling of direct completions is correctly implemented.
src/draive/stages/stage.py (2)
756-884: Excellent implementation of multiple tool handling modes.The enhanced
tool_callmethod properly supports all four handling modes with clear semantics:
"auto": Standard tool execution"direct": Tool result replaces stage result"extend": Tool result extends existing result"spawn": Asynchronous execution without waitingNote that the
"spawn"mode (lines 852-861) creates a detached task that may outlive the stage execution. Ensure this is the intended behavior and document any cleanup requirements.
1956-2021: Consistent implementation of turn-based tool handling.The
_lmm_completionfunction follows the same pattern as other modules with proper handling of completions, extensions, and regular responses. Debug logging provides good observability.src/draive/tools/function.py (4)
35-167: Consistent naming improvements throughout the class.The renaming from "failure" to "error" and "availability_check" to "availability" improves API consistency. All changes are properly propagated through constructors, methods, and type annotations.
168-184: Correct implementation of turn-based availability checking.The
available()method now properly accepts and forwards thetools_turnparameter to the availability checking function, enabling sophisticated turn-based tool availability logic.
186-254: Well-structured tool call implementation with proper error handling.The renamed
call()method includes excellent additions:
- Event reporting for observability (lines 193-200)
- Proper distinction between
Exception(formatted and wrapped) andBaseException(re-raised)- Comprehensive telemetry recording
The separation of
BaseExceptionhandling (lines 240-253) is particularly good as it ensures system exceptions likeKeyboardInterruptaren't caught and formatted as tool errors.
267-271: Helper functions correctly updated.The
_available()function signature and_default_error_format()naming are properly updated to match the new API design.Also applies to: 466-468
src/draive/tools/toolbox.py (3)
34-34: LGTM! The parameter renaming improves clarity.The change from
repeated_calls_limittotool_turns_limitbetter represents the concept of tool usage across multiple conversation turns rather than just repeated calls.Also applies to: 39-39, 46-46, 66-68, 93-93, 116-116, 122-122
199-202: LGTM! Method call updated correctly.The change from
tool.tool_calltotool.callaligns with the updatedToolprotocol.
227-248: LGTM! Tool handling modes properly implemented.The new handling logic correctly implements the extended tool response modes:
- "spawn" appropriately uses
ctx.spawn()for detached execution- Error handling is simplified and catches the specific
LMMToolError- The f-string usage is more idiomatic
Also applies to: 260-266, 270-270
KaQuMiQ
force-pushed
the
feature/tool_handling
branch
4 times, most recently
from
6e8fd56 to
cd2d567
Compare
There was a problem hiding this comment.
Actionable comments posted: 5
🔭 Outside diff range comments (3)
pyproject.toml (1)
48-52: Optional: pin the extra’ed dependency to the same base version for clarity.The
opentelemetryextra currently pulls"haiway[opentelemetry]"without a version spec, which resolves fine thanks to the base requirement (haiway~=0.27.2). Pinning it explicitly avoids accidental drifts if the base line is ever removed.- "haiway[opentelemetry]", + "haiway[opentelemetry]~=0.27.2",src/draive/generation/model/default.py (1)
27-36: Type inconsistency in decoder parameter.The function signature declares
decoder: ModelGeneratorDecoder | Nonewithout the generic type parameter, but line 121 uses it asdecoder(result)expecting it to returnGenerated. This should bedecoder: ModelGeneratorDecoder[Generated] | Noneto match the protocol definition.- decoder: ModelGeneratorDecoder | None, + decoder: ModelGeneratorDecoder[Generated] | None,Also applies to: 120-122
src/draive/generation/text/default.py (1)
23-30: Remove remaining AsyncIterable return types in text generation modulesThe streaming types weren’t fully cleaned up—
AsyncIterableis still imported and used in several text‐generation files. Please update these signatures to returnstronly and remove the unused import:• src/draive/generation/text/state.py
– At top: change
from collections.abc import AsyncIterable, Iterable
to
from collections.abc import Iterable
– Around line 26, update
) -> AsyncIterable[str] | str:
to
) -> str:• src/draive/generation/text/types.py
– RemoveAsyncIterablefrom the import on line 1.
– At line 22, change
) -> AsyncIterable[str] | str:
to
) -> str:• src/draive/generation/text/default.py
– RemoveAsyncIterablefrom the import on line 1.
– At line 30, change
) -> AsyncIterable[str] | str:
to
) -> str:Fix these to fully eliminate the deprecated streaming API.
♻️ Duplicate comments (3)
src/draive/generation/model/default.py (1)
82-84: Good fix for the uninitialized variable issue.The
resultvariable is now properly initialized withMultimodalContent.empty, addressing the previous review comment about potential access of uninitialized variable.src/draive/choice/default.py (1)
77-79: Properly initialized variables.Good that
resultis initialized withMultimodalContent.empty, addressing the previous review comment about uninitialized variable access.src/draive/lmm/types.py (1)
236-262: Consider documenting the aggregation logicThe
completion()method implements complex aggregation logic. Consider adding a docstring to clarify the precedence and merging behavior.def completion( self, *, extension: MultimodalContent | None = None, ) -> LMMCompletion | None: + """Aggregate completion and extension responses into a single completion. + + Extensions are prepended to completions in the final content. + Metadata is merged from all completion responses. + + Returns: + Aggregated completion or None if no completion responses exist. + """ completion_parts: list[MultimodalContent] = []
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (24)
pyproject.toml(1 hunks)src/draive/__init__.py(2 hunks)src/draive/anthropic/lmm_generation.py(1 hunks)src/draive/choice/default.py(2 hunks)src/draive/conversation/completion/default.py(3 hunks)src/draive/conversation/realtime/default.py(4 hunks)src/draive/generation/model/default.py(2 hunks)src/draive/generation/model/state.py(1 hunks)src/draive/generation/model/types.py(2 hunks)src/draive/generation/text/default.py(3 hunks)src/draive/generation/text/state.py(1 hunks)src/draive/generation/text/types.py(0 hunks)src/draive/lmm/types.py(3 hunks)src/draive/mcp/client.py(1 hunks)src/draive/mistral/lmm_generation.py(2 hunks)src/draive/ollama/lmm_generation.py(2 hunks)src/draive/openai/lmm_generation.py(2 hunks)src/draive/stages/stage.py(2 hunks)src/draive/tools/__init__.py(2 hunks)src/draive/tools/function.py(15 hunks)src/draive/tools/toolbox.py(10 hunks)src/draive/tools/types.py(4 hunks)src/draive/vllm/lmm_generation.py(2 hunks)tests/test_tool_call.py(2 hunks)
📓 Path-based instructions (3)
**/*.py
📄 CodeRabbit Inference Engine (CLAUDE.md)
**/*.py: Use absolute imports fromdraivepackage
Follow Ruff import ordering (standard library, third party, local)
Use Python 3.12+ type features (type unions with|, generic syntax)
Use base and abstract types likeSequenceorIterableinstead of concrete types
Use custom exceptions for specific errors
Use NumPy docstring convention for all functions, classes, and methods
Skip module-level docstrings unless explicitly requested
Include sections: Parameters, Returns, Raises, Notes (if needed) in docstrings
ALWAYS useSequence[T]instead oflist[T],Mapping[K,V]instead ofdict[K,V], andSet[T]instead ofset[T]for collections in State, Config and DataModel classes
Immutable updates to State, Config, and DataModel classes should be performed through the.updated()method
Files:
src/draive/ollama/lmm_generation.pytests/test_tool_call.pysrc/draive/mcp/client.pysrc/draive/__init__.pysrc/draive/tools/__init__.pysrc/draive/openai/lmm_generation.pysrc/draive/anthropic/lmm_generation.pysrc/draive/generation/model/state.pysrc/draive/mistral/lmm_generation.pysrc/draive/conversation/realtime/default.pysrc/draive/generation/model/types.pysrc/draive/generation/text/state.pysrc/draive/conversation/completion/default.pysrc/draive/generation/model/default.pysrc/draive/choice/default.pysrc/draive/vllm/lmm_generation.pysrc/draive/tools/types.pysrc/draive/lmm/types.pysrc/draive/generation/text/default.pysrc/draive/tools/function.pysrc/draive/stages/stage.pysrc/draive/tools/toolbox.py
tests/**/*.py
📄 CodeRabbit Inference Engine (CLAUDE.md)
tests/**/*.py: Tests should use pytest with async support and be located in thetests/directory
Mock dependencies within scope using stubbed functionality state in tests
Files:
tests/test_tool_call.py
**/__init__.py
📄 CodeRabbit Inference Engine (CLAUDE.md)
Put exported symbols into
__init__.py
Files:
src/draive/__init__.pysrc/draive/tools/__init__.py
🧠 Learnings (8)
src/draive/__init__.py (1)
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/*.py : Use absolute imports from draive package
src/draive/tools/__init__.py (1)
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/*.py : Use absolute imports from draive package
src/draive/conversation/realtime/default.py (1)
Learnt from: KaQuMiQ
PR: #338
File: src/draive/bedrock/lmm_generation.py:102-107
Timestamp: 2025-06-16T09:44:08.688Z
Learning: LMMInstruction is a type alias for str in the draive codebase, so LMMInstruction and str are equivalent types and can be used interchangeably.
src/draive/generation/text/state.py (1)
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/*.py : Use base and abstract types like Sequence or Iterable instead of concrete types
src/draive/conversation/completion/default.py (1)
Learnt from: KaQuMiQ
PR: #338
File: src/draive/bedrock/lmm_generation.py:102-107
Timestamp: 2025-06-16T09:44:08.688Z
Learning: LMMInstruction is a type alias for str in the draive codebase, so LMMInstruction and str are equivalent types and can be used interchangeably.
src/draive/tools/types.py (1)
Learnt from: KaQuMiQ
PR: #327
File: src/draive/helpers/instruction_preparation.py:28-34
Timestamp: 2025-05-28T17:41:57.460Z
Learning: The draive project uses and requires Python 3.12+, so PEP-695 generic syntax with square brackets (e.g., def func[T: Type]()) is valid and should be used instead of the older TypeVar approach.
src/draive/lmm/types.py (1)
Learnt from: KaQuMiQ
PR: #338
File: src/draive/bedrock/lmm_generation.py:102-107
Timestamp: 2025-06-16T09:44:08.688Z
Learning: LMMInstruction is a type alias for str in the draive codebase, so LMMInstruction and str are equivalent types and can be used interchangeably.
src/draive/generation/text/default.py (1)
Learnt from: KaQuMiQ
PR: #338
File: src/draive/bedrock/lmm_generation.py:102-107
Timestamp: 2025-06-16T09:44:08.688Z
Learning: LMMInstruction is a type alias for str in the draive codebase, so LMMInstruction and str are equivalent types and can be used interchangeably.
🧬 Code Graph Analysis (8)
src/draive/__init__.py (1)
src/draive/tools/types.py (1)
ToolError(24-31)
src/draive/tools/__init__.py (1)
src/draive/tools/types.py (1)
ToolError(24-31)
src/draive/openai/lmm_generation.py (2)
src/draive/lmm/types.py (10)
LMMToolRequests(293-310)of(85-108)of(144-153)of(164-173)of(195-211)of(222-231)of(271-285)of(295-306)of(320-331)of(386-396)src/draive/ollama/lmm_generation.py (3)
lmm_completion(41-51)lmm_completion(54-64)lmm_completion(66-175)
src/draive/generation/model/state.py (1)
src/draive/generation/model/types.py (1)
ModelGeneratorDecoder(17-21)
src/draive/conversation/completion/default.py (5)
src/draive/lmm/types.py (16)
LMMCompletion(162-179)completion(236-261)LMMToolRequests(293-310)LMMToolResponses(220-266)completion_extension(263-266)LMMToolRequest(269-290)LMMToolResponse(193-217)of(85-108)of(144-153)of(164-173)of(195-211)of(222-231)of(271-285)of(295-306)of(320-331)of(386-396)src/draive/multimodal/content.py (4)
MultimodalContent(24-236)appending(189-219)meta(151-164)of(28-50)src/draive/tools/toolbox.py (4)
available_tools(156-189)respond_all(207-217)respond(219-276)of(28-124)src/draive/conversation/types.py (5)
ConversationMessage(79-227)model(46-64)model(163-177)ConversationMessageChunk(23-76)of(232-249)src/draive/instructions/types.py (7)
Instruction(45-224)formatted(48-53)formatted(57-62)formatted(65-83)of(87-96)of(100-109)of(112-146)
src/draive/generation/model/default.py (5)
src/draive/parameters/model.py (2)
DataModel(386-756)from_json(723-738)src/draive/multimodal/content.py (2)
MultimodalContent(24-236)appending(189-219)src/draive/lmm/state.py (4)
LMM(27-84)completion(30-39)completion(43-52)completion(55-82)src/draive/lmm/types.py (5)
completion(236-261)LMMCompletion(162-179)LMMToolRequests(293-310)LMMToolResponses(220-266)completion_extension(263-266)src/draive/tools/toolbox.py (2)
available_tools(156-189)respond_all(207-217)
src/draive/generation/text/default.py (4)
src/draive/multimodal/content.py (3)
MultimodalContent(24-236)appending(189-219)to_str(176-187)src/draive/lmm/state.py (4)
LMM(27-84)completion(30-39)completion(43-52)completion(55-82)src/draive/lmm/types.py (5)
completion(236-261)LMMCompletion(162-179)LMMToolRequests(293-310)LMMToolResponses(220-266)completion_extension(263-266)src/draive/tools/toolbox.py (2)
available_tools(156-189)respond_all(207-217)
src/draive/tools/function.py (3)
src/draive/tools/types.py (8)
ToolAvailabilityChecking(35-40)ToolResultFormatting(44-48)ToolErrorFormatting(52-56)handling(77-77)meta(74-74)available(79-82)call(84-89)name(62-62)src/draive/multimodal/content.py (2)
meta(151-164)MultimodalContent(24-236)src/draive/utils/processing.py (3)
event(164-168)event(172-180)event(183-205)
💤 Files with no reviewable changes (1)
- src/draive/generation/text/types.py
🧰 Additional context used
📓 Path-based instructions (3)
**/*.py
📄 CodeRabbit Inference Engine (CLAUDE.md)
**/*.py: Use absolute imports fromdraivepackage
Follow Ruff import ordering (standard library, third party, local)
Use Python 3.12+ type features (type unions with|, generic syntax)
Use base and abstract types likeSequenceorIterableinstead of concrete types
Use custom exceptions for specific errors
Use NumPy docstring convention for all functions, classes, and methods
Skip module-level docstrings unless explicitly requested
Include sections: Parameters, Returns, Raises, Notes (if needed) in docstrings
ALWAYS useSequence[T]instead oflist[T],Mapping[K,V]instead ofdict[K,V], andSet[T]instead ofset[T]for collections in State, Config and DataModel classes
Immutable updates to State, Config, and DataModel classes should be performed through the.updated()method
Files:
src/draive/ollama/lmm_generation.pytests/test_tool_call.pysrc/draive/mcp/client.pysrc/draive/__init__.pysrc/draive/tools/__init__.pysrc/draive/openai/lmm_generation.pysrc/draive/anthropic/lmm_generation.pysrc/draive/generation/model/state.pysrc/draive/mistral/lmm_generation.pysrc/draive/conversation/realtime/default.pysrc/draive/generation/model/types.pysrc/draive/generation/text/state.pysrc/draive/conversation/completion/default.pysrc/draive/generation/model/default.pysrc/draive/choice/default.pysrc/draive/vllm/lmm_generation.pysrc/draive/tools/types.pysrc/draive/lmm/types.pysrc/draive/generation/text/default.pysrc/draive/tools/function.pysrc/draive/stages/stage.pysrc/draive/tools/toolbox.py
tests/**/*.py
📄 CodeRabbit Inference Engine (CLAUDE.md)
tests/**/*.py: Tests should use pytest with async support and be located in thetests/directory
Mock dependencies within scope using stubbed functionality state in tests
Files:
tests/test_tool_call.py
**/__init__.py
📄 CodeRabbit Inference Engine (CLAUDE.md)
Put exported symbols into
__init__.py
Files:
src/draive/__init__.pysrc/draive/tools/__init__.py
🧠 Learnings (8)
src/draive/__init__.py (1)
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/*.py : Use absolute imports from draive package
src/draive/tools/__init__.py (1)
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/*.py : Use absolute imports from draive package
src/draive/conversation/realtime/default.py (1)
Learnt from: KaQuMiQ
PR: #338
File: src/draive/bedrock/lmm_generation.py:102-107
Timestamp: 2025-06-16T09:44:08.688Z
Learning: LMMInstruction is a type alias for str in the draive codebase, so LMMInstruction and str are equivalent types and can be used interchangeably.
src/draive/generation/text/state.py (1)
Learnt from: CR
PR: miquido/draive#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-21T11:09:57.256Z
Learning: Applies to **/*.py : Use base and abstract types like Sequence or Iterable instead of concrete types
src/draive/conversation/completion/default.py (1)
Learnt from: KaQuMiQ
PR: #338
File: src/draive/bedrock/lmm_generation.py:102-107
Timestamp: 2025-06-16T09:44:08.688Z
Learning: LMMInstruction is a type alias for str in the draive codebase, so LMMInstruction and str are equivalent types and can be used interchangeably.
src/draive/tools/types.py (1)
Learnt from: KaQuMiQ
PR: #327
File: src/draive/helpers/instruction_preparation.py:28-34
Timestamp: 2025-05-28T17:41:57.460Z
Learning: The draive project uses and requires Python 3.12+, so PEP-695 generic syntax with square brackets (e.g., def func[T: Type]()) is valid and should be used instead of the older TypeVar approach.
src/draive/lmm/types.py (1)
Learnt from: KaQuMiQ
PR: #338
File: src/draive/bedrock/lmm_generation.py:102-107
Timestamp: 2025-06-16T09:44:08.688Z
Learning: LMMInstruction is a type alias for str in the draive codebase, so LMMInstruction and str are equivalent types and can be used interchangeably.
src/draive/generation/text/default.py (1)
Learnt from: KaQuMiQ
PR: #338
File: src/draive/bedrock/lmm_generation.py:102-107
Timestamp: 2025-06-16T09:44:08.688Z
Learning: LMMInstruction is a type alias for str in the draive codebase, so LMMInstruction and str are equivalent types and can be used interchangeably.
🧬 Code Graph Analysis (8)
src/draive/__init__.py (1)
src/draive/tools/types.py (1)
ToolError(24-31)
src/draive/tools/__init__.py (1)
src/draive/tools/types.py (1)
ToolError(24-31)
src/draive/openai/lmm_generation.py (2)
src/draive/lmm/types.py (10)
LMMToolRequests(293-310)of(85-108)of(144-153)of(164-173)of(195-211)of(222-231)of(271-285)of(295-306)of(320-331)of(386-396)src/draive/ollama/lmm_generation.py (3)
lmm_completion(41-51)lmm_completion(54-64)lmm_completion(66-175)
src/draive/generation/model/state.py (1)
src/draive/generation/model/types.py (1)
ModelGeneratorDecoder(17-21)
src/draive/conversation/completion/default.py (5)
src/draive/lmm/types.py (16)
LMMCompletion(162-179)completion(236-261)LMMToolRequests(293-310)LMMToolResponses(220-266)completion_extension(263-266)LMMToolRequest(269-290)LMMToolResponse(193-217)of(85-108)of(144-153)of(164-173)of(195-211)of(222-231)of(271-285)of(295-306)of(320-331)of(386-396)src/draive/multimodal/content.py (4)
MultimodalContent(24-236)appending(189-219)meta(151-164)of(28-50)src/draive/tools/toolbox.py (4)
available_tools(156-189)respond_all(207-217)respond(219-276)of(28-124)src/draive/conversation/types.py (5)
ConversationMessage(79-227)model(46-64)model(163-177)ConversationMessageChunk(23-76)of(232-249)src/draive/instructions/types.py (7)
Instruction(45-224)formatted(48-53)formatted(57-62)formatted(65-83)of(87-96)of(100-109)of(112-146)
src/draive/generation/model/default.py (5)
src/draive/parameters/model.py (2)
DataModel(386-756)from_json(723-738)src/draive/multimodal/content.py (2)
MultimodalContent(24-236)appending(189-219)src/draive/lmm/state.py (4)
LMM(27-84)completion(30-39)completion(43-52)completion(55-82)src/draive/lmm/types.py (5)
completion(236-261)LMMCompletion(162-179)LMMToolRequests(293-310)LMMToolResponses(220-266)completion_extension(263-266)src/draive/tools/toolbox.py (2)
available_tools(156-189)respond_all(207-217)
src/draive/generation/text/default.py (4)
src/draive/multimodal/content.py (3)
MultimodalContent(24-236)appending(189-219)to_str(176-187)src/draive/lmm/state.py (4)
LMM(27-84)completion(30-39)completion(43-52)completion(55-82)src/draive/lmm/types.py (5)
completion(236-261)LMMCompletion(162-179)LMMToolRequests(293-310)LMMToolResponses(220-266)completion_extension(263-266)src/draive/tools/toolbox.py (2)
available_tools(156-189)respond_all(207-217)
src/draive/tools/function.py (3)
src/draive/tools/types.py (8)
ToolAvailabilityChecking(35-40)ToolResultFormatting(44-48)ToolErrorFormatting(52-56)handling(77-77)meta(74-74)available(79-82)call(84-89)name(62-62)src/draive/multimodal/content.py (2)
meta(151-164)MultimodalContent(24-236)src/draive/utils/processing.py (3)
event(164-168)event(172-180)event(183-205)
🔇 Additional comments (35)
pyproject.toml (1)
24-24: Verified:haiway 0.27.2is available and supports Python 3.12
The PyPI release metadata includes wheels tagged for cp312/py3, so you can safely merge.src/draive/__init__.py (1)
213-213: LGTM! Consistent exception class renaming.The import and export of
ToolErrorcorrectly replacesToolException, aligning with the broader codebase refactoring where the exception class was renamed and now requires aMultimodalContentparameter.Also applies to: 376-376
tests/test_tool_call.py (1)
54-58: LGTM! Test method calls updated correctly.The test calls have been properly updated from
tool_calltocall, maintaining consistency with the renamed tool invocation interface while preserving all test logic and assertions.Also applies to: 75-79
src/draive/mcp/client.py (1)
603-605: LGTM! Parameter names updated to match new FunctionTool interface.The parameter renaming from
availability_checktoavailabilityandformat_failuretoformat_errorcorrectly aligns with the updatedFunctionToolconstructor signature while maintaining the same functionality.src/draive/tools/__init__.py (1)
7-7: LGTM! Exception class import and export updated consistently.The replacement of
ToolExceptionwithToolErrorin both import and export statements correctly reflects the exception class renaming and maintains proper module interface consistency.Also applies to: 18-18
src/draive/ollama/lmm_generation.py (1)
147-159: LGTM! Standardized LMMToolRequests instantiation using factory method.The refactoring to use
LMMToolRequests.of()instead of direct constructor call standardizes the creation pattern across the codebase while maintaining the same functionality. The tool requests list and content parameter are properly passed to the factory method.src/draive/anthropic/lmm_generation.py (1)
264-275: LGTM! Clean refactoring to factory method pattern.The change standardizes
LMMToolRequestsinstantiation by using theof()factory method instead of direct constructor calls. This improves consistency across LMM generation modules and makes the conditional check more explicit.src/draive/vllm/lmm_generation.py (1)
289-301: Consistent factory method adoption.This change aligns with the broader refactoring effort to standardize
LMMToolRequestsinstantiation across all LMM generation modules using theof()factory method.src/draive/mistral/lmm_generation.py (1)
269-281: Consistent refactoring to factory method.This change continues the pattern of standardizing
LMMToolRequestsinstantiation across the codebase using theof()factory method, maintaining consistency with other LMM generation modules.src/draive/openai/lmm_generation.py (1)
292-304: Completes consistent factory method adoption.This change rounds out the standardization effort across all LMM generation modules, ensuring consistent
LMMToolRequestsinstantiation using theof()factory method.src/draive/generation/model/state.py (1)
29-29: Improved type annotation with explicit generic constraint.The change makes the generic type relationship explicit by constraining the
decoderparameter toModelGeneratorDecoder[Generated], which improves type safety and aligns with the generic protocol definition insrc/draive/generation/model/types.py.src/draive/generation/model/types.py (1)
17-21: Good type safety improvements!The generic parameterization of
ModelGeneratorDecoderensures type consistency between the decoder output and the expected generated model type. This change will help catch type mismatches at compile time.Also applies to: 36-36
src/draive/conversation/realtime/default.py (1)
76-81: Good improvements to logging and control flow.The changes enhance observability:
- Warning logs clearly indicate unsupported tool handling modes in realtime
- Debug logs help trace various event types during realtime conversation
- Explicit
continuestatement clarifies the asynchronous handling patternAlso applies to: 102-103, 112-113, 120-121, 126-128, 140-140
src/draive/generation/model/default.py (1)
86-116: Well-structured tool response handling.The new approach properly accumulates partial results:
- Direct completions from tools break the loop immediately
- Extensions are accumulated in
result_extension- Debug logging provides good visibility into the flow
src/draive/choice/default.py (1)
77-111: Consistent implementation with model generation.The tool response handling pattern matches the implementation in
generate_model, maintaining consistency across the codebase:
- Accumulates extensions before final resolution
- Handles direct completions appropriately
- Good debug logging for traceability
src/draive/lmm/types.py (2)
190-190: LGTM! The extended response handling types enhance tool interaction flexibility.The addition of "completion", "extension", and "detached" handling modes provides more granular control over tool response processing, aligning well with the PR's objective to refine tool handling.
304-304: Good improvement: Consistent non-null content handling.Using
MultimodalContent.emptyas the default instead ofNoneeliminates the need for null checks and improves API consistency.Also applies to: 309-309
src/draive/tools/types.py (5)
21-21: LGTM! Extended tool handling modes enable richer invocation patterns.The addition of "extend" and "spawn" modes provides valuable flexibility for tool invocation strategies, supporting both result accumulation and detached execution patterns.
24-32: Good refactor: Consistent error handling with required content.Renaming to
ToolErrorfollows Python conventions better, and requiringMultimodalContentensures errors are always properly formatted for model consumption.
36-40: LGTM! Turn-based availability checking improves clarity.Using
tools_turninstead of repetition-based tracking provides clearer semantics for tool availability checks.
48-48: Good type unification: Consistent use of MultimodalContent.Standardizing on
MultimodalContentfor both result and error formatting improves API consistency.Also applies to: 56-56
79-89: Excellent protocol improvements for better tool lifecycle management.The changes enhance the Tool protocol:
- Dynamic availability checking with
tools_turnparameter- Cleaner method naming with
callinstead oftool_call- Consistent
MultimodalContentreturn typesrc/draive/conversation/completion/default.py (3)
1-3: LGTM! Appropriate imports for incremental streaming.The new imports support the improved streaming implementation with incremental tool response processing.
111-165: Excellent refactor: Turn-based tracking with proper extension accumulation.The synchronous completion function has been cleanly refactored to:
- Use turn-based tool tracking for better clarity
- Properly accumulate and prepend extensions to the final result
- Handle both direct completions and tool responses with clear control flow
- Add helpful debug logging throughout
The implementation correctly uses the new
completion()andcompletion_extension()methods fromLMMToolResponses.
259-294: Excellent incremental streaming design for tool responses.The implementation effectively uses
asyncio.waitwithFIRST_COMPLETEDto:
- Process tool responses as they complete rather than waiting for all
- Stream completions and extensions immediately to the output
- Properly track completion status to end streaming when appropriate
This design significantly improves user experience by reducing latency.
src/draive/generation/text/default.py (1)
57-102: Clean implementation with consistent tool handling pattern.The synchronous text generation implementation:
- Uses the same turn-based tracking pattern as conversation completion
- Properly accumulates extensions before appending to results
- Includes helpful debug logging
- Cleanly converts the final MultimodalContent to string
The removal of streaming significantly improves code clarity.
src/draive/tools/function.py (4)
25-25: Consistent terminology update: failure → error.The renaming from "failure" to "error" improves consistency across the codebase and aligns with Python conventions.
Also applies to: 42-44, 106-106
167-183: Good implementation of turn-based availability checking.The method correctly:
- Accepts the required
tools_turnparameter- Passes it to the availability check function
- Handles exceptions gracefully by returning False and logging the error
226-229: Error handling looks correct.The code properly uses
LMMToolError(imported fromdraive.lmm) to wrap tool execution errors with formatted content.
307-309: Thorough and consistent updates to tool decorator.All decorator overloads and the implementation have been properly updated with:
- Renamed
availabilityandformat_errorparameters- Updated default formatting functions
- Consistent parameter documentation
Also applies to: 360-361, 409-411, 424-425
src/draive/stages/stage.py (2)
804-878: Well-structured implementation of the new tool handling modesThe refactored
tool_callmethod correctly implements the four handling modes (auto,direct,extend,spawn) with clear separation of concerns. The error handling is properly unified, and the use ofctx.spawnfor detached execution is appropriate.
1966-2020: Excellent refactoring of the completion loop with result extension supportThe updated
_lmm_completionfunction properly handles the new tool turn semantics and result extension accumulation. The enhanced debug logging at key points will greatly aid in troubleshooting tool interactions.src/draive/tools/toolbox.py (3)
34-34: Consistent and semantically clearer parameter namingThe renaming from
repeated_calls_limit/repetition_leveltotool_turns_limit/tools_turnthroughout the file better represents the concept of conversation turns with tool interactions rather than simple repetition counting.Also applies to: 67-68, 93-94, 116-117, 122-123, 132-132, 147-147, 153-153, 159-167, 287-287, 316-316, 324-324
227-248: Well-implemented extended tool handling modesThe response handling correctly maps the four tool handling modes to their corresponding LMM response types. The spawn implementation properly uses
ctx.spawnfor fire-and-forget execution and returns an appropriate placeholder response.
260-276: Simplified and consistent error handlingThe error handling is now more straightforward with a unified "error" handling type for all error cases. The f-string usage for logging is more Pythonic, and the error messages are clearer with the "ERROR:" prefix.
KaQuMiQ
force-pushed
the
feature/tool_handling
branch
from
cd2d567 to
c76c72d
Compare
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
1 participant
No description provided.