agenta 0.51.5__py3-none-any.whl → 0.51.6__py3-none-any.whl

agenta/sdk/decorators/serving.py

@@ -1,5 +1,12 @@
 from typing import Type, Any, Callable, Dict, Optional, Tuple, List
-from inspect import signature, iscoroutinefunction, Signature, Parameter
+from inspect import (
+    iscoroutinefunction,
+    isgenerator,
+    isasyncgen,
+    signature,
+    Signature,
+    Parameter,
+)
 from functools import wraps
 from traceback import format_exception
 from asyncio import sleep
@@ -7,6 +14,10 @@ from uuid import UUID
 from pydantic import BaseModel, HttpUrl, ValidationError
 from os import environ
 
+from starlette.responses import (
+    Response as StarletteResponse,
+    StreamingResponse,
+)
 from fastapi import Body, FastAPI, HTTPException, Request
 
 from agenta.sdk.middleware.mock import MockMiddleware
@@ -33,6 +44,7 @@ from agenta.sdk.utils.helpers import get_current_version
 from agenta.sdk.types import (
     MultipleChoice,
     BaseResponse,
+    StreamResponse,
     MCField,
 )
 
@@ -380,6 +392,37 @@ class entrypoint:
                 span_id,
             ) = await self.fetch_inline_trace(inline)
 
+        try:
+            if isinstance(result, StarletteResponse):
+                result.headers.setdefault("X-ag-version", "3.0")
+                if content_type:
+                    result.headers.setdefault("X-ag-content-type", content_type)
+                if tree_id:
+                    result.headers.setdefault("X-ag-tree-id", tree_id)
+                if trace_id:
+                    result.headers.setdefault("X-ag-trace-id", trace_id)
+                if span_id:
+                    result.headers.setdefault("X-ag-span-id", span_id)
+
+                return result
+        except:
+            return result
+
+        try:
+            if isasyncgen(result) or isgenerator(result):
+                return StreamResponse(
+                    content=result,
+                    content_type=content_type,
+                    tree_id=tree_id,
+                    trace_id=trace_id,
+                    span_id=span_id,
+                )
+        except:
+            return StreamingResponse(
+                result,
+                media_type="text/event-stream",
+            )
+
         try:
             return BaseResponse(
                 data=data,

agenta/sdk/decorators/tracing.py

@@ -18,12 +18,12 @@ from opentelemetry import baggage
 from opentelemetry.context import attach, detach, get_current
 from opentelemetry.baggage import set_baggage, get_all
 
+from agenta.sdk.utils.logging import get_module_logger
 from agenta.sdk.utils.exceptions import suppress
+from agenta.sdk.utils.otel import debug_otel_context
 from agenta.sdk.context.tracing import tracing_context
 from agenta.sdk.tracing.conventions import parse_span_kind
 
-from agenta.sdk.utils.logging import get_module_logger
-
 import agenta as ag
 
 log = get_module_logger(__name__)
@@ -66,40 +66,43 @@ class instrument:  # pylint: disable=invalid-name
 
             @wraps(func)
             def astream_wrapper(*args, **kwargs):
-                self._parse_type_and_kind()
+                # debug_otel_context("[BEFORE STREAM] [BEFORE SETUP]")
 
                 captured_ctx = otel_context.get_current()
-                parent_ctx = trace.set_span_in_context(trace.get_current_span())
 
-                token = self._attach_baggage()
+                self._parse_type_and_kind()
+
+                self._attach_baggage()
+
                 ctx = self._get_traceparent()
 
-                agen = func(*args, **kwargs)
+                # debug_otel_context("[BEFORE STREAM] [AFTER SETUP]")
 
                 async def wrapped_generator():
-                    token_ctx = otel_context.attach(captured_ctx)
+                    # debug_otel_context("[WITHIN STREAM] [BEFORE ATTACH]")
+
+                    otel_token = otel_context.attach(captured_ctx)
+
+                    # debug_otel_context("[WITHIN STREAM] [AFTER ATTACH]")
 
                     try:
                         with ag.tracer.start_as_current_span(
                             name=func.__name__,
                             kind=self.kind,
-                            context=parent_ctx or ctx,
+                            context=ctx,
                         ):
                             self._set_link()
                             self._pre_instrument(func, *args, **kwargs)
 
                             _result = []
 
-                            try:
-                                while True:
-                                    chunk = await agen.__anext__()
+                            agen = func(*args, **kwargs)
 
+                            try:
+                                async for chunk in agen:
                                     _result.append(chunk)
                                     yield chunk
 
-                            except StopAsyncIteration:
-                                pass
-
                             finally:
                                 if self.aggregate:
                                     result = self.aggregate(_result)
@@ -113,9 +116,11 @@ class instrument:  # pylint: disable=invalid-name
                                 self._post_instrument(result)
 
                     finally:
-                        otel_context.detach(token_ctx)
+                        # debug_otel_context("[WITHIN STREAM] [BEFORE DETACH]")
 
-                        self._detach_baggage(token)
+                        otel_context.detach(otel_token)
+
+                        # debug_otel_context("[WITHIN STREAM] [AFTER DETACH]")
 
                 return wrapped_generator()
 
@@ -128,43 +133,36 @@ class instrument:  # pylint: disable=invalid-name
             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)
