agenta 0.51.2__py3-none-any.whl → 0.51.4__py3-none-any.whl

agenta/sdk/decorators/tracing.py

@@ -1,8 +1,18 @@
 from typing import Callable, Optional, Any, Dict, List, Union
 
+from opentelemetry import context as otel_context
+from opentelemetry.context import attach, detach
+from opentelemetry import trace
+
+
 from functools import wraps
 from itertools import chain
-from inspect import iscoroutinefunction, getfullargspec
+from inspect import (
+    getfullargspec,
+    iscoroutinefunction,
+    isgeneratorfunction,
+    isasyncgenfunction,
+)
 
 from opentelemetry import baggage
 from opentelemetry.context import attach, detach, get_current
@@ -31,6 +41,7 @@ class instrument:  # pylint: disable=invalid-name
         redact: Optional[Callable[..., Any]] = None,
         redact_on_error: Optional[bool] = True,
         max_depth: Optional[int] = 2,
+        aggregate: Optional[Callable[[List[Any]], Any]] = None,
         # DEPRECATING
         kind: str = "task",
         spankind: Optional[str] = "TASK",
@@ -43,67 +54,213 @@ class instrument:  # pylint: disable=invalid-name
         self.redact = redact
         self.redact_on_error = redact_on_error
         self.max_depth = max_depth
+        self.aggregate = aggregate
 
     def __call__(self, func: Callable[..., Any]):
         is_coroutine_function = iscoroutinefunction(func)
+        is_sync_generator = isgeneratorfunction(func)
+        is_async_generator = isasyncgenfunction(func)
 
-        @wraps(func)
-        async def awrapper(*args, **kwargs):
-            async def aauto_instrumented(*args, **kwargs):
+        # ---- ASYNC GENERATOR ----
+        if is_async_generator:
+
+            @wraps(func)
+            def astream_wrapper(*args, **kwargs):
                 self._parse_type_and_kind()
 
-                token = self._attach_baggage()
+                captured_ctx = otel_context.get_current()
+                parent_ctx = trace.set_span_in_context(trace.get_current_span())
 
+                token = self._attach_baggage()
                 ctx = self._get_traceparent()
 
-                with ag.tracer.start_as_current_span(
-                    name=func.__name__,
-                    kind=self.kind,
-                    context=ctx,
-                ):
-                    self._set_link()
+                agen = func(*args, **kwargs)
 
-                    self._pre_instrument(func, *args, **kwargs)
+                async def wrapped_generator():
+                    token_ctx = otel_context.attach(captured_ctx)
 
-                    result = await func(*args, **kwargs)
+                    try:
+                        with ag.tracer.start_as_current_span(
+                            name=func.__name__,
+                            kind=self.kind,
+                            context=parent_ctx or ctx,
+                        ):
+                            self._set_link()
+                            self._pre_instrument(func, *args, **kwargs)
 
-                    self._post_instrument(result)
+                            _result = []
 
-                    return result
+                            try:
+                                while True:
+                                    chunk = await agen.__anext__()
 
-                self._detach_baggage(token)
+                                    _result.append(chunk)
+                                    yield chunk
 
-            return await aauto_instrumented(*args, **kwargs)
+                            except StopAsyncIteration:
+                                pass
 
-        @wraps(func)
-        def wrapper(*args, **kwargs):
-            def auto_instrumented(*args, **kwargs):
+                            finally:
+                                if self.aggregate:
+                                    result = self.aggregate(_result)
+                                elif all(isinstance(r, str) for r in _result):
+                                    result = "".join(_result)
+                                elif all(isinstance(r, bytes) for r in _result):
+                                    result = b"".join(_result)
+                                else:
+                                    result = _result
+
+                                self._post_instrument(result)
+
+                    finally:
+                        otel_context.detach(token_ctx)
+
+                        self._detach_baggage(token)
+
+                return wrapped_generator()
+
+            return astream_wrapper
+
+        # ---- SYNC GENERATOR ----
+        if is_sync_generator:
+
+            @wraps(func)
+            def stream_wrapper(*args, **kwargs):
                 self._parse_type_and_kind()
 
+                captured_ctx = otel_context.get_current()
+                parent_ctx = trace.set_span_in_context(trace.get_current_span())
+
                 token = self._attach_baggage()
+                ctx = self._get_traceparent()
+
+                agen = func(*args, **kwargs)
+
+                def wrapped_generator():
+                    token_ctx = otel_context.attach(captured_ctx)
+
+                    try:
+                        with ag.tracer.start_as_current_span(
+                            name=func.__name__,
+                            kind=self.kind,
+                            context=parent_ctx or ctx,
+                        ):
+                            self._set_link()
+                            self._pre_instrument(func, *args, **kwargs)
+
+                            _result = []
+
+                            gen_return = None
+                            stop_iteration_raised = False
+
+                            try:
+                                while True:
+                                    try:
+                                        chunk = next(agen)
 
