agenta 0.17.1__py3-none-any.whl → 0.17.3__py3-none-any.whl

agenta/__init__.py

@@ -13,12 +13,15 @@ from .sdk.types import (
     FileInputURL,
     BinaryParam,
 )
+
 from .sdk.tracing.llm_tracing import Tracing
 from .sdk.decorators.tracing import instrument
 from .sdk.decorators.llm_entrypoint import entrypoint, app
 from .sdk.agenta_init import Config, AgentaSingleton, init
 from .sdk.utils.helper.openai_cost import calculate_token_usage
 from .sdk.client import Agenta
+from .sdk.tracing import callbacks
+
 
 config = PreInitObject("agenta.config", Config)
 DEFAULT_AGENTA_SINGLETON_INSTANCE = AgentaSingleton()

agenta/client/backend/types/llm_tokens.py

@@ -33,6 +33,6 @@ class LlmTokens(pydantic.BaseModel):
         return super().dict(**kwargs_with_defaults)
 
     class Config:
-        frozen = True
+        frozen = False
         smart_union = True
         json_encoders = {dt.datetime: serialize_datetime}

agenta/sdk/__init__.py

@@ -13,13 +13,13 @@ from .types import (
     FileInputURL,
     BinaryParam,
 )
+
 from .tracing.llm_tracing import Tracing
 from .decorators.tracing import instrument
 from .decorators.llm_entrypoint import entrypoint, app
 from .agenta_init import Config, AgentaSingleton, init
 from .utils.helper.openai_cost import calculate_token_usage
 
-
 config = PreInitObject("agenta.config", Config)
 DEFAULT_AGENTA_SINGLETON_INSTANCE = AgentaSingleton()
 tracing = DEFAULT_AGENTA_SINGLETON_INSTANCE.tracing  # type: ignore

agenta/sdk/decorators/tracing.py

@@ -1,5 +1,6 @@
 # Stdlib Imports
 import inspect
+import traceback
 from functools import wraps
 from typing import Any, Callable, Optional
 
@@ -59,6 +60,9 @@ class instrument(BaseDecorator):
                 self.tracing.update_span_status(span=span, value="OK")
             except Exception as e:
                 result = str(e)
