agenta 0.51.4__tar.gz → 0.51.6__tar.gz

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

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

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agenta"
-version = "0.51.4"
+version = "0.51.6"
 description = "The SDK for agenta is an open-source LLMOps platform."
 readme = "README.md"
 authors = ["Mahmoud Mabrouk <mahmoud@agenta.ai>"]