chore(weave): Move accumulator exts into op #3874
Conversation
WalkthroughThis pull request updates several integration modules by replacing the use of the public Changes
Suggested reviewers
Poem
Tip ⚡🧪 Multi-step agentic review comment chat (experimental)
📜 Recent review detailsConfiguration used: .coderabbit.yaml 📒 Files selected for processing (2)
✅ Files skipped from review due to trivial changes (1)
🧰 Additional context used📓 Path-based instructions (1)`**/*.py`: Focus on pythonic code patterns. Check for proper...
⏰ Context from checks skipped due to timeout of 90000ms (875)
🔇 Additional comments (2)
✨ Finishing Touches
🪧 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
Documentation and Community
|
|
Preview this PR with FeatureBee: https://beta.wandb.ai/?betaVersion=5502206b46ff446d03f180d9e0d0b38cae6fb2fc |
![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 (1)
weave/trace/op.py (1)
939-951: Dynamic attribute delegation.
Using__getattr__to delegate to the wrapped iterator is a neat approach. Just ensure that any special attributes or dunder methods needed by advanced consumers of this wrapper are explicitly handled.Consider clarifying the docstring that only normal attributes are forwarded, so devs do not rely on unimplemented dunder methods from the underlying iterator.
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (16)
weave/integrations/anthropic/anthropic_sdk.py(4 hunks)weave/integrations/bedrock/bedrock_sdk.py(2 hunks)weave/integrations/cohere/cohere_sdk.py(3 hunks)weave/integrations/google_ai_studio/google_ai_studio_sdk.py(3 hunks)weave/integrations/groq/groq_sdk.py(2 hunks)weave/integrations/huggingface/huggingface_inference_client_sdk.py(3 hunks)weave/integrations/instructor/instructor_iterable_utils.py(3 hunks)weave/integrations/instructor/instructor_partial_utils.py(2 hunks)weave/integrations/langchain_nvidia_ai_endpoints/langchain_nv_ai_endpoints.py(2 hunks)weave/integrations/litellm/litellm.py(2 hunks)weave/integrations/mistral/v0/mistral.py(2 hunks)weave/integrations/mistral/v1/mistral.py(2 hunks)weave/integrations/openai/openai_sdk.py(3 hunks)weave/integrations/vertexai/vertexai_sdk.py(3 hunks)weave/trace/op.py(6 hunks)weave/trace/op_extensions/accumulator.py(0 hunks)
💤 Files with no reviewable changes (1)
- weave/trace/op_extensions/accumulator.py
🧰 Additional context used
📓 Path-based instructions (1)
`**/*.py`: Focus on pythonic code patterns. Check for proper...
**/*.py: Focus on pythonic code patterns.
Check for proper exception handling.
Verify type hints usage where applicable.
Look for potential performance improvements.
Don't comment on formatting if black/isort is configured.
Check for proper dependency injection patterns.
Verify proper async handling if applicable.
weave/integrations/mistral/v0/mistral.pyweave/integrations/openai/openai_sdk.pyweave/integrations/instructor/instructor_partial_utils.pyweave/integrations/litellm/litellm.pyweave/integrations/langchain_nvidia_ai_endpoints/langchain_nv_ai_endpoints.pyweave/integrations/groq/groq_sdk.pyweave/integrations/mistral/v1/mistral.pyweave/integrations/huggingface/huggingface_inference_client_sdk.pyweave/integrations/bedrock/bedrock_sdk.pyweave/integrations/anthropic/anthropic_sdk.pyweave/integrations/vertexai/vertexai_sdk.pyweave/integrations/instructor/instructor_iterable_utils.pyweave/integrations/google_ai_studio/google_ai_studio_sdk.pyweave/integrations/cohere/cohere_sdk.pyweave/trace/op.py
⏰ Context from checks skipped due to timeout of 90000ms (21)
- GitHub Check: Trace nox tests (3, 12, scorers)
- GitHub Check: Trace nox tests (3, 10, scorers)
- GitHub Check: Trace nox tests (3, 9, scorers)
- GitHub Check: Trace nox tests (3, 12, scorers)
- GitHub Check: Trace nox tests (3, 9, scorers)
- GitHub Check: Trace nox tests (3, 12, scorers)
- GitHub Check: Trace nox tests (3, 9, scorers)
- GitHub Check: Trace nox tests (3, 12, scorers)
- GitHub Check: Trace nox tests (3, 9, scorers)
- GitHub Check: Trace nox tests (3, 12, scorers)
- GitHub Check: Trace nox tests (3, 9, scorers)
- GitHub Check: Trace nox tests (3, 12, scorers)
- GitHub Check: Trace nox tests (3, 9, scorers)
- GitHub Check: Trace nox tests (3, 12, scorers)
- GitHub Check: Trace nox tests (3, 9, scorers)
- GitHub Check: Trace nox tests (3, 12, scorers)
- GitHub Check: Trace nox tests (3, 9, scorers)
- GitHub Check: Trace nox tests (3, 12, scorers)
- GitHub Check: Trace nox tests (3, 9, scorers)
- GitHub Check: Trace nox tests (3, 12, scorers)
- GitHub Check: Trace nox tests (3, 9, scorers)
🔇 Additional comments (52)
weave/integrations/mistral/v0/mistral.py (2)
9-9: Import updated to use private function from centralized location.The import has been changed to use
_add_accumulatorfromweave.trace.opinstead ofadd_accumulatorfromweave.trace.op_extensions.accumulator. This change centralizes accumulator functionality in the core op module.
84-84: Function call updated to match import change.The function call has been updated to use
_add_accumulatorinstead ofadd_accumulator, maintaining consistent behavior while using the newly centralized implementation.weave/integrations/instructor/instructor_partial_utils.py (2)
7-7: Import updated to use private function from centralized location.The import has been changed to use
_add_accumulatorfromweave.trace.opinstead of importing from the extensions directory, aligning with the refactoring to consolidate accumulator functionality.
22-26: Function call updated while preserving behavior.The function call has been updated to use
_add_accumulatorwhile maintaining the same parameter structure and behavior. The change is consistent with the import update.weave/integrations/google_ai_studio/google_ai_studio_sdk.py (3)
10-10: Import updated to use private function from centralized location.The import has been changed to use
_add_accumulatorfromweave.trace.opinstead of importing from the extensions directory, consistent with the refactoring pattern across multiple files.
103-108: Function call updated in synchronous wrapper.The function call has been updated to use
_add_accumulatorinstead ofadd_accumulator, maintaining the same parameters and behavior for the accumulator in the synchronous wrapper.
128-133: Function call updated in asynchronous wrapper.The function call has been updated to use
_add_accumulatorinstead ofadd_accumulator, maintaining the same parameters and behavior for the accumulator in the asynchronous wrapper.weave/integrations/cohere/cohere_sdk.py (3)
10-10: Import updated to use private function from centralized location.The import has been changed to use
_add_accumulatorfromweave.trace.opinstead of importing from the extensions directory, consistent with the refactoring pattern across integration modules.
170-170: Function call updated in stream wrapper.The function call has been updated to use
_add_accumulatorinstead ofadd_accumulatorwhile maintaining the same parameter pattern and behavior.
179-179: Function call updated in v2 stream wrapper.The function call has been updated to use
_add_accumulatorinstead ofadd_accumulatorwhile preserving the same behavior for the v2 accumulator.weave/integrations/litellm/litellm.py (2)
9-9: Import update looks good.The import has been updated to use
_add_accumulatorfromweave.trace.opinstead ofadd_accumulatorfrom the extensions module, which aligns with the PR objective of moving accumulator functionality into the main op module.
94-99: Function call updated correctly.The function call to
add_accumulatorhas been properly replaced with_add_accumulatorwhile maintaining all the same parameters and functionality. This ensures backward compatibility while supporting the architectural change.weave/integrations/openai/openai_sdk.py (3)
10-10: Import update looks good.Import statement correctly updated to include
_add_accumulatorfrom the centralizedweave.trace.opmodule. This follows the PR's objective of consolidating accumulator functionality.
334-341: Function call updated correctly in sync wrapper.The call to
add_accumulatorhas been properly replaced with_add_accumulatorin the synchronous wrapper implementation while maintaining the same parameter structure and functionality.
370-377: Function call updated correctly in async wrapper.The call to
add_accumulatorhas been properly replaced with_add_accumulatorin the asynchronous wrapper implementation. This maintains consistency between sync and async implementations.weave/integrations/huggingface/huggingface_inference_client_sdk.py (3)
8-8: Import update looks good.The import has been correctly updated to use
_add_accumulatorfrom the centralizedweave.trace.opmodule instead of the previous extension module location.
82-87: Function call updated correctly in sync wrapper.The call to
add_accumulatorhas been properly replaced with_add_accumulatorin the synchronous wrapper implementation while maintaining the same parameter structure and functionality.
106-111: Function call updated correctly in async wrapper.The call to
add_accumulatorhas been properly replaced with_add_accumulatorin the asynchronous wrapper implementation, maintaining consistency with the synchronous version.weave/integrations/bedrock/bedrock_sdk.py (2)
4-4: Import update looks good.Import statement properly updated to include both
_add_accumulatorand_IteratorWrapperfrom the centralizedweave.trace.opmodule, which aligns with the PR's objective of consolidating these utilities in one place.
136-141: Function call updated correctly.The function call has been properly updated from
add_accumulatorto_add_accumulatorwhile maintaining all the same parameters, including the customiterator_wrapperimplementation. This ensures backward compatibility while supporting the architectural change.weave/integrations/langchain_nvidia_ai_endpoints/langchain_nv_ai_endpoints.py (2)
17-17: Import updated to use private API.The import has been updated to use
_add_accumulatorfromweave.trace.opinstead of using the publicadd_accumulatorfromweave.trace.op_extensions.accumulator. This change is part of the effort to move accumulator functionality into the main op module.
164-169: Function call updated to use private API.The function call has been updated to use
_add_accumulatorinstead ofadd_accumulator, maintaining the same parameters and functionality. This change is consistent with the import change on line 17.weave/integrations/mistral/v1/mistral.py (2)
9-9: Import updated to use private API.The import has been updated to use
_add_accumulatorfromweave.trace.opinstead of the previous public API. This change aligns with the PR objective of moving accumulator extensions into the op module.
91-91: Function call updated to use private API.The function call has been updated to use
_add_accumulatorinstead ofadd_accumulator, maintaining the same functionality with identical parameters. This modification is consistent with the import change on line 9.weave/integrations/vertexai/vertexai_sdk.py (3)
10-10: Import updated to use private API.The import statement has been updated to use
_add_accumulatorfromweave.trace.opinstead of importing it from the previous location. This change is part of the effort to consolidate accumulator functionality into the op module.
98-103: Function call updated in sync wrapper.The
vertexai_wrapper_syncfunction now uses_add_accumulatorinstead ofadd_accumulatorwhile maintaining the same parameters and functionality. This change is consistent with the import update on line 10.
123-128: Function call updated in async wrapper.The
vertexai_wrapper_asyncfunction now uses_add_accumulatorinstead ofadd_accumulatorwhile maintaining the same parameters and functionality. This change ensures consistent usage of the private API across both synchronous and asynchronous wrappers.weave/integrations/anthropic/anthropic_sdk.py (4)
11-11: Import updated to use private APIs.The import has been updated to use both
_add_accumulatorand_IteratorWrapperfromweave.trace.opmodule instead of their previous locations. This change aligns with the PR objective of moving accumulator extensions into the op module.
80-84: Function call updated in sync wrapper.The
create_wrapper_syncfunction now uses_add_accumulatorinstead ofadd_accumulatorwhile maintaining the same parameters and functionality. This change is consistent with the import update on line 11.
104-108: Function call updated in async wrapper.The
create_wrapper_asyncfunction now uses_add_accumulatorinstead ofadd_accumulatorwhile maintaining the same parameters and functionality. This ensures consistent usage of the private API across different wrapper types.
173-178: Function call updated in stream wrapper.The
create_stream_wrapperfunction now uses_add_accumulatorinstead ofadd_accumulatorwhile maintaining the same parameters, including theAnthropicIteratorWrapper. This change completes the transition to using the private API throughout the file.weave/integrations/instructor/instructor_iterable_utils.py (3)
8-8: Confirm stability of private function import.
Replacingadd_accumulatorwith_add_accumulatoris consistent with the broader refactor in this PR. Note, however, that_add_accumulatoris a private or semi-private function; ensure it is officially supported for future maintenance.We can verify official usage support by checking the project’s documentation or release notes. Would you like to search for references in the repository or open an issue to confirm its stability?
35-39: No functional concerns.
Switching to_add_accumulatorretains the existing parameters (make_accumulator,should_accumulate) without changing the core logic. This is a straightforward swap that appears correct.
55-59: No functional concerns (async wrapper).
The asynchronous wrapper also appropriately transitions to_add_accumulator. The arguments and functionality remain aligned with the synchronous counterpart, so no issues noted.weave/integrations/groq/groq_sdk.py (2)
96-100: Straightforward accumulator replacement.
The call to_add_accumulatorwith the same parameters (make_accumulator,should_accumulate) is consistent with the pattern elsewhere in this PR, ensuring the refactor remains uniform.
9-9:✅ Verification successful
Ensure consistency with the new import.
The import line now pulls_add_accumulatorfromweave.trace.op. Verify that all references to_add_accumulatorelsewhere in this file and across the codebase match the updated import path.
🏁 Script executed:
#!/bin/bash # Check all references to `_add_accumulator` in the codebase to confirm consistency. rg '_add_accumulator'Length of output: 3533
Reference Consistency Verified
The grep output confirms that all references to
_add_accumulatorconsistently use the updated import (from weave.trace.op import _add_accumulator) both inweave/integrations/groq/groq_sdk.pyand across other modules in the codebase. No further changes are necessary.weave/trace/op.py (16)
5-5: New import for process lifecycle management.
Importingatexitcan help manage resources at interpreter shutdown. Ensure it does not introduce side effects in multi-threaded or long-running environments if used outside typical shutdown flows.
11-12: Weak reference & additional collection imports.
Addingweakrefand importing these abstract collection types (e.g.,AsyncIterator) broadens the module’s capabilities for generic, iterator-like patterns. No issues spotted in principle.
20-20: TypeVar imports for generic typing.
Defining generic typesSandVpaves the way for strongly typed accumulators and iterators. Looks good for maintainability.Also applies to: 24-24
45-46: Introduction of type variables S and V.
These type variables help enforce correct usage of accumulators and iterators, improving type safety. No concerns.
617-621: New optional parameter 'accumulator' and updated docstring.
The decorator now supports anaccumulatorparameter for enhanced iterator handling. The docstring clarifies the feature adequately. Validate that existing code paths with no accumulator supplied remain unaffected.Do you want to run integration tests or generate them to confirm backward compatibility and ensure no breakage for code not passing an accumulator?
632-636: Iterator detection logic.
The lines now explicitly distinguish generator functions from async generators. This combined check setsis_iteratoreffectively. The approach is correct and aligns well with Python’s introspection capabilities.
702-706: Runtime flags for op wrappers.
Marking_is_async,_is_iterator,_is_generator, and_is_async_generatorclarifies runtime type checks. This is a clean approach to unify handling of different function types.
710-711: Refactor step for wrapper creation.
Assigning the wrapped op towrapped_opappears consistent with the prior logic, preserving the function’s identity and allowing subsequent modifications. No concerns.
713-723: Conditional accumulator assignment for iterators.
This code conditionally applies the accumulator ifis_iteratoris true and an accumulator is provided. This is a neat extension that does not disrupt non-iterator paths. Looks valid.
804-812: New definitions for yield/error/close callbacks.
Introducing_OnYieldType,_OnErrorType,_OnCloseTypesets up strong type hints for iterator lifecycle hooks. The new constants clarify logging messages.
815-935: _IteratorWrapper(Generic[V])
This large class manages iterator lifecycle—including error handling, closure, and context management. Theatexitusage ensures a fallback closure, but confirm that it won’t lead to unexpected calls in embedded or partial interpreter kicks. Overall, the design is robust for capturing iteration outputs.
952-971: Context management (sync).
These methods properly manage errors and final close calls on synchronous iteration. The usage of_call_on_error_onceand_call_on_close_onceis consistent with the error-lifecycle pattern.
972-990: Context management (async).
The same pattern extends to async context management, ensuring symmetrical handling for asynchronous iterators. Implementation is aligned with best practices.
992-1015: _Accumulator(Generic[S, V])
A container for maintaining shared state across yielded values. It cleanly separates accumulator logic from the iteration mechanics, promoting modular design.
1017-1035: _build_iterator_from_accumulator_for_op(...)
Wraps the raw iterator to pipe its items through the accumulator. Theon_yieldcallback is particularly helpful, catching each item without halting normal iteration if issues occur. Nicely done.
1037-1094: _add_accumulator(...)
Provides an internal mechanism for hooking an accumulator into any op that returns an iterator. This final bridging function unifies the logic elegantly, using the new_IteratorWrapperand_Accumulatorclasses.
Summary by CodeRabbit
Refactor
New Features