+                ctx = self._get_traceparent()
 
                 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,
+                            context=ctx,
                         ):
                             self._set_link()
+
                             self._pre_instrument(func, *args, **kwargs)
 
                             _result = []
 
+                            gen = func(*args, **kwargs)
+
                             gen_return = None
-                            stop_iteration_raised = False
 
                             try:
                                 while True:
                                     try:
-                                        chunk = next(agen)
-
+                                        chunk = next(gen)
                                     except StopIteration as e:
                                         gen_return = e.value
-                                        stop_iteration_raised = True
                                         break
 
                                     _result.append(chunk)
-
                                     yield chunk
 
                             finally:
@@ -179,11 +177,9 @@ class instrument:  # pylint: disable=invalid-name
 
                                 self._post_instrument(result)
 
-                            return gen_return if stop_iteration_raised else None
+                            return gen_return
 
                     finally:
-                        otel_context.detach(token_ctx)
-
                         self._detach_baggage(token)
 
                 return wrapped_generator()
@@ -197,21 +193,18 @@ class instrument:  # pylint: disable=invalid-name
             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)
+                ctx = self._get_traceparent()
 
                 try:
                     with ag.tracer.start_as_current_span(
                         name=func.__name__,
                         kind=self.kind,
-                        context=parent_ctx or ctx,
+                        context=ctx,
                     ):
                         self._set_link()
+
                         self._pre_instrument(func, *args, **kwargs)
 
                         result = await func(*args, **kwargs)
@@ -219,8 +212,6 @@ class instrument:  # pylint: disable=invalid-name
                         self._post_instrument(result)
 
                 finally:
-                    otel_context.detach(token_ctx)
-
                     self._detach_baggage(token)
 
                 return result
@@ -232,21 +223,18 @@ class instrument:  # pylint: disable=invalid-name
         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)
+            ctx = self._get_traceparent()
 
             try:
                 with ag.tracer.start_as_current_span(
                     name=func.__name__,
                     kind=self.kind,
-                    context=parent_ctx or ctx,
+                    context=ctx,
                 ):
                     self._set_link()
+
                     self._pre_instrument(func, *args, **kwargs)
 
                     result = func(*args, **kwargs)
@@ -254,8 +242,6 @@ class instrument:  # pylint: disable=invalid-name
                     self._post_instrument(result)
 
             finally:
-                otel_context.detach(token_ctx)
-
                 self._detach_baggage(token)
 
             return result

agenta/sdk/types.py

@@ -5,6 +5,9 @@ from typing import List, Union, Optional, Dict, Literal, Any
 from pydantic import ConfigDict, BaseModel, HttpUrl
 from pydantic import BaseModel, Field, model_validator
 
+from starlette.responses import StreamingResponse
+
+
 from agenta.sdk.assets import supported_llm_models
 from agenta.client.backend.types import AgentaNodesResponse, AgentaNodeDto
 
@@ -45,6 +48,42 @@ class BaseResponse(BaseModel):
     model_config = ConfigDict(use_enum_values=True, exclude_none=True)
 
 
+class StreamResponse(StreamingResponse):
+    def __init__(
+        self,
+        content,
+        media_type: str = "text/event-stream",
+        *,
+        version: Optional[str] = "3.0",
+        tree_id: Optional[str] = None,
+        trace_id: Optional[str] = None,
+        span_id: Optional[str] = None,
+        content_type: Optional[str] = None,
+        extra_headers: Optional[Dict[str, str]] = None,
+        status_code: int = 200,
+        background=None,
+    ):
+        headers = dict(extra_headers or {})
+        if version is not None:
+            headers["X-ag-version"] = version
+        if content_type:
+            headers["X-ag-content-type"] = content_type
+        if tree_id:
+            headers["X-ag-tree-id"] = tree_id
+        if trace_id:
+            headers["X-ag-trace-id"] = trace_id
+        if span_id:
+            headers["X-ag-span-id"] = span_id
+
+        super().__init__(
+            content=content,
+            media_type=media_type,
+            status_code=status_code,
+            headers=headers,
+            background=background,
+        )
+
+
 class DictInput(dict):
     def __new__(cls, default_keys: Optional[List[str]] = None):
         instance = super().__new__(cls, default_keys)

agenta/sdk/utils/otel.py