+                self.tracing.set_span_attribute(
+                    {"traceback_exception": traceback.format_exc()}
+                )
                 self.tracing.update_span_status(span=span, value="ERROR")
             finally:
                 self.tracing.end_span(
@@ -87,6 +91,9 @@ class instrument(BaseDecorator):
                 self.tracing.update_span_status(span=span, value="OK")
             except Exception as e:
                 result = str(e)
+                self.tracing.set_span_attribute(
+                    {"traceback_exception": traceback.format_exc()}
+                )
                 self.tracing.update_span_status(span=span, value="ERROR")
             finally:
                 self.tracing.end_span(

agenta/sdk/tracing/callbacks.py

@@ -0,0 +1,157 @@
+import agenta as ag
+
+
+def litellm_handler():
+    try:
+        from litellm.utils import ModelResponse
+        from litellm.integrations.custom_logger import (
+            CustomLogger as LitellmCustomLogger,
+        )
+    except ImportError as exc:
+        raise ImportError(
+            "The litellm SDK is not installed. Please install it using `pip install litellm`."
+        ) from exc
+    except Exception as exc:
+        raise Exception(
+            "Unexpected error occurred when importing litellm: {}".format(exc)
+        ) from exc
+
+    class LitellmHandler(LitellmCustomLogger):
+        """This handler is responsible for instrumenting certain events when using litellm to call LLMs.
+
+        Args:
+            LitellmCustomLogger (object): custom logger that allows us to override the events to capture.
+        """
+
+        @property
+        def _trace(self):
+            return ag.tracing
+
+        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"
+            )
+            self._trace.start_span(
+                name=f"{span_kind}_call",
+                input={"messages": kwargs["messages"]},
+                spankind=span_kind,
+            )
+            self._trace.set_span_attribute(
+                {
+                    "model_config": {
+                        "model": kwargs.get("model"),
+                        **kwargs.get(
+                            "optional_params"
+                        ),  # model-specific params passed in
+                    },
+                }
+            )
+
+        def log_stream_event(self, kwargs, response_obj, start_time, end_time):
+            self._trace.update_span_status(span=self._trace.active_span, value="OK")
+            self._trace.end_span(
+                outputs={
+                    "message": kwargs.get(
+                        "complete_streaming_response"
+                    ),  # the complete streamed response (only set if `completion(..stream=True)`)
+                    "usage": response_obj.usage.dict(),  # litellm calculates usage
+                    "cost": kwargs.get(
+                        "response_cost"
+                    ),  # litellm calculates response cost
+                },
+            )
+
+        def log_success_event(
+            self, kwargs, response_obj: ModelResponse, start_time, end_time
+        ):
+            self._trace.update_span_status(span=self._trace.active_span, value="OK")
+            self._trace.end_span(
+                outputs={
+                    "message": response_obj.choices[0].message.content,
+                    "usage": response_obj.usage.dict(),  # litellm calculates usage
+                    "cost": kwargs.get(
+                        "response_cost"
+                    ),  # litellm calculates response cost
+                },
+            )
+
+        def log_failure_event(
+            self, kwargs, response_obj: ModelResponse, start_time, end_time
+        ):
+            self._trace.update_span_status(span=self._trace.active_span, value="ERROR")
+            self._trace.set_span_attribute(
+                {
+                    "traceback_exception": kwargs[
+                        "traceback_exception"
+                    ],  # the traceback generated via `traceback.format_exc()`
+                    "call_end_time": kwargs[
+                        "end_time"
+                    ],  # datetime object of when call was completed
+                },
+            )
+            self._trace.end_span(
+                outputs={
+                    "message": kwargs["exception"],  # the Exception raised
+                    "usage": response_obj.usage.dict(),  # litellm calculates usage
+                    "cost": kwargs.get(
+                        "response_cost"
+                    ),  # litellm calculates response cost
+                },
+            )
+
+        async def async_log_stream_event(
+            self, kwargs, response_obj, start_time, end_time
+        ):
+            self._trace.update_span_status(span=self._trace.active_span, value="OK")
+            self._trace.end_span(
+                outputs={
+                    "message": kwargs.get(
+                        "complete_streaming_response"
+                    ),  # the complete streamed response (only set if `completion(..stream=True)`)
+                    "usage": response_obj.usage.dict(),  # litellm calculates usage
+                    "cost": kwargs.get(
+                        "response_cost"
+                    ),  # litellm calculates response cost
+                },
+            )
+
+        async def async_log_success_event(
+            self, kwargs, response_obj, start_time, end_time
+        ):
+            self._trace.update_span_status(span=self._trace.active_span, value="OK")
+            self._trace.end_span(
+                outputs={
+                    "message": response_obj.choices[0].message.content,
+                    "usage": response_obj.usage.dict(),  # litellm calculates usage
+                    "cost": kwargs.get(
+                        "response_cost"
+                    ),  # litellm calculates response cost
+                },
+            )
+
+        async def async_log_failure_event(
+            self, kwargs, response_obj, start_time, end_time
+        ):
+            self._trace.update_span_status(span=self._trace.active_span, value="ERROR")
+            self._trace.set_span_attribute(
+                {
+                    "traceback_exception": kwargs[
+                        "traceback_exception"
+                    ],  # the traceback generated via `traceback.format_exc()`
+                    "call_end_time": kwargs[
+                        "end_time"
+                    ],  # datetime object of when call was completed
+                },
+            )
+            self._trace.end_span(
+                outputs={
+                    "message": kwargs["exception"],  # the Exception raised
+                    "usage": response_obj.usage.dict(),  # litellm calculates usage
+                    "cost": kwargs.get(
+                        "response_cost"
+                    ),  # litellm calculates response cost
+                },
+            )
+
+    return LitellmHandler()

agenta/sdk/tracing/llm_tracing.py

@@ -1,16 +1,21 @@
 import os
 from threading import Lock
 from datetime import datetime, timezone
-from typing import Optional, Dict, Any, List, Union
+from typing import Optional, Dict, Any, List
 
 from agenta.sdk.tracing.logger import llm_logger
 from agenta.sdk.tracing.tasks_manager import TaskQueue
 from agenta.client.backend.client import AsyncAgentaApi
 from agenta.client.backend.client import AsyncObservabilityClient
-from agenta.client.backend.types.create_span import CreateSpan, SpanKind, SpanStatusCode
+from agenta.client.backend.types.create_span import (
+    CreateSpan,
+    LlmTokens,
+    SpanStatusCode,
+)
 
 from bson.objectid import ObjectId
 
+
 VARIANT_TRACKING_FEATURE_FLAG = False
 
 
@@ -92,6 +97,13 @@ class Tracing(metaclass=SingletonMeta):
         self,
         attributes: Dict[str, Any] = {},
     ):
