Conversation
WalkthroughThis change introduces several enhancements and refactorings across the draive Python framework. A new documentation file, ✨ 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 (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🔭 Outside diff range comments (3)
src/draive/stages/stage.py (3)
1684-1688:⚠️ Potential issueSame
isinstanceunion issue inwith_volatile_tools_contextThe comprehension repeats the incorrect usage that raises
TypeError. Apply the same tuple-of-types pattern:- if not isinstance(element, LMMToolRequests | LMMToolResponses) + if not isinstance(element, (LMMToolRequests, LMMToolResponses))
1605-1611:⚠️ Potential issueFallback never triggers – wrong check for caught exceptions
if exc in exceptions:compares an instance (exc) with a set/tuple of types, so it is alwaysFalse.
Useisinstanceinstead:- if exc in exceptions: + if isinstance(exc, tuple(exceptions)):Without this fix the fallback stage is unreachable, defeating the purpose of the helper.
294-317: 🛠️ Refactor suggestionConsider validating recalled context ends with
LMMCompletion
memory_recall(replace|extend)blindly trusts the recalled context.
Many downstream assertions (assert ... isinstance(context[-1], LMMCompletion)) assume the invariant that the last element is a completion. If a faultyMemoryimplementation returns an invalid context, later stages will raise.Add a defensive check, or at least log a warning:
+ recalled = await memory.recall() + assert not recalled or isinstance(recalled[-1], LMMCompletion), ( + "Recalled context must end with LMMCompletion" + ) + return StageState( + context=recalled, + result=result, + )Also applies to: 299-304
🧹 Nitpick comments (4)
src/draive/tools/function.py (1)
114-133: Consider refactoring for cleaner code.The metadata assignment logic is correct but could be simplified for better readability.
Apply this diff to simplify the conditional logic:
- self.meta: Meta - if description: - object.__setattr__( - self, - "meta", - meta.updated( - kind="tool", - name=name, - description=description, - ), - ) - - else: - object.__setattr__( - self, - "meta", - meta.updated( - kind="tool", - name=name, - ), - ) + self.meta: Meta + meta_updates = {"kind": "tool", "name": name} + if description: + meta_updates["description"] = description + object.__setattr__( + self, + "meta", + meta.updated(**meta_updates), + )CLAUDE.md (1)
47-49: Add language specification to the code block.The static analysis tool correctly identified that the fenced code block should specify a language for proper syntax highlighting.
Apply this diff to specify the language:
-``` +```mermaid Generation APIs → Stages → LMM → Provider ImplementationAlternatively, if it's meant as plain text flow diagram: ```diff -``` +```text Generation APIs → Stages → LMM → Provider Implementation<details> <summary>🧰 Tools</summary> <details> <summary>🪛 markdownlint-cli2 (0.17.2)</summary> 47-47: Fenced code blocks should have a language specified null (MD040, fenced-code-language) </details> </details> </blockquote></details> <details> <summary>src/draive/stages/stage.py (1)</summary><blockquote> `1138-1147`: **Duplicate route keys silently overwrite earlier stages** When two routed stages share the same `.meta.name` (case-insensitive) the later one silently replaces the former in `routes`/`options`. This is hard to debug. Consider raising instead: ```diff - routes[key] = route.execution - options[key] = route.meta + if key in routes: + raise ValueError(f"Duplicate stage name '{key}' in Stage.router() options") + routes[key] = route.execution + options[key] = route.metasrc/draive/stages/types.py (1)
8-14: Minor nit:Mappingisn’t used
Mappingis imported but not referenced in this module (only in the other file).
Removing the unused import avoids linter warnings.-from collections.abc import Mapping, Sequence +from collections.abc import Sequence
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
Knowledge Base: Disabled due to Reviews > Disable Knowledge Base setting
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (10)
CLAUDE.md(1 hunks)pyproject.toml(1 hunks)src/draive/__init__.py(2 hunks)src/draive/commons/metadata.py(1 hunks)src/draive/stages/__init__.py(2 hunks)src/draive/stages/stage.py(42 hunks)src/draive/stages/types.py(5 hunks)src/draive/tools/function.py(3 hunks)src/draive/tools/toolbox.py(9 hunks)src/draive/tools/types.py(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (5)
src/draive/stages/__init__.py (1)
src/draive/stages/types.py (2)
StageRouting(208-234)StageState(28-45)
src/draive/__init__.py (1)
src/draive/stages/types.py (1)
StageState(28-45)
src/draive/tools/types.py (1)
src/draive/commons/metadata.py (1)
Meta(20-275)
src/draive/tools/function.py (5)
src/draive/commons/metadata.py (4)
description(121-127)updated(197-201)kind(79-85)name(100-106)src/draive/tools/types.py (4)
description(62-62)meta(68-68)name(59-59)available(71-71)src/draive/parameters/model.py (1)
updated(554-558)src/draive/prompts/template.py (1)
available(67-72)src/draive/resources/template.py (1)
available(79-84)
src/draive/stages/types.py (2)
src/draive/multimodal/content.py (1)
MultimodalContent(24-198)src/draive/commons/metadata.py (1)
Meta(20-275)
🪛 markdownlint-cli2 (0.17.2)
CLAUDE.md
47-47: Fenced code blocks should have a language specified
null
(MD040, fenced-code-language)
🔇 Additional comments (14)
pyproject.toml (1)
8-8: Version bump looks appropriate for the scope of changes.The patch version increment from 0.66.1 to 0.66.2 aligns with the introduction of new stage routing capabilities and the StageState type. Note that this PR also includes a breaking change to the
ToolAvailabilityCheckingprotocol, which technically might warrant a minor version bump according to semantic versioning principles.src/draive/stages/__init__.py (2)
9-10: LGTM! New stage routing types properly exposed.The addition of
StageRoutingandStageStateto the public API is well-structured and aligns with the stage routing functionality being introduced. TheStageStateNamedTuple provides a clean abstraction for stage execution results, whileStageRoutingdefines the protocol for routing logic.
22-23: Exports correctly added to maintain API consistency.The new types are properly included in the
__all__list, following the module's export conventions.src/draive/__init__.py (2)
192-192: LGTM! StageState properly imported for package-level exposure.The import of
StageStatefrom the stages module is appropriate for exposing this new core type at the package level.
337-337: StageState correctly added to public API exports.Including
StageStatein the__all__list makes this important new type easily accessible to users without requiring deep imports. This follows the established pattern for core framework types.src/draive/commons/metadata.py (1)
78-118: LGTM! Consistent implementation of new metadata properties.The
kindandnameproperties and their correspondingwith_*methods follow the established pattern used fordescriptionandtags. The implementation correctly:
- Uses pattern matching for safe type checking
- Maintains immutability by creating new instances
- Follows the same code structure as existing metadata fields
These additions provide solid foundation for the enhanced metadata usage across the draive framework.
src/draive/tools/function.py (3)
27-27: LGTM! Proper attribute declaration.Adding
descriptionto__slots__is correct for the new explicit description attribute.
74-79: LGTM! Explicit description attribute assignment.The explicit assignment of the description attribute follows the established pattern for other immutable attributes in this class.
138-138: LGTM! Correct meta parameter passing.The availability check now correctly receives the
metaparameter, aligning with the updatedToolAvailabilityCheckingprotocol.src/draive/tools/toolbox.py (5)
35-35: LGTM! Proper meta parameter addition.The
metaparameter is correctly added to support metadata propagation in toolbox creation.
40-125: LGTM! Comprehensive metadata handling.The metadata handling logic correctly covers all scenarios:
- New toolboxes: Use provided meta or default empty
- Extended toolboxes: Merge existing meta with new meta when provided
- Tool collections: Use provided meta or default empty
The implementation ensures metadata is consistently preserved and propagated across all toolbox creation paths.
136-151: LGTM! Consistent meta parameter in async factory method.The
fetchedmethod correctly accepts and propagates the meta parameter, maintaining consistency with theofmethod.
156-156: LGTM! Meta attribute declaration.The meta attribute is properly declared as part of the class state.
290-290: LGTM! Metadata preservation in instance methods.The
with_toolsandfilteredmethods correctly propagate the existing metadata, ensuring consistency across toolbox transformations.Also applies to: 319-319, 327-327
| def __call__( | ||
| self, | ||
| meta: Meta, | ||
| ) -> bool: ... |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Breaking change: Verify all implementations are updated.
The ToolAvailabilityChecking protocol signature has been changed to accept a meta: Meta parameter. This is a breaking change that will affect all existing implementations.
While this change enables more sophisticated metadata-driven availability checking, ensure that all implementations throughout the codebase have been updated to match the new signature.
🏁 Script executed:
#!/bin/bash
# Description: Find all implementations of ToolAvailabilityChecking to verify they've been updated
# Expected: All implementations should accept the meta parameter
# Search for class definitions that might implement ToolAvailabilityChecking
ast-grep --pattern 'class $_ {
$$$
def __call__(
self,
$$$
) -> bool:
$$$
$$$
}'
# Search for function definitions that might be used as ToolAvailabilityChecking
rg -A 10 "def.*__call__.*bool" --type py
# Search for any remaining old-style availability checking patterns
rg -A 5 "ToolAvailabilityChecking" --type pyLength of output: 6310
🏁 Script executed:
#!/bin/bash
# Inspect ToolAvailabilityChecking definition
echo "===== src/draive/tools/types.py (first 200 lines) ====="
sed -n '1,200p' src/draive/tools/types.py
# Context around availability_check parameter in function definitions
echo "===== availability_check contexts in src/draive/tools/function.py ====="
rg -n "availability_check" -C 3 src/draive/tools/function.py
# Context around default _check_availability assignment
echo "===== object.__setattr__ usage in src/draive/tools/function.py ====="
rg -n "object.__setattr__" -C 5 src/draive/tools/function.pyLength of output: 7227
Fix default ToolAvailabilityChecking signature in function.py.
The ToolAvailabilityChecking protocol now requires a meta: Meta parameter. While no custom implementations with the old signature were found, the default availability checker is still defined as lambda: True, which will break when called with a meta argument.
Please update the default in src/draive/tools/function.py:
--- a/src/draive/tools/function.py
+++ b/src/draive/tools/function.py
@@ -92,7 +92,7 @@
object.__setattr__(
self,
"_check_availability",
- availability_check
- or (
- lambda: True # available by default
- ),
+ availability_check
+ or (
+ lambda meta: True # available by default
+ ),
)- No other
ToolAvailabilityCheckingimplementations with the old signature were detected. ResourceTemplateandPromptTemplate’s__call__methods are unrelated to this protocol.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def __call__( | |
| self, | |
| meta: Meta, | |
| ) -> bool: ... | |
| object.__setattr__( | |
| self, | |
| "_check_availability", | |
| availability_check | |
| or ( | |
| lambda meta: True # available by default | |
| ), | |
| ) |
🤖 Prompt for AI Agents
In src/draive/tools/function.py around the default availability checker
definition, update the default lambda from taking no parameters to accepting a
single meta parameter (meta: Meta) to match the new ToolAvailabilityChecking
protocol signature. Change the default from lambda: True to lambda meta: True so
it can be called with the meta argument without breaking.
| return StageState( | ||
| context=tuple( | ||
| element | ||
| for element in context | ||
| if not isinstance(element, LMMToolRequests | LMMToolResponses) | ||
| ), |
There was a problem hiding this comment.
isinstance with PEP-604 union causes a runtime TypeError
isinstance(element, LMMToolRequests | LMMToolResponses) is invalid (PEP-604 unions are not accepted by isinstance / issubclass).
At runtime Python raises: TypeError: isinstance() argument 2 cannot be a union type, which would break strip_context_tools() on the first call.
- if not isinstance(element, LMMToolRequests | LMMToolResponses)
+ if not isinstance(element, (LMMToolRequests, LMMToolResponses))📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| return StageState( | |
| context=tuple( | |
| element | |
| for element in context | |
| if not isinstance(element, LMMToolRequests | LMMToolResponses) | |
| ), | |
| return StageState( | |
| context=tuple( | |
| element | |
| for element in context | |
| if not isinstance(element, (LMMToolRequests, LMMToolResponses)) | |
| ), |
🤖 Prompt for AI Agents
In src/draive/stages/stage.py around lines 779 to 784, the use of isinstance
with a PEP-604 union type causes a runtime TypeError because isinstance does not
accept union types. To fix this, replace the union in isinstance with a tuple of
the classes, i.e., use isinstance(element, (LMMToolRequests, LMMToolResponses))
instead of isinstance(element, LMMToolRequests | LMMToolResponses).
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
src/draive/stages/stage.py (1)
783-783:⚠️ Potential issueFix isinstance with PEP-604 union type (duplicate issue).
The use of
isinstance(element, LMMToolRequests | LMMToolResponses)causes a runtimeTypeErrorbecauseisinstancedoesn't accept PEP-604 union types. This is the same issue identified in the previous review.Apply this fix:
- if not isinstance(element, LMMToolRequests | LMMToolResponses) + if not isinstance(element, (LMMToolRequests, LMMToolResponses))
🧹 Nitpick comments (1)
CLAUDE.md (1)
46-48: Fix markdown linting issue by specifying code block language.The fenced code block lacks a language specification, which triggers a markdown linting warning.
Since this appears to be a simple text flow diagram, specify it as text:
-``` +```text Generation APIs → Stages → LMM → Provider Implementation<details> <summary>🧰 Tools</summary> <details> <summary>🪛 markdownlint-cli2 (0.17.2)</summary> 46-46: Fenced code blocks should have a language specified null (MD040, fenced-code-language) </details> </details> </blockquote></details> </blockquote></details> <details> <summary>📜 Review details</summary> **Configuration used: CodeRabbit UI** **Review profile: CHILL** **Plan: Pro** **Knowledge Base: Disabled due to Reviews > Disable Knowledge Base setting** <details> <summary>📥 Commits</summary> Reviewing files that changed from the base of the PR and between 842a7e1e9c99955038b27e9a3062e28c4d822a93 and bba6679c415a92f57ff9b4e88b64f187df39ce2e. </details> <details> <summary>⛔ Files ignored due to path filters (1)</summary> * `uv.lock` is excluded by `!**/*.lock` </details> <details> <summary>📒 Files selected for processing (11)</summary> * `CLAUDE.md` (1 hunks) * `pyproject.toml` (1 hunks) * `src/draive/__init__.py` (2 hunks) * `src/draive/commons/metadata.py` (1 hunks) * `src/draive/prompts/template.py` (1 hunks) * `src/draive/stages/__init__.py` (2 hunks) * `src/draive/stages/stage.py` (42 hunks) * `src/draive/stages/types.py` (5 hunks) * `src/draive/tools/function.py` (3 hunks) * `src/draive/tools/toolbox.py` (9 hunks) * `src/draive/tools/types.py` (1 hunks) </details> <details> <summary>🚧 Files skipped from review as they are similar to previous changes (7)</summary> * src/draive/stages/__init__.py * src/draive/tools/types.py * src/draive/__init__.py * src/draive/tools/toolbox.py * pyproject.toml * src/draive/tools/function.py * src/draive/commons/metadata.py </details> <details> <summary>🧰 Additional context used</summary> <details> <summary>🧬 Code Graph Analysis (2)</summary> <details> <summary>src/draive/prompts/template.py (2)</summary><blockquote> <details> <summary>src/draive/tools/types.py (1)</summary> * `meta` (68-68) </details> <details> <summary>src/draive/multimodal/content.py (1)</summary> * `meta` (125-138) </details> </blockquote></details> <details> <summary>src/draive/stages/types.py (2)</summary><blockquote> <details> <summary>src/draive/multimodal/content.py (1)</summary> * `MultimodalContent` (24-198) </details> <details> <summary>src/draive/commons/metadata.py (1)</summary> * `Meta` (20-275) </details> </blockquote></details> </details><details> <summary>🪛 markdownlint-cli2 (0.17.2)</summary> <details> <summary>CLAUDE.md</summary> 46-46: Fenced code blocks should have a language specified null (MD040, fenced-code-language) </details> </details> </details> <details> <summary>🔇 Additional comments (6)</summary><blockquote> <details> <summary>src/draive/stages/types.py (3)</summary> `21-39`: **Excellent abstraction design.** The `StageState` named tuple is a well-designed abstraction that encapsulates the dual nature of stage transformations. This improves type safety and makes the code more readable compared to raw tuples. --- `200-228`: **Well-designed routing protocol.** The `StageRouting` protocol provides a clean interface for dynamic stage selection. The design with context, result, and options mapping as inputs and string key as output enables flexible routing implementations. --- `249-255`: **Good enhancement to exception handling.** Adding the optional `execution_result` parameter to `StageException` allows exceptions to carry partial results, which is valuable for debugging and error recovery scenarios. </details> <details> <summary>src/draive/stages/stage.py (3)</summary> `78-133`: **Well-implemented LLM-based routing function.** The `_lmm_routing` function provides intelligent stage selection using LLM analysis. The implementation correctly: - Formats options for LLM consumption - Uses XML tags for structured response parsing - Validates selections against available options - Provides clear error messages The approach of wrapping the current result in `<RESULT>` tags and asking the LLM to respond with `<SELECTION>` tags is a solid pattern for structured LLM interactions. --- `1060-1172`: **Excellent routing stage implementation.** The `Stage.router` class method is well-designed with: - Proper validation of required metadata (name and description) - Smart handling of single-option case (no routing needed) - Normalization of stage names to lowercase keys - Clear error handling for invalid selections The requirement for stages to have name and description metadata makes sense for routing decisions. --- `1200-1245`: **Essential metadata management method.** The `with_meta` method provides a clean way to update stage metadata, which is crucial for the routing mechanism. The implementation correctly creates a new instance with merged metadata while preserving immutability. </details> </blockquote></details> </details> <!-- This is an auto-generated comment by CodeRabbit for review status -->
| availability_check | ||
| or ( | ||
| lambda: True # available by default | ||
| lambda meta: True # available by default |
There was a problem hiding this comment.
Fix the protocol inconsistency causing a runtime error.
The default lambda now expects a meta parameter, but the PromptAvailabilityCheck protocol (lines 18-19) still defines functions as taking no parameters, and the available property (line 69) calls _check_availability() with no arguments. This will cause a TypeError at runtime.
To fix this inconsistency, you need to update both the protocol definition and the function call:
class PromptAvailabilityCheck(Protocol):
- def __call__(self) -> bool: ...
+ def __call__(self, meta: Meta) -> bool: ...And update the available property to pass the metadata:
@property
def available(self) -> bool:
try:
- return self._check_availability()
+ return self._check_availability(self.declaration.meta)
except Exception:
return False📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| lambda meta: True # available by default | |
| class PromptAvailabilityCheck(Protocol): | |
| - def __call__(self) -> bool: ... | |
| + def __call__(self, meta: Meta) -> bool: ... | |
| @property | |
| def available(self) -> bool: | |
| try: | |
| - return self._check_availability() | |
| + return self._check_availability(self.declaration.meta) | |
| except Exception: | |
| return False |
🤖 Prompt for AI Agents
In src/draive/prompts/template.py around line 62, the default lambda function
expects a 'meta' parameter, but the PromptAvailabilityCheck protocol (lines
18-19) defines functions with no parameters, and the 'available' property (line
69) calls _check_availability() without arguments, causing a runtime TypeError.
To fix this, update the PromptAvailabilityCheck protocol to define the function
as accepting a 'meta' parameter, and modify the 'available' property to call
_check_availability() with the appropriate metadata argument.
1 participant
No description provided.