agenta 0.20.0a0__py3-none-any.whl → 0.20.0a1__py3-none-any.whl

agenta/sdk/decorators/llm_entrypoint.py

@@ -281,6 +281,11 @@ class entrypoint(BaseDecorator):
             For synchronous functions, it calls them directly, while for asynchronous functions,
             it awaits their execution.
             """
+            WAIT_FOR_SPANS = True
+            TIMEOUT = 10
+            TIMESTEP = 0.01
+            NOFSTEPS = TIMEOUT / TIMESTEP
+
             data = None
             trace = None
 
@@ -296,7 +301,15 @@ class entrypoint(BaseDecorator):
                 result = func(*args, **func_params["params"])
 
             if token is not None:
+                if WAIT_FOR_SPANS:
+                    remaining_steps = NOFSTEPS
+
+                    while not ag.tracing.is_trace_ready() and remaining_steps > 0:
+                        await asyncio.sleep(0.01)
+                        remaining_steps -= 1
+
                 trace = ag.tracing.dump_trace()
+                ag.tracing.flush_spans()
                 tracing_context.reset(token)
 
             if isinstance(result, Context):

agenta/sdk/tracing/callbacks.py

@@ -1,5 +1,9 @@
 import agenta as ag
 
+from agenta.sdk.tracing.tracing_context import tracing_context, TracingContext
+
+from agenta.sdk.utils.debug import debug
+
 
 def litellm_handler():
     try:
@@ -23,20 +27,25 @@ def litellm_handler():
             LitellmCustomLogger (object): custom logger that allows us to override the events to capture.
         """
 
+        def __init__(self):
+            self.span = None
+
         @property
         def _trace(self):
             return ag.tracing
 
+        @debug()
         def log_pre_api_call(self, model, messages, kwargs):
             call_type = kwargs.get("call_type")
             span_kind = (
                 "llm" if call_type in ["completion", "acompletion"] else "embedding"
             )
 