+                                    except StopIteration as e:
+                                        gen_return = e.value
+                                        stop_iteration_raised = True
+                                        break
+
+                                    _result.append(chunk)
+
+                                    yield chunk
+
+                            finally:
+                                if self.aggregate:
+                                    result = self.aggregate(_result)
+                                elif all(isinstance(r, str) for r in _result):
+                                    result = "".join(_result)
+                                elif all(isinstance(r, bytes) for r in _result):
+                                    result = b"".join(_result)
+                                else:
+                                    result = _result
+
+                                self._post_instrument(result)
+
+                            return gen_return if stop_iteration_raised else None
+
+                    finally:
+                        otel_context.detach(token_ctx)
+
+                        self._detach_baggage(token)
+
+                return wrapped_generator()
+
+            return stream_wrapper
+
+        # ---- ASYNC FUNCTION (non-generator) ----
+        if is_coroutine_function:
+
+            @wraps(func)
+            async def awrapper(*args, **kwargs):
+                self._parse_type_and_kind()
+
+                captured_ctx = otel_context.get_current()
+                parent_ctx = trace.set_span_in_context(trace.get_current_span())
+
+                token = self._attach_baggage()
                 ctx = self._get_traceparent()
 
+                token_ctx = otel_context.attach(captured_ctx)
+
+                try:
+                    with ag.tracer.start_as_current_span(
+                        name=func.__name__,
+                        kind=self.kind,
+                        context=parent_ctx or ctx,
+                    ):
+                        self._set_link()
+                        self._pre_instrument(func, *args, **kwargs)
+
+                        result = await func(*args, **kwargs)
+
+                        self._post_instrument(result)
+
+                finally:
+                    otel_context.detach(token_ctx)
+
+                    self._detach_baggage(token)
+
+                return result
+
+            return awrapper
+
+        # ---- SYNC FUNCTION (non-generator) ----
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            self._parse_type_and_kind()
+
+            captured_ctx = otel_context.get_current()
+            parent_ctx = trace.set_span_in_context(trace.get_current_span())
+
+            token = self._attach_baggage()
+            ctx = self._get_traceparent()
+
+            token_ctx = otel_context.attach(captured_ctx)
+
+            try:
                 with ag.tracer.start_as_current_span(
                     name=func.__name__,
                     kind=self.kind,
-                    context=ctx,
+                    context=parent_ctx or ctx,
                 ):
                     self._set_link()
-
                     self._pre_instrument(func, *args, **kwargs)
 
                     result = func(*args, **kwargs)
 
                     self._post_instrument(result)
 
-                    return result
+            finally:
+                otel_context.detach(token_ctx)
 
                 self._detach_baggage(token)
 
-            return auto_instrumented(*args, **kwargs)
+            return result
 
-        return awrapper if is_coroutine_function else wrapper
+        return wrapper
 
     def _parse_type_and_kind(self):
         if not ag.tracing.get_current_span().is_recording():

agenta/sdk/tracing/processors.py

@@ -8,10 +8,6 @@ from opentelemetry.sdk.trace.export import (
     SpanExporter,
     ReadableSpan,
     BatchSpanProcessor,
-    _DEFAULT_MAX_QUEUE_SIZE,
-    _DEFAULT_SCHEDULE_DELAY_MILLIS,
-    _DEFAULT_MAX_EXPORT_BATCH_SIZE,
-    _DEFAULT_EXPORT_TIMEOUT_MILLIS,
 )
 
 from agenta.sdk.utils.logging import get_module_logger
@@ -42,10 +38,10 @@ class TraceProcessor(SpanProcessor):
         if not self.inline:
             self._delegate = BatchSpanProcessor(
                 span_exporter,
-                max_queue_size or _DEFAULT_MAX_QUEUE_SIZE,
-                schedule_delay_millis or _DEFAULT_SCHEDULE_DELAY_MILLIS,
-                max_export_batch_size or _DEFAULT_MAX_EXPORT_BATCH_SIZE,
-                export_timeout_millis or _DEFAULT_EXPORT_TIMEOUT_MILLIS,
+                max_queue_size,
+                schedule_delay_millis,
+                max_export_batch_size,
+                export_timeout_millis,
             )
         # --- DISTRIBUTED
 