+        """
+        Set attributes for the active span.
+
+        Args:
+            attributes (Dict[str, Any], optional): A dictionary of attributes to set. Defaults to {}.
+        """
+
         if (
             self.active_span is None
         ):  # This is the case where entrypoint wants to save the trace information but the parent span has not been initialized yet
@@ -99,7 +111,7 @@ class Tracing(metaclass=SingletonMeta):
                 self.trace_config_cache[key] = value
         else:
             for key, value in attributes.items():
-                self.active_span.attributes[key] = value
+                self.active_span.attributes[key] = value  # type: ignore
 
     def set_trace_tags(self, tags: List[str]):
         self.tags.extend(tags)
@@ -150,11 +162,12 @@ class Tracing(metaclass=SingletonMeta):
             )
             if VARIANT_TRACKING_FEATURE_FLAG:
                 # TODO: we should get the variant_id and variant_name (and environment) from the config object
-                span.variant_id = config.variant_id
-                span.variant_name = (config.variant_name,)
+                span.variant_id = config.variant_id  # type: ignore
+                span.variant_name = (config.variant_name,)  # type: ignore
 
         else:
-            span.parent_span_id = self.active_span.id
+            span.parent_span_id = self.active_span.id  # type: ignore
+
         self.span_dict[span.id] = span
         self.active_span = span
 
@@ -164,28 +177,111 @@ class Tracing(metaclass=SingletonMeta):
     def update_span_status(self, span: CreateSpan, value: str):
         span.status = value
 
+    def _update_span_cost(self, span: CreateSpan, cost: Optional[float]):
+        if cost is not None and isinstance(cost, float):
+            if span.cost is None:
+                span.cost = cost
+            else:
+                span.cost += cost
+
+    def _update_span_tokens(self, span: CreateSpan, tokens: Optional[dict]):
+        if isinstance(tokens, LlmTokens):
+            tokens = tokens.dict()
+        if tokens is not None and isinstance(tokens, dict):
+            if span.tokens is None:
+                span.tokens = LlmTokens(**tokens)
+            else:
+                span.tokens.prompt_tokens += tokens["prompt_tokens"]
+                span.tokens.completion_tokens += tokens["completion_tokens"]
+                span.tokens.total_tokens += tokens["total_tokens"]
+
     def end_span(self, outputs: Dict[str, Any]):
         """
         Ends the active span, if it is a parent span, ends the trace too.
+
+        Args:
+            outputs (Dict[str, Any]): A dictionary containing the outputs of the span.
+                It should have the following keys:
+                - "message" (str): The message output of the span.
+                - "cost" (Optional[Any]): The cost of the span.
+                - "usage" (Optional[Any]): The number of tokens used in the span.
+
+        Raises:
+            ValueError: If there is no active span to end.
+
+        Returns:
+            None
         """
+
         if self.active_span is None:
             raise ValueError("There is no active span to end.")
+
         self.active_span.end_time = datetime.now(timezone.utc)
         self.active_span.outputs = [outputs.get("message", "")]
-        self.active_span.cost = outputs.get("cost", None)
-        self.active_span.tokens = outputs.get("usage", None)
+        if self.active_span.spankind in [
+            "LLM",
+            "RETRIEVER",
+        ]:  # TODO: Remove this whole part. Setting the cost should be done through set_span_attribute
+            self._update_span_cost(self.active_span, outputs.get("cost", None))
+            self._update_span_tokens(self.active_span, outputs.get("usage", None))
 
         # Push span to list of recorded spans
         self.pending_spans.append(self.active_span)