-            ag.tracing.start_span(
+            self.span = ag.tracing.open_span(
                 name=f"{span_kind}_call",
                 input={"messages": kwargs["messages"]},
                 spankind=span_kind,
+                active=False,
             )
             ag.tracing.set_attributes(
                 {
@@ -49,9 +58,10 @@ def litellm_handler():
                 }
             )
 
+        @debug()
         def log_stream_event(self, kwargs, response_obj, start_time, end_time):
-            ag.tracing.set_status(status="OK")
-            ag.tracing.end_span(
+            ag.tracing.set_status(status="OK", span_id=self.span.id)
+            ag.tracing.store_outputs(
                 outputs={
                     "message": kwargs.get(
                         "complete_streaming_response"
@@ -65,13 +75,16 @@ def litellm_handler():
                         "response_cost"
                     ),  # litellm calculates response cost
                 },
+                span_id=self.span.id,
             )
+            ag.tracing.close_span(span_id=self.span.id)
 
+        @debug()
         def log_success_event(
             self, kwargs, response_obj: ModelResponse, start_time, end_time
         ):
-            ag.tracing.set_status(status="OK")
-            ag.tracing.end_span(
+            ag.tracing.set_status(status="OK", span_id=self.span.id)
+            ag.tracing.store_outputs(
                 outputs={
                     "message": response_obj.choices[0].message.content,
                     "usage": (
@@ -83,12 +96,15 @@ def litellm_handler():
                         "response_cost"
                     ),  # litellm calculates response cost
                 },
+                span_id=self.span.id,
             )
+            ag.tracing.close_span(span_id=self.span.id)
 
+        @debug()
         def log_failure_event(
             self, kwargs, response_obj: ModelResponse, start_time, end_time
         ):
-            ag.tracing.set_status(status="ERROR")
+            ag.tracing.set_status(status="ERROR", span_id=self.span.id)
             ag.tracing.set_attributes(
                 {
                     "traceback_exception": repr(
@@ -98,10 +114,11 @@ def litellm_handler():
                         "end_time"
                     ],  # datetime object of when call was completed
                 },
+                span_id=self.span.id,
             )
-            ag.tracing.end_span(
+            ag.tracing.store_outputs(
                 outputs={
-                    "message": kwargs["exception"],  # the Exception raised
+                    "message": repr(kwargs["exception"]),  # the Exception raised
                     "usage": (
                         response_obj.usage.dict()
                         if hasattr(response_obj, "usage")
@@ -111,13 +128,16 @@ def litellm_handler():
                         "response_cost"
                     ),  # litellm calculates response cost
                 },
+                span_id=self.span.id,
             )
+            ag.tracing.close_span(span_id=self.span.id)
 
+        @debug()
         async def async_log_stream_event(
             self, kwargs, response_obj, start_time, end_time
         ):
-            ag.tracing.set_status(status="OK")
-            ag.tracing.end_span(
+            ag.tracing.set_status(status="OK", span_id=self.span.id)
+            ag.tracing.store_outputs(
                 outputs={
                     "message": kwargs.get(
                         "complete_streaming_response"
@@ -131,13 +151,16 @@ def litellm_handler():
                         "response_cost"
                     ),  # litellm calculates response cost
                 },
+                span_id=self.span.id,
             )
+            ag.tracing.close_span(span_id=self.span.id)
 
+        @debug()
         async def async_log_success_event(
             self, kwargs, response_obj, start_time, end_time
         ):
-            ag.tracing.set_status(status="OK")
-            ag.tracing.end_span(
+            ag.tracing.set_status(status="OK", span_id=self.span.id)
+            ag.tracing.store_outputs(
                 outputs={
                     "message": response_obj.choices[0].message.content,
                     "usage": (
@@ -149,12 +172,15 @@ def litellm_handler():
                         "response_cost"
                     ),  # litellm calculates response cost
                 },
+                span_id=self.span.id,
             )
+            ag.tracing.close_span(span_id=self.span.id)
 
+        @debug()
         async def async_log_failure_event(
             self, kwargs, response_obj, start_time, end_time
         ):
-            ag.tracing.set_status(status="ERROR")
+            ag.tracing.set_status(status="ERROR", span_id=self.span.id)
             ag.tracing.set_attributes(
                 {
                     "traceback_exception": kwargs[
@@ -164,8 +190,9 @@ def litellm_handler():
                         "end_time"
                     ],  # datetime object of when call was completed
                 },
+                span_id=self.span.id,
             )
-            ag.tracing.end_span(
+            ag.tracing.store_outputs(
                 outputs={
                     "message": repr(kwargs["exception"]),  # the Exception raised
                     "usage": (
@@ -177,6 +204,8 @@ def litellm_handler():
                         "response_cost"
                     ),  # litellm calculates response cost
                 },
+                span_id=self.span.id,
             )
+            ag.tracing.close_span(span_id=self.span.id)
 
     return LitellmHandler()

agenta/sdk/tracing/llm_tracing.py

@@ -25,7 +25,7 @@ from bson.objectid import ObjectId
 
 VARIANT_TRACKING_FEATURE_FLAG = False
 
-from agenta.sdk.utils.debug import debug, DEBUG, SHIFT
+from agenta.sdk.utils.debug import debug
 
 
 logging.setLevel("DEBUG")
@@ -195,6 +195,14 @@ class Tracing(metaclass=SingletonMeta):
 
         tracing.trace_tags.extend(tags)
 
+    @debug()
+    def is_trace_ready(self):
+        tracing = tracing_context.get()
+
+        are_spans_ready = [span.end_time is not None for span in tracing.spans.values()]
+
+        return all(are_spans_ready)
+
     @debug()
     def close_trace(self) -> None:
         """
@@ -234,6 +242,7 @@ class Tracing(metaclass=SingletonMeta):
         name: str,
         spankind: str,
         input: Dict[str, Any],
+        active: bool = True,
         config: Optional[Dict[str, Any]] = None,
         **kwargs,
     ) -> CreateSpan:
@@ -270,7 +279,14 @@ class Tracing(metaclass=SingletonMeta):
         else:
             span.parent_span_id = tracing.active_span.id  # type: ignore
 
-        tracing.push(span)
+        tracing.spans[span.id] = span
+
+        if active:
+            tracing.active_span = span
+        else:
+            # DETACHED SPAN
+            pass
+
         ### --- TO BE CLEANED --- <<<
 
         logging.info(f"Opened  span  {span_id} {spankind.upper()}")
@@ -279,8 +295,7 @@ class Tracing(metaclass=SingletonMeta):
 
     @debug(req=True)
     def set_attributes(
-        self,
-        attributes: Dict[str, Any] = {},
+        self, attributes: Dict[str, Any] = {}, span_id: Optional[str] = None
     ) -> None:
         """
         Set attributes for the active span.
@@ -289,41 +304,31 @@ class Tracing(metaclass=SingletonMeta):
             attributes (Dict[str, Any], optional): A dictionary of attributes to set. Defaults to {}.
         """
 
-        tracing = tracing_context.get()
-
-        if tracing.active_span is None:
-            logging.error(f"Cannot set attributes ({set(attributes)}), no active span")
-            return
+        span = self._get_target_span(span_id)
 
         logging.info(
-            f"Setting span  {tracing.active_span.id} {tracing.active_span.spankind.upper()} attributes={attributes}"
+            f"Setting span  {span.id} {span.spankind.upper()} attributes={attributes}"
         )
 
         for key, value in attributes.items():
-            tracing.active_span.attributes[key] = value  # type: ignore
+            span.attributes[key] = value  # type: ignore
 
     @debug()
-    def set_status(self, status: str) -> None:
+    def set_status(self, status: str, span_id: Optional[str] = None) -> None:
         """
         Set status for the active span.
 
         Args:
             status: Enum ( UNSET, OK, ERROR )
         """
-        tracing = tracing_context.get()
-
-        if tracing.active_span is None:
-            logging.error(f"Cannot set status ({status}), no active span")
-            return
+        span = self._get_target_span(span_id)
 
-        logging.info(
-            f"Setting span  {tracing.active_span.id} {tracing.active_span.spankind.upper()} status={status}"
-        )
+        logging.info(f"Setting span  {span.id} {span.spankind.upper()} status={status}")
 
-        tracing.active_span.status = status
+        span.status = status
 
     @debug()
-    def close_span(self) -> None:
+    def close_span(self, span_id: Optional[str] = None) -> None:
         """
         Ends the active span, if it is a parent span, ends the trace too.
 
@@ -340,86 +345,76 @@ class Tracing(metaclass=SingletonMeta):
         Returns:
             None
         """
+        span = self._get_target_span(span_id)
 
-        tracing = tracing_context.get()
-
-        if tracing.active_span is None:
-            logging.error("Cannot close span, no active span")
-
-        span_id = tracing.active_span.id
-        spankind = tracing.active_span.spankind
+        spankind = span.spankind
 
-        logging.info(f"Closing span  {span_id} {spankind}")
+        logging.info(f"Closing span  {span.id} {spankind}")
 
         ### --- TO BE CLEANED --- >>>
-        tracing.active_span.end_time = datetime.now(timezone.utc)
+        span.end_time = datetime.now(timezone.utc)
 
         # TODO: Remove this whole part. Setting the cost should be done through set_span_attribute
-        if isinstance(tracing.active_span.outputs, dict):
-            self._update_span_cost(
-                tracing.active_span, tracing.active_span.outputs.get("cost", None)
-            )
-            self._update_span_tokens(
-                tracing.active_span, tracing.active_span.outputs.get("usage", None)
-            )
+        if isinstance(span.outputs, dict):
+            self._update_span_cost(span, span.outputs.get("cost", None))
+            self._update_span_tokens(span, span.outputs.get("usage", None))
 
-        active_span_parent_id = tracing.active_span.parent_span_id
+        span_parent_id = span.parent_span_id
 
-        if active_span_parent_id is not None:
-            parent_span = tracing.spans[active_span_parent_id]
-            self._update_span_cost(parent_span, tracing.active_span.cost)
-            self._update_span_tokens(parent_span, tracing.active_span.tokens)
-            tracing.active_span = parent_span
+        if span_parent_id is not None:
+            tracing = tracing_context.get()
+
+            parent_span = tracing.spans[span_parent_id]
+            self._update_span_cost(parent_span, span.cost)
+            self._update_span_tokens(parent_span, span.tokens)
+
+            if span_id is None:
+                tracing.active_span = parent_span
         ### --- TO BE CLEANED --- <<<
 
         logging.info(f"Closed  span  {span_id} {spankind}")
 
     @debug()
-    def store_internals(self, internals: Dict[str, Any] = {}) -> None:
+    def store_internals(
+        self, internals: Dict[str, Any] = {}, span_id: Optional[str] = None
+    ) -> None:
         """
         Set internals for the active span.
 
         Args:
             internals (Dict[str, Any], optional): A dictionary of local variables to set. Defaults to {}.
         """
-
-        tracing = tracing_context.get()
-
-        if tracing.active_span is None:
-            logging.error(f"Cannot set internals ({set(internals)}), no active span")
-            return
+        span = self._get_target_span(span_id)
 
         logging.info(
-            f"Setting span  {tracing.active_span.id} {tracing.active_span.spankind.upper()} internals={internals}"
+            f"Setting span  {span.id} {span.spankind.upper()} internals={internals}"
         )
 
-        if tracing.active_span.internals is None:
-            tracing.active_span.internals = dict()
+        if span.internals is None:
+            span.internals = dict()
 
         for key, value in internals.items():
-            tracing.active_span.internals[key] = value  # type: ignore
+            span.internals[key] = value  # type: ignore
 
     @debug()
-    def store_outputs(self, outputs: Dict[str, Any] = {}) -> None:
+    def store_outputs(
+        self, outputs: Dict[str, Any] = {}, span_id: Optional[str] = None
+    ) -> None:
         """
         Set outputs for the active span.
 
         Args:
             outputs (Dict[str, Any], optional): A dictionary of output variables to set. Defaults to {}.
         """
-
-        tracing = tracing_context.get()
-
-        if tracing.active_span is None:
-            logging.error(f"Cannot set outputs ({set(outputs)}), no active span")
-            return
+        span = self._get_target_span(span_id)
 
         logging.info(
-            f"Setting span  {tracing.active_span.id} {tracing.active_span.spankind.upper()} outputs={outputs}"
+            f"Setting span  {span.id} {span.spankind.upper()} outputs={outputs}"
         )
 
-        tracing.active_span.outputs = outputs
+        span.outputs = outputs
 
+    @debug()
     def dump_trace(self):
         """
         Collects and organizes tracing information into a dictionary.
@@ -443,7 +438,10 @@ class Tracing(metaclass=SingletonMeta):
                     trace["usage"] = (
                         None if span.tokens is None else json.loads(span.tokens.json())
                     )
-                    trace["latency"] = (span.end_time - span.start_time).total_seconds()
+                    trace["latency"] = (
+                        (span.end_time if span.end_time else span.start_time)
+                        - span.start_time
+                    ).total_seconds()
 
             spans = (
                 []
@@ -526,6 +524,26 @@ class Tracing(metaclass=SingletonMeta):
         # return uuid4().hex[:16]
         return str(ObjectId())
 
+    def _get_target_span(self, span_id) -> CreateSpan:
+        tracing = tracing_context.get()
+
+        span = None
+
+        if span_id is None:
+            if tracing.active_span is None:
+                logging.error(f"Cannot set attributes, no active span")
+                return
+
+            span = tracing.active_span
+        else:
+            if span_id not in tracing.spans.keys():
+                logging.error(f"Cannot set attributes, span ({span_id}) not found")
+                return
+
+            span = tracing.spans[span_id]
+
+        return span
+
     def _process_spans(self) -> None:
         tracing = tracing_context.get()
 

agenta/sdk/tracing/tracing_context.py

@@ -23,9 +23,5 @@ class TracingContext:
     def __str__(self) -> str:
         return self.__repr__()
 
-    def push(self, span) -> None:
-        self.active_span = span
-        self.spans[span.id] = span
-
 
 tracing_context = ContextVar(CURRENT_TRACING_CONTEXT_KEY, default=None)

{agenta-0.20.0a0.dist-info → agenta-0.20.0a1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: agenta
-Version: 0.20.0a0
+Version: 0.20.0a1
 Summary: The SDK for agenta is an open-source LLMOps platform.
 Home-page: https://agenta.ai
 Keywords: LLMOps,LLM,evaluation,prompt engineering

{agenta-0.20.0a0.dist-info → agenta-0.20.0a1.dist-info}/RECORD

@@ -131,16 +131,16 @@ agenta/sdk/agenta_init.py,sha256=8MfDuypxohd0qRTdtGjX7L17KW-1UGmzNVdiqF15_ak,979
 agenta/sdk/client.py,sha256=trKyBOYFZRk0v5Eptxvh87yPf50Y9CqY6Qgv4Fy-VH4,2142
 agenta/sdk/context.py,sha256=q-PxL05-I84puunUAs9LGsffEXcYhDxhQxjuOz2vK90,901
 agenta/sdk/decorators/base.py,sha256=9aNdX5h8a2mFweuhdO-BQPwXGKY9ONPIdLRhSGAGMfY,217
-agenta/sdk/decorators/llm_entrypoint.py,sha256=cyXUHJ-cldbXmBOo3CF5XTSz_Pr2vAA5Inv5LHdigyE,26910
+agenta/sdk/decorators/llm_entrypoint.py,sha256=ukdw3jyKarkXHKNaQyQmzVnJsYCntzmgoAFccvCZQnA,27343
 agenta/sdk/decorators/tracing.py,sha256=BzXa_2b2VvbO5ZNTh126OkL3LFq8dMK-QFug8BHUBKA,2687
 agenta/sdk/router.py,sha256=0sbajvn5C7t18anH6yNo7-oYxldHnYfwcbmQnIXBePw,269
 agenta/sdk/tracing/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-agenta/sdk/tracing/callbacks.py,sha256=0rqkW-PGZSaNSMp_t4bGI7R9HQRgTyIy0gxmpGVJWpE,6915
+agenta/sdk/tracing/callbacks.py,sha256=xWP3RRtmj3ogXeqYI6zdj6gF5QGu70UlQ17WNwdXIx8,8063
 agenta/sdk/tracing/context_manager.py,sha256=HskDaiORoOhjeN375gm05wYnieQzh5UnoIsnSAHkAyc,252
-agenta/sdk/tracing/llm_tracing.py,sha256=VZskBPxa4S-0810KhsTnSy7XdrG3ovcpsKjfxQEGUx8,17167
+agenta/sdk/tracing/llm_tracing.py,sha256=A7_-Xx3Rioh_nfT7OjZvDPbKUvxEqruKhYRfdGwyo4w,17312
 agenta/sdk/tracing/logger.py,sha256=GfH7V-jBHcn7h5dbdrnkDMe_ml3wkXFBeoQiqR4KVRc,474
 agenta/sdk/tracing/tasks_manager.py,sha256=FBSFOWIKBycyA4ShB2ZVMzrzYQ8pWGWWBClFX8nlZFA,3726
-agenta/sdk/tracing/tracing_context.py,sha256=dTbsBMbIAmSOaWIxSfhKHOB1JjrAZ0eSHAlfzZs08Z4,891
+agenta/sdk/tracing/tracing_context.py,sha256=nt3ewa-TK9BRJviGIZYazsAQUiG4daWxjtsbjeaDprs,789
 agenta/sdk/types.py,sha256=1rVy8ob-rTOrIFcSSseXrt0JQtqqhlJfVgVxCB2ErCk,5754
 agenta/sdk/utils/debug.py,sha256=QyuPsSoN0425UD13x_msPxSF_VT6YwHiQunZUibI-jg,2149
 agenta/sdk/utils/globals.py,sha256=JmhJcCOSbwvjQ6GDyUc2_SYR27DZk7YcrRH80ktHHOM,435
@@ -161,7 +161,7 @@ agenta/templates/simple_prompt/app.py,sha256=kODgF6lhzsaJPdgL5b21bUki6jkvqjWZzWR
 agenta/templates/simple_prompt/env.example,sha256=g9AE5bYcGPpxawXMJ96gh8oenEPCHTabsiOnfQo3c5k,70
 agenta/templates/simple_prompt/requirements.txt,sha256=ywRglRy7pPkw8bljmMEJJ4aOOQKrt9FGKULZ-DGkoBU,23
 agenta/templates/simple_prompt/template.toml,sha256=DQBtRrF4GU8LBEXOZ-GGuINXMQDKGTEG5y37tnvIUIE,60
-agenta-0.20.0a0.dist-info/METADATA,sha256=xMd6YzEzMtRW5i6uO62QRRePYN0r91gLpIgaF8QrzIE,26462
-agenta-0.20.0a0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-agenta-0.20.0a0.dist-info/entry_points.txt,sha256=PDiu8_8AsL7ibU9v4iNoOKR1S7F2rdxjlEprjM9QOgo,46
-agenta-0.20.0a0.dist-info/RECORD,,
+agenta-0.20.0a1.dist-info/METADATA,sha256=s33kh03kzj-9YuQZipN4uhW57Tyhoj9O5g-g70w_aCY,26462
+agenta-0.20.0a1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+agenta-0.20.0a1.dist-info/entry_points.txt,sha256=PDiu8_8AsL7ibU9v4iNoOKR1S7F2rdxjlEprjM9QOgo,46
+agenta-0.20.0a1.dist-info/RECORD,,