test

Author: andrewtruong

weave/integrations/anthropic/anthropic_sdk.py

@@ -8,7 +8,7 @@
 import weave
 from weave.integrations.patcher import MultiPatcher, NoOpPatcher, SymbolPatcher
 from weave.trace.autopatch import IntegrationSettings, OpSettings
-from weave.trace.op import _IteratorWrapper, add_accumulator
+from weave.trace.op import _add_accumulator, _IteratorWrapper
 
 if TYPE_CHECKING:
     from anthropic.lib.streaming import MessageStream
@@ -77,7 +77,7 @@ def wrapper(fn: Callable) -> Callable:
         "We need to do this so we can check if `stream` is used"
         op_kwargs = settings.model_dump()
         op = weave.op(fn, **op_kwargs)
-        return add_accumulator(
+        return _add_accumulator(
             op,  # type: ignore
             make_accumulator=lambda inputs: anthropic_accumulator,
             should_accumulate=should_use_accumulator,
@@ -101,7 +101,7 @@ async def _async_wrapper(*args: Any, **kwargs: Any) -> Any:
         "We need to do this so we can check if `stream` is used"
         op_kwargs = settings.model_dump()
         op = weave.op(_fn_wrapper(fn), **op_kwargs)
-        return add_accumulator(
+        return _add_accumulator(
             op,  # type: ignore
             make_accumulator=lambda inputs: anthropic_accumulator,
             should_accumulate=should_use_accumulator,
@@ -170,7 +170,7 @@ def create_stream_wrapper(settings: OpSettings) -> Callable[[Callable], Callable
     def wrapper(fn: Callable) -> Callable:
         op_kwargs = settings.model_dump()
         op = weave.op(fn, **op_kwargs)
-        return add_accumulator(
+        return _add_accumulator(
             op,  # type: ignore
             make_accumulator=lambda _: anthropic_stream_accumulator,
             should_accumulate=lambda _: True,

weave/integrations/bedrock/bedrock_sdk.py

@@ -1,7 +1,7 @@
 from typing import TYPE_CHECKING, Any, Callable, Optional
 
 import weave
-from weave.trace.op import _IteratorWrapper, add_accumulator
+from weave.trace.op import _add_accumulator, _IteratorWrapper
 from weave.trace.weave_client import Call
 
 if TYPE_CHECKING:
@@ -133,7 +133,7 @@ def get(self, key: str, default: Any = None) -> Any:
                         )
                     return self
 
-        return add_accumulator(
+        return _add_accumulator(
             op,
             make_accumulator=lambda _: bedrock_stream_accumulator,
             should_accumulate=lambda _: True,

weave/integrations/cohere/cohere_sdk.py

@@ -7,7 +7,7 @@
 import weave
 from weave.integrations.patcher import MultiPatcher, NoOpPatcher, SymbolPatcher
 from weave.trace.autopatch import IntegrationSettings, OpSettings
-from weave.trace.op import add_accumulator
+from weave.trace.op import _add_accumulator
 
 if TYPE_CHECKING:
     from cohere.types.non_streamed_chat_response import NonStreamedChatResponse
@@ -167,7 +167,7 @@ def cohere_stream_wrapper(settings: OpSettings) -> Callable:
     def wrapper(fn: Callable) -> Callable:
         op_kwargs = settings.model_dump()
         op = weave.op(fn, **op_kwargs)
-        return add_accumulator(op, lambda inputs: cohere_accumulator)
+        return _add_accumulator(op, lambda inputs: cohere_accumulator)
 
     return wrapper
 
@@ -176,7 +176,7 @@ def cohere_stream_wrapper_v2(settings: OpSettings) -> Callable:
     def wrapper(fn: Callable) -> Callable:
         op_kwargs = settings.model_dump()
         op = weave.op(fn, **op_kwargs)
-        return add_accumulator(op, lambda inputs: cohere_accumulator_v2)
+        return _add_accumulator(op, lambda inputs: cohere_accumulator_v2)
 
     return wrapper
 

weave/integrations/google_ai_studio/google_ai_studio_sdk.py

@@ -7,7 +7,7 @@
 import weave
 from weave.integrations.patcher import MultiPatcher, NoOpPatcher, SymbolPatcher
 from weave.trace.autopatch import IntegrationSettings, OpSettings
-from weave.trace.op import add_accumulator
+from weave.trace.op import _add_accumulator
 from weave.trace.serialization.serialize import dictify
 from weave.trace.weave_client import Call
 
@@ -100,7 +100,7 @@ def wrapper(fn: Callable) -> Callable:
 
         op = weave.op(fn, **op_kwargs)
         op._set_on_finish_handler(gemini_on_finish)
-        return add_accumulator(
+        return _add_accumulator(
             op,  # type: ignore
             make_accumulator=lambda inputs: gemini_accumulator,
             should_accumulate=lambda inputs: isinstance(inputs, dict)
@@ -125,7 +125,7 @@ async def _async_wrapper(*args: Any, **kwargs: Any) -> Any:
 
         op = weave.op(_fn_wrapper(fn), **op_kwargs)
         op._set_on_finish_handler(gemini_on_finish)
-        return add_accumulator(
+        return _add_accumulator(
             op,  # type: ignore
             make_accumulator=lambda inputs: gemini_accumulator,
             should_accumulate=lambda inputs: isinstance(inputs, dict)

weave/integrations/groq/groq_sdk.py

@@ -6,7 +6,7 @@
 import weave
 from weave.integrations.patcher import MultiPatcher, NoOpPatcher, SymbolPatcher
 from weave.trace.autopatch import IntegrationSettings, OpSettings
-from weave.trace.op import add_accumulator
+from weave.trace.op import _add_accumulator
 
 if TYPE_CHECKING:
     from groq.types.chat import ChatCompletion, ChatCompletionChunk
@@ -93,7 +93,7 @@ def groq_wrapper(settings: OpSettings) -> Callable[[Callable], Callable]:
     def wrapper(fn: Callable) -> Callable:
         op_kwargs = settings.model_dump()
         op = weave.op(fn, **op_kwargs)
-        return add_accumulator(
+        return _add_accumulator(
             op,  # type: ignore
             make_accumulator=lambda inputs: groq_accumulator,
             should_accumulate=should_use_accumulator,

weave/integrations/huggingface/huggingface_inference_client_sdk.py

@@ -5,7 +5,7 @@
 import weave
 from weave.integrations.patcher import MultiPatcher, NoOpPatcher, SymbolPatcher
 from weave.trace.autopatch import IntegrationSettings, OpSettings
-from weave.trace.op import add_accumulator
+from weave.trace.op import _add_accumulator
 from weave.trace.serialization.serialize import dictify
 
 if TYPE_CHECKING:
@@ -79,7 +79,7 @@ def wrapper(fn: Callable) -> Callable:
             op_kwargs["postprocess_inputs"] = huggingface_postprocess_inputs
 
         op = weave.op(fn, **op_kwargs)
-        return add_accumulator(
+        return _add_accumulator(
             op,  # type: ignore
             make_accumulator=lambda inputs: huggingface_accumulator,
             should_accumulate=lambda inputs: isinstance(inputs, dict)
@@ -103,7 +103,7 @@ async def _async_wrapper(*args: Any, **kwargs: Any) -> Any:
             op_kwargs["postprocess_inputs"] = huggingface_postprocess_inputs
 
         op = weave.op(_fn_wrapper(fn), **op_kwargs)
-        return add_accumulator(
+        return _add_accumulator(
             op,  # type: ignore
             make_accumulator=lambda inputs: huggingface_accumulator,
             should_accumulate=lambda inputs: isinstance(inputs, dict)

weave/integrations/instructor/instructor_iterable_utils.py

@@ -5,7 +5,7 @@
 
 import weave
 from weave.trace.autopatch import OpSettings
-from weave.trace.op import add_accumulator
+from weave.trace.op import _add_accumulator
 
 
 def instructor_iterable_accumulator(
@@ -32,7 +32,7 @@ def instructor_wrapper_sync(settings: OpSettings) -> Callable[[Callable], Callab
     def wrapper(fn: Callable) -> Callable:
         op_kwargs = settings.model_dump()
         op = weave.op(fn, **op_kwargs)
-        return add_accumulator(
+        return _add_accumulator(
             op,  # type: ignore
             make_accumulator=lambda inputs: instructor_iterable_accumulator,
             should_accumulate=should_accumulate_iterable,
@@ -52,7 +52,7 @@ async def _async_wrapper(*args: Any, **kwargs: Any) -> Any:
 
         op_kwargs = settings.model_dump()
         op = weave.op(_fn_wrapper(fn), **op_kwargs)
-        return add_accumulator(
+        return _add_accumulator(
             op,  # type: ignore
             make_accumulator=lambda inputs: instructor_iterable_accumulator,
             should_accumulate=should_accumulate_iterable,

weave/integrations/instructor/instructor_partial_utils.py

@@ -4,7 +4,7 @@
 
 import weave
 from weave.trace.autopatch import OpSettings
-from weave.trace.op import add_accumulator
+from weave.trace.op import _add_accumulator
 
 
 def instructor_partial_accumulator(
@@ -19,7 +19,7 @@ def instructor_wrapper_partial(settings: OpSettings) -> Callable[[Callable], Cal
     def wrapper(fn: Callable) -> Callable:
         op_kwargs = settings.model_dump()
         op = weave.op(fn, **op_kwargs)
-        return add_accumulator(
+        return _add_accumulator(
             op,  # type: ignore
             make_accumulator=lambda inputs: instructor_partial_accumulator,
             should_accumulate=lambda inputs: True,

weave/integrations/langchain_nvidia_ai_endpoints/langchain_nv_ai_endpoints.py

@@ -14,7 +14,7 @@
 import weave
 from weave.integrations.patcher import MultiPatcher, NoOpPatcher, SymbolPatcher
 from weave.trace.autopatch import IntegrationSettings, OpSettings
-from weave.trace.op import Op, ProcessedInputs, add_accumulator
+from weave.trace.op import Op, ProcessedInputs, _add_accumulator
 
 _lc_nvidia_patcher: MultiPatcher | None = None
 
@@ -161,7 +161,7 @@ def wrapper(fn: Callable) -> Callable:
         op_kwargs = settings.model_dump()
         op = weave.op(fn, **op_kwargs)
         op._set_on_input_handler(postprocess_inputs_to_openai_format)
-        return add_accumulator(
+        return _add_accumulator(
             op,
             make_accumulator=lambda inputs: nvidia_accumulator,
             should_accumulate=should_use_accumulator,

weave/integrations/litellm/litellm.py

@@ -6,7 +6,7 @@
 import weave
 from weave.integrations.patcher import MultiPatcher, NoOpPatcher, SymbolPatcher
 from weave.trace.autopatch import IntegrationSettings, OpSettings
-from weave.trace.op import add_accumulator
+from weave.trace.op import _add_accumulator
 
 if TYPE_CHECKING:
     from litellm.utils import ModelResponse
@@ -91,7 +91,7 @@ def make_wrapper(settings: OpSettings) -> Callable:
     def litellm_wrapper(fn: Callable) -> Callable:
         op_kwargs = settings.model_dump()
         op = weave.op(fn, **op_kwargs)
-        return add_accumulator(
+        return _add_accumulator(
             op,  # type: ignore
             make_accumulator=lambda inputs: litellm_accumulator,
             should_accumulate=should_use_accumulator,

weave/integrations/mistral/v0/mistral.py

@@ -6,7 +6,7 @@
 import weave
 from weave.integrations.patcher import MultiPatcher, NoOpPatcher, SymbolPatcher
 from weave.trace.autopatch import IntegrationSettings, OpSettings
-from weave.trace.op import add_accumulator
+from weave.trace.op import _add_accumulator
 
 if TYPE_CHECKING:
     from mistralai.models.chat_completion import (
@@ -81,7 +81,7 @@ def mistral_stream_wrapper(settings: OpSettings) -> Callable:
     def wrapper(fn: Callable) -> Callable:
         op_kwargs = settings.model_dump()
         op = weave.op(fn, **op_kwargs)
-        acc_op = add_accumulator(op, lambda inputs: mistral_accumulator)  # type: ignore
+        acc_op = _add_accumulator(op, lambda inputs: mistral_accumulator)  # type: ignore
         return acc_op
 
     return wrapper

weave/integrations/mistral/v1/mistral.py

@@ -6,7 +6,7 @@
 import weave
 from weave.integrations.patcher import MultiPatcher, NoOpPatcher, SymbolPatcher
 from weave.trace.autopatch import IntegrationSettings, OpSettings
-from weave.trace.op import add_accumulator
+from weave.trace.op import _add_accumulator
 
 if TYPE_CHECKING:
     from mistralai.models import (
@@ -88,7 +88,7 @@ def mistral_stream_wrapper(settings: OpSettings) -> Callable:
     def wrapper(fn: Callable) -> Callable:
         op_kwargs = settings.model_dump()
         op = weave.op(fn, **op_kwargs)
-        acc_op = add_accumulator(op, lambda inputs: mistral_accumulator)  # type: ignore
+        acc_op = _add_accumulator(op, lambda inputs: mistral_accumulator)  # type: ignore
         return acc_op
 
     return wrapper

weave/integrations/openai/openai_sdk.py

@@ -7,7 +7,7 @@
 import weave
 from weave.integrations.patcher import MultiPatcher, NoOpPatcher, SymbolPatcher
 from weave.trace.autopatch import IntegrationSettings, OpSettings
-from weave.trace.op import Op, ProcessedInputs, add_accumulator
+from weave.trace.op import Op, ProcessedInputs, _add_accumulator
 
 if TYPE_CHECKING:
     from openai.types.chat import ChatCompletionChunk
@@ -331,7 +331,7 @@ def _openai_stream_options_is_set(inputs: dict) -> bool:
         op_kwargs = settings.model_dump()
         op = weave.op(_add_stream_options(fn), **op_kwargs)
         op._set_on_input_handler(openai_on_input_handler)
-        return add_accumulator(
+        return _add_accumulator(
             op,  # type: ignore
             make_accumulator=lambda inputs: lambda acc, value: openai_accumulator(
                 acc, value, skip_last=not _openai_stream_options_is_set(inputs)
@@ -367,7 +367,7 @@ def _openai_stream_options_is_set(inputs: dict) -> bool:
         op_kwargs = settings.model_dump()
         op = weave.op(_add_stream_options(fn), **op_kwargs)
         op._set_on_input_handler(openai_on_input_handler)
-        return add_accumulator(
+        return _add_accumulator(
             op,  # type: ignore
             make_accumulator=lambda inputs: lambda acc, value: openai_accumulator(
                 acc, value, skip_last=not _openai_stream_options_is_set(inputs)

weave/integrations/vertexai/vertexai_sdk.py

@@ -7,7 +7,7 @@
 import weave
 from weave.integrations.patcher import MultiPatcher, NoOpPatcher, SymbolPatcher
 from weave.trace.autopatch import IntegrationSettings, OpSettings
-from weave.trace.op import add_accumulator
+from weave.trace.op import _add_accumulator
 from weave.trace.serialization.serialize import dictify
 from weave.trace.weave_client import Call
 
@@ -95,7 +95,7 @@ def wrapper(fn: Callable) -> Callable:
 
         op = weave.op(fn, **op_kwargs)
         op._set_on_finish_handler(vertexai_on_finish)
-        return add_accumulator(
+        return _add_accumulator(
             op,  # type: ignore
             make_accumulator=lambda inputs: vertexai_accumulator,
             should_accumulate=lambda inputs: isinstance(inputs, dict)
@@ -120,7 +120,7 @@ async def _async_wrapper(*args: Any, **kwargs: Any) -> Any:
 
         op = weave.op(_fn_wrapper(fn), **op_kwargs)
         op._set_on_finish_handler(vertexai_on_finish)
-        return add_accumulator(
+        return _add_accumulator(
             op,  # type: ignore
             make_accumulator=lambda inputs: vertexai_accumulator,
             should_accumulate=lambda inputs: isinstance(inputs, dict)

weave/trace/op.py

@@ -718,7 +718,7 @@ def make_accumulator(inputs: dict) -> Callable:
                 return accumulator
 
             # Apply the accumulator
-            add_accumulator(wrapped_op, make_accumulator)
+            _add_accumulator(wrapped_op, make_accumulator)
 
         return wrapped_op
 
@@ -1034,7 +1034,7 @@ def on_close() -> None:
     return iterator_wrapper(value, on_yield, on_error, on_close)
 
 
-def add_accumulator(
+def _add_accumulator(
     op: Op,
     make_accumulator: Callable[[dict], Callable[[S, V], S]],
     *,