-        self.llm_logger.info(
-            f"Pushed {self.active_span.spankind} span {self.active_span.id} to recorded spans."
-        )
-        if self.active_span.parent_span_id is None:
+
+        active_span_parent_id = self.active_span.parent_span_id
+        if (
+            self.active_span.status == SpanStatusCode.ERROR.value
+            and active_span_parent_id is not None
+        ):
+            self.record_exception_and_end_trace(span_parent_id=active_span_parent_id)
+
+        if active_span_parent_id is None:
             self.end_trace(parent_span=self.active_span)
+
         else:
-            self.active_span = self.span_dict[self.active_span.parent_span_id]
+            parent_span = self.span_dict[active_span_parent_id]
+            self._update_span_cost(parent_span, self.active_span.cost)
+            self._update_span_tokens(parent_span, self.active_span.tokens)
+            self.active_span = parent_span
+
+    def record_exception_and_end_trace(self, span_parent_id: str):
+        """
+        Record an exception and end the trace.
+
+        Args:
+            span_parent_id (str): The ID of the parent span.
+
+        Returns:
+            None
+        """
+
+        parent_span = self.span_dict.get(span_parent_id)
+        if parent_span is not None:
+            # Update parent span of active span
+            parent_span.outputs = self.active_span.outputs  # type: ignore
+            parent_span.status = "ERROR"
+            parent_span.end_time = datetime.now(timezone.utc)
+
+            # Push parent span to list of recorded spans and end trace
+            self.pending_spans.append(parent_span)
+            self.end_trace(parent_span=parent_span)
+
+        # TODO: improve exception logic here.
 
     def end_trace(self, parent_span: CreateSpan):
+        """
+        Ends the active trace and sends the recorded spans for processing.
+
+        Args:
+            parent_span (CreateSpan): The parent span of the trace.
+
+        Raises:
+            RuntimeError: If there is no active trace to end.
+
+        Returns:
+            None
+        """
+
         if self.api_key == "":
             return
 

{agenta-0.17.1.dist-info → agenta-0.17.3.dist-info}/METADATA

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

{agenta-0.17.1.dist-info → agenta-0.17.3.dist-info}/RECORD

@@ -1,4 +1,4 @@
-agenta/__init__.py,sha256=O7kq034_S58b2yzrtR_m3Qd-0t7FAh9q7Z7_HsoD-9Y,796
+agenta/__init__.py,sha256=geaQThRuZzebZncauJW2SZzKqIT7-ZkIJJf-zZlWHCs,833
 agenta/cli/evaluation_commands.py,sha256=fs6492tprPId9p8eGO02Xy-NCBm2RZNJLZWcUxugwd8,474
 agenta/cli/helper.py,sha256=vRxHyeNaltzNIGrfU2vO0H28_rXDzx9QqIZ_S-W6zL4,6212
 agenta/cli/main.py,sha256=Wz0ODhoeKK3Qg_CFUhu6D909szk05tc8ZVBB6H1-w7k,9763
@@ -85,7 +85,7 @@ agenta/client/backend/types/image.py,sha256=p7Vmp7HlMV3YyXe8SFdXYJjCbPNIypW6NfVG
 agenta/client/backend/types/invite_request.py,sha256=1nJTUHspzw2WYpUSd4UUtRnjDHM-dqDBvYewgw-hCQE,993
 agenta/client/backend/types/list_api_keys_response.py,sha256=ZNh7jKwHEMKNp8OV5WJ5XxtKn39DwqK1f8vlFKl54x4,1097
 agenta/client/backend/types/llm_run_rate_limit.py,sha256=mfT4lTczPxrJvd8ZCOAjPvw58QoM151p_uZT0PWNOJ4,1045
-agenta/client/backend/types/llm_tokens.py,sha256=PzJZM5RdXCHGDOkzCN2jQc3iAvCdr0iazpsrNUOE6kk,1069
+agenta/client/backend/types/llm_tokens.py,sha256=J236Fgmz5TeFO0MQA1ZA1QozvR6d3kt8aEUyWkq3jLI,1070
 agenta/client/backend/types/new_human_evaluation.py,sha256=lIgMjVccSp22RRfMGGLH4-yKjMtJeQvjhlwX9EtAxmY,1150
 agenta/client/backend/types/new_testset.py,sha256=9NOC1-f_UZASy4ptzidLNcRU6Odq609ayvSQxEva-40,1009
 agenta/client/backend/types/organization.py,sha256=vJf6Gbz8WCnqabPQmt_t_gfrWPpuvTXgTxKCJKJsrmc,1218
