agenta 0.51.3__tar.gz → 0.51.4__tar.gz

{agenta-0.51.3 → agenta-0.51.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: agenta
-Version: 0.51.3
+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.3 → agenta-0.51.4}/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-0.51.3 → agenta-0.51.4}/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.3 → agenta-0.51.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agenta"
-version = "0.51.3"
+version = "0.51.4"
 description = "The SDK for agenta is an open-source LLMOps platform."
 readme = "README.md"
 authors = ["Mahmoud Mabrouk <mahmoud@agenta.ai>"]
@@ -29,7 +29,7 @@ importlib-metadata = ">=8.0.0,<9.0"
 httpx = ">=0.28.0"
 pyyaml = "^6.0.2"
 toml = "^0.10.2"
-litellm = ">=1.75.9"
+litellm = "==1.76.0"
 jinja2 = "^3.1.6"
 opentelemetry-api = ">=1.27.0,<2.0.0"
 opentelemetry-sdk = ">=1.27.0,<2.0.0"