{agenta-0.51.2.dist-info → agenta-0.51.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: agenta
-Version: 0.51.2
+Version: 0.51.4
 Summary: The SDK for agenta is an open-source LLMOps platform.
 Keywords: LLMOps,LLM,evaluation,prompt engineering
 Author: Mahmoud Mabrouk
@@ -22,7 +22,7 @@ Requires-Dist: httpx (>=0.28.0)
 Requires-Dist: huggingface-hub (<0.31.0)
 Requires-Dist: importlib-metadata (>=8.0.0,<9.0)
 Requires-Dist: jinja2 (>=3.1.6,<4.0.0)
-Requires-Dist: litellm (>=1.75.9)
+Requires-Dist: litellm (==1.76.0)
 Requires-Dist: openai (>=1.100.0)
 Requires-Dist: opentelemetry-api (>=1.27.0,<2.0.0)
 Requires-Dist: opentelemetry-exporter-otlp-proto-http (>=1.27.0,<2.0.0)

{agenta-0.51.2.dist-info → agenta-0.51.4.dist-info}/RECORD

@@ -315,7 +315,7 @@ agenta/sdk/context/tracing.py,sha256=tM1T-HBt9D1qD8VsVKWlLe6MBRQrtHUA33E3KNM4GXQ
 agenta/sdk/decorators/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 agenta/sdk/decorators/running.py,sha256=j077-4u4XytNV57g0UY8ddG_M2k99Q829yPejiuGeg8,4034
 agenta/sdk/decorators/serving.py,sha256=jU8nP0btR4tytPkjiEmrmVSvShSSEFX7tKilHZfr-XE,26186
-agenta/sdk/decorators/tracing.py,sha256=Ug1ybanpW6CsvT41eh2wDrjNy81amDqGU0Yo2FN5sPk,10222
+agenta/sdk/decorators/tracing.py,sha256=4QkOaQzRIkR9_OTMZ31rC2Ad3G7MZXIx7SoF194QwGM,15832
 agenta/sdk/litellm/__init__.py,sha256=Bpz1gfHQc0MN1yolWcjifLWznv6GjHggvRGQSpxpihM,37
 agenta/sdk/litellm/litellm.py,sha256=dCjw0H_jD3L3UQ3l9SbLm5dfeIGnel6dPtQYJ78beYM,10202
 agenta/sdk/litellm/mockllm.py,sha256=F10Sbj2V1Cihkg6bgHkIOuw8irYEYpIp47UsO8LGOeE,2322
@@ -345,7 +345,7 @@ agenta/sdk/tracing/attributes.py,sha256=DwjjOk3mGOvz0jYu8EYr3hhItvawK5GX80_Mfciq
 agenta/sdk/tracing/conventions.py,sha256=JBtznBXZ3aRkGKkLl7cPwdMNh3w1G-H2Ta2YrAxbr38,950
 agenta/sdk/tracing/exporters.py,sha256=b1Aq5AcQYnmAxapQMYpN_MfaLOG4IGTzlQw3PfBiKgI,3486
 agenta/sdk/tracing/inline.py,sha256=y2S_MGGqmXgyUgbkNNyrb8_X-QtGuDy8JwxlwWibIx8,31507
-agenta/sdk/tracing/processors.py,sha256=wZWF8vTEouhbGxbfPGXVQJI53CcXFByqxHiCSIzJV-4,5407
+agenta/sdk/tracing/processors.py,sha256=ZE6baO2U3ER73rckvXhdbNmLgs3jn6le7iAUQc6osbY,5141
 agenta/sdk/tracing/propagation.py,sha256=EeOqDMqnh_MoEhGd1do_vy_tQBYUcoC8kpLqVoZeqg0,2561
 agenta/sdk/tracing/spans.py,sha256=nqUOjjirBxB8Eacv8Qj4Ra_6rknGi3lbJdNyKmk5ODQ,3707
 agenta/sdk/tracing/tracing.py,sha256=eqQLcvgmR6Tz0WSwLCERkGKiLzsKfDyR_ah-Tyz1qGE,9256
@@ -365,6 +365,6 @@ agenta/sdk/workflows/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hS
 agenta/sdk/workflows/registry.py,sha256=OucCfmhyEENAF5OwjRBLjbDIy3Qar-prFXLhKMU0BZA,944
 agenta/sdk/workflows/types.py,sha256=nklaaIrK0o5Nloy1l6hnPXjtN7H4-OUv0ivAUf278Og,11682
 agenta/sdk/workflows/utils.py,sha256=ILfY8DSBWLrdWIuKg6mq7rANwKiiY6sxEeFiBFhjLYM,413
-agenta-0.51.2.dist-info/METADATA,sha256=iSFL_mwtya1RW2QZauN4et_Xnd--hEmYTlj0n3zDuL8,31641
-agenta-0.51.2.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-agenta-0.51.2.dist-info/RECORD,,
+agenta-0.51.4.dist-info/METADATA,sha256=tsHuDehdC8pfhngzFsGDP4qX5nuKOOnXCcKx-4FfOgc,31641
+agenta-0.51.4.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+agenta-0.51.4.dist-info/RECORD,,