@@ -126,16 +126,18 @@ agenta/docker/docker-assets/entrypoint.sh,sha256=29XK8VQjQsx4hN2j-4JDy-6kQb5y4LC
 agenta/docker/docker-assets/lambda_function.py,sha256=h4UZSSfqwpfsCgERv6frqwm_4JrYu9rLz3I-LxCfeEg,83
 agenta/docker/docker-assets/main.py,sha256=7MI-21n81U7N7A0GxebNi0cmGWtJKcR2sPB6FcH2QfA,251
 agenta/docker/docker_utils.py,sha256=5uHMCzXkCvIsDdEiwbnnn97KkzsFbBvyMwogCsv_Z5U,3509
-agenta/sdk/__init__.py,sha256=oHgl-qoEyi3d2VI_Kv-rIMSx9zgs6b5MP62PLq5GqYI,762
+agenta/sdk/__init__.py,sha256=cF0de6DiH-NZWEm0XvPN8_TeC1whPBnDf1WYYE1qK2g,762
 agenta/sdk/agenta_init.py,sha256=j7qwyDtXfLozWpnayJHPz2aQOzHSGvHo9V6s0FXeUe8,9937
 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=o1kD14dfXLV3p1OyzgCUA6mIoyDV_YuW__kfXaAKW_I,22754
-agenta/sdk/decorators/tracing.py,sha256=bC-YlPQUrHBEqvhLJxr63N0qlo1jvrbt7ro2AMGXXZw,3160
+agenta/sdk/decorators/tracing.py,sha256=_6lbnhoJgz6q6OkMDNWoLstuv-ZvTz9ExcetW2sTiEY,3447
 agenta/sdk/router.py,sha256=0sbajvn5C7t18anH6yNo7-oYxldHnYfwcbmQnIXBePw,269
+agenta/sdk/tracing/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+agenta/sdk/tracing/callbacks.py,sha256=E5cA6dzut90WSQ38Zy2y1XkiP6wRQ1j1CGcEo5oCsWc,6312
 agenta/sdk/tracing/context_manager.py,sha256=HskDaiORoOhjeN375gm05wYnieQzh5UnoIsnSAHkAyc,252
-agenta/sdk/tracing/llm_tracing.py,sha256=PmMYQ5N8atYut85Rk2hZ1jmvSF80Duuy6Clf7URcTCA,8193
+agenta/sdk/tracing/llm_tracing.py,sha256=fHtc1tKgdR2CDnFMrfDIxSvyjwFW3Qk2RE8P-y_tCM0,11361
 agenta/sdk/tracing/logger.py,sha256=GfH7V-jBHcn7h5dbdrnkDMe_ml3wkXFBeoQiqR4KVRc,474
 agenta/sdk/tracing/tasks_manager.py,sha256=ROrWIaqS2J2HHiJtRWiHKlLY8CCsqToP5VeXu7mamck,3748
 agenta/sdk/types.py,sha256=KMnQUOdjaHSWctDLIiMHnk0o3c-C47Vm4Mn2kIZ88YI,5740
@@ -157,7 +159,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.17.1.dist-info/METADATA,sha256=iPlbIzVHb16h_3Pnuik3WLtwPYCOtNi61kqvM3vNhuY,26458
-agenta-0.17.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-agenta-0.17.1.dist-info/entry_points.txt,sha256=PDiu8_8AsL7ibU9v4iNoOKR1S7F2rdxjlEprjM9QOgo,46
-agenta-0.17.1.dist-info/RECORD,,
+agenta-0.17.3.dist-info/METADATA,sha256=EwPxAqK2BdKSykVU91MNVPsShfg9hWqrs9RyHXEWr_g,26458
+agenta-0.17.3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+agenta-0.17.3.dist-info/entry_points.txt,sha256=PDiu8_8AsL7ibU9v4iNoOKR1S7F2rdxjlEprjM9QOgo,46
+agenta-0.17.3.dist-info/RECORD,,