@@ -0,0 +1,21 @@
+from opentelemetry import context as otel_context
+from opentelemetry.trace import get_current_span
+from opentelemetry.baggage import get_all as get_all_baggage
+
+
+def debug_otel_context(label: str = "OTEL CONTEXT"):
+    ctx = otel_context.get_current()
+    span = get_current_span()
+    baggage = get_all_baggage(ctx)
+
+    print("\n===== {} =====".format(label))
+    print("Context object:", ctx)
+    print("Baggage:", baggage or "{}")
+    if span and span.get_span_context().is_valid:
+        print("Current Span:")
+        print("  Trace ID:", format(span.get_span_context().trace_id, "032x"))
+        print("  Span  ID:", format(span.get_span_context().span_id, "016x"))
+        # print("  Is Recording:", span.is_recording())
+    else:
+        print("Current Span: None or Invalid")
+    print("=" * 40)

{agenta-0.51.5.dist-info → agenta-0.51.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: agenta
-Version: 0.51.5
+Version: 0.51.6
 Summary: The SDK for agenta is an open-source LLMOps platform.
 Keywords: LLMOps,LLM,evaluation,prompt engineering
 Author: Mahmoud Mabrouk

{agenta-0.51.5.dist-info → agenta-0.51.6.dist-info}/RECORD

@@ -314,8 +314,8 @@ agenta/sdk/context/serving.py,sha256=f0Iq8Sv_CBJF7838tDecy_89K6HI7k6joWBkWHTM7yE
 agenta/sdk/context/tracing.py,sha256=tM1T-HBt9D1qD8VsVKWlLe6MBRQrtHUA33E3KNM4GXQ,1237
 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=4QkOaQzRIkR9_OTMZ31rC2Ad3G7MZXIx7SoF194QwGM,15832
+agenta/sdk/decorators/serving.py,sha256=N-oIKo39AD9bOw4dQibqc4YLW9rqa778e-2OlxLOkUw,27439
+agenta/sdk/decorators/tracing.py,sha256=Nz1Y2na1Uw4S56qWZI9EXAI6LAxlru9SG5PdXJM6W_k,15105
 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
@@ -349,7 +349,7 @@ agenta/sdk/tracing/processors.py,sha256=ZE6baO2U3ER73rckvXhdbNmLgs3jn6le7iAUQc6o
 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
-agenta/sdk/types.py,sha256=Cq0U3gI8ll65KhsnVWkuLJvwZGEsuTe3WdsHX1UX6zQ,20753
+agenta/sdk/types.py,sha256=0-_eR7BZjLZRHmo6HSTsfOhNFmbX_a6GfigqaM7dFX4,21883
 agenta/sdk/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 agenta/sdk/utils/cache.py,sha256=q6SYVzHW9RCwtZ7z4Ue3uMwXroroLlko5AvQ1MouvFM,1429
 agenta/sdk/utils/constants.py,sha256=zW3R4rjXOo2L5lz6q84l_zYuOM9u4mpPRHw_B1Dr_hI,67
@@ -358,6 +358,7 @@ agenta/sdk/utils/exceptions.py,sha256=Dq3lh5Ug7IhF4nvYRXXlwhgxEuByOnT1w0uVMZx1ea
 agenta/sdk/utils/globals.py,sha256=9ixKS8aI6ZWgsjT8WngvTg4dsnP2cKErNJngHxzdK9U,256
 agenta/sdk/utils/helpers.py,sha256=nqUrTfRQfzQhP9ZuJ72MXVRtIFgAoIhzymFN2wGZYhQ,1198
 agenta/sdk/utils/logging.py,sha256=j4NzpFk2ilOM10sRBOxCjmansDHSx6HwMV8IAEreRb8,8386
+agenta/sdk/utils/otel.py,sha256=X5oNv3PqTghDPa0OjKeFJdFNcfDtI40N4jv5WMNrZ38,806
 agenta/sdk/utils/preinit.py,sha256=1TAAHhYyYnLLwvzwnf33Qwkou7tI3iITlVt-1kcyGaM,1186
 agenta/sdk/utils/singleton.py,sha256=17Ph7LGnnV8HkPjImruKita2ni03Ari5jr0jqm__4sc,312
 agenta/sdk/utils/timing.py,sha256=nZR-kudVUtKFlHuBhztgSGxj7FVnCB4Uv6sfg-1dkrQ,1556
@@ -365,6 +366,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.5.dist-info/METADATA,sha256=mWw6SmWBaGsUFleUTe-rpldEzjwto991eMyScJW2FeY,31641
-agenta-0.51.5.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-agenta-0.51.5.dist-info/RECORD,,
+agenta-0.51.6.dist-info/METADATA,sha256=UpiHrNBK2TAXDbxHOKqmqmWAKE3ktZXraOnzge3tX_4,31641
+agenta-0.51.6.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+agenta-0.51.6.dist-info/RECORD,,