agenta 0.27.0__py3-none-any.whl → 0.27.0a0__py3-none-any.whl

agenta/sdk/decorators/tracing.py

@@ -1,250 +1,283 @@
-from typing import Callable, Optional, Any, Dict, List, Union
+import inspect
+import traceback
 from functools import wraps
 from itertools import chain
-from inspect import iscoroutinefunction, getfullargspec
+from typing import Callable, Optional, Any, Dict, List
+
+import agenta as ag
+
 
 from agenta.sdk.utils.exceptions import suppress
+
 from agenta.sdk.context.tracing import tracing_context
 from agenta.sdk.tracing.conventions import parse_span_kind
 
-import agenta as ag
-
 
-class instrument:  # pylint: disable=invalid-name
+class instrument:
     DEFAULT_KEY = "__default__"
 
     def __init__(
         self,
-        type: str = "task",  # pylint: disable=redefined-builtin
+        kind: str = "task",
         config: Optional[Dict[str, Any]] = None,
         ignore_inputs: Optional[bool] = None,
         ignore_outputs: Optional[bool] = None,
         max_depth: Optional[int] = 2,
-        # DEPRECATING
-        kind: str = "task",
+        # DEPRECATED
         spankind: Optional[str] = "TASK",
     ) -> None:
-        self.type = spankind or kind or type
-        self.kind = None
+        self.kind = spankind if spankind is not None else kind
         self.config = config
         self.ignore_inputs = ignore_inputs
         self.ignore_outputs = ignore_outputs
         self.max_depth = max_depth
 
     def __call__(self, func: Callable[..., Any]):
-        is_coroutine_function = iscoroutinefunction(func)
+        is_coroutine_function = inspect.iscoroutinefunction(func)
+
+        def parse(*args, **kwargs) -> Dict[str, Any]:
+            inputs = {
+                key: value
+                for key, value in chain(
+                    zip(inspect.getfullargspec(func).args, args),
+                    kwargs.items(),
+                )
+            }
+
+            return inputs
+
+        def redact(
+            io: Dict[str, Any], ignore: List[str] | bool = False
+        ) -> Dict[str, Any]:
+            """
+            Redact user-defined sensitive information from inputs and outputs as defined by the ignore list or boolean flag.
+
+            Example:
+            - ignore = ["password"] -> {"username": "admin", "password": "********"} -> {"username": "admin"}
+            - ignore = True -> {"username": "admin", "password": "********"} -> {}
+            - ignore = False -> {"username": "admin", "password": "********"} -> {"username": "admin", "password": "********"}
+            """
+            io = {
+                key: value
+                for key, value in io.items()
+                if key
+                not in (
+                    ignore
+                    if isinstance(ignore, list)
+                    else io.keys()
+                    if ignore is True
+                    else []
+                )
+            }
+
+            return io
+
+        def patch(result: Any) -> Dict[str, Any]:
+            """
+            Patch the result to ensure that it is a dictionary, with a default key when necessary.
+
+            Example:
+            - result = "Hello, World!" -> {"__default__": "Hello, World!"}
+            - result = {"message": "Hello, World!", "cost": 0.0, "usage": {}} -> {"__default__": "Hello, World!"}
+            - result = {"message": "Hello, World!"} -> {"message": "Hello, World!"}
+            """
+            outputs = (
+                {instrument.DEFAULT_KEY: result}
+                if not isinstance(result, dict)
+                else (
+                    {instrument.DEFAULT_KEY: result["message"]}
+                    if all(key in result for key in ["message", "cost", "usage"])
+                    else result
+                )
+            )
+
+            return outputs
 
         @wraps(func)
         async def async_wrapper(*args, **kwargs):
-            async def _async_auto_instrumented(*args, **kwargs):
-                self._parse_type_and_kind()
-
-                with ag.tracer.start_as_current_span(func.__name__, kind=self.kind):
-                    self._pre_instrument(func, *args, **kwargs)
+            async def wrapped_func(*args, **kwargs):
+                if not ag.tracing.get_current_span().is_recording():
+                    self.kind = "workflow"
+
+                kind = parse_span_kind(self.kind)
+
+                with ag.tracer.start_as_current_span(func.__name__, kind=kind):
+                    span = ag.tracing.get_current_span()
+
+                    with suppress():
+                        span.set_attributes(
+                            attributes={"node": self.kind},
+                            namespace="type",
+                        )
+
+                        if span.parent is None:
+                            rctx = tracing_context.get()
+
+                            span.set_attributes(
+                                attributes={"configuration": rctx.get("config", {})},
+                                namespace="meta",
+                            )
+                            span.set_attributes(
+                                attributes={"environment": rctx.get("environment", {})},
+                                namespace="meta",
+                            )
+                            span.set_attributes(
+                                attributes={"version": rctx.get("version", {})},
+                                namespace="meta",
+                            )
+                            span.set_attributes(
+                                attributes={"variant": rctx.get("variant", {})},
+                                namespace="meta",
+                            )
+
+                        _inputs = redact(parse(*args, **kwargs), self.ignore_inputs)
+                        span.set_attributes(
+                            attributes={"inputs": _inputs},
+                            namespace="data",
+                            max_depth=self.max_depth,
+                        )
 
                     result = await func(*args, **kwargs)
 
-                    self._post_instrument(result)
+                    with suppress():
+                        cost = None
+                        usage = {}
+
+                        if isinstance(result, dict):
+                            cost = result.get("cost", None)
+                            usage = result.get("usage", {})
+
+                        span.set_attributes(
+                            attributes={"total": cost},
+                            namespace="metrics.unit.costs",
+                        )
+                        span.set_attributes(
+                            attributes=(
+                                {
+                                    "prompt": usage.get("prompt_tokens", None),
+                                    "completion": usage.get("completion_tokens", None),
+                                    "total": usage.get("total_tokens", None),
+                                }
+                            ),
+                            namespace="metrics.unit.tokens",
+                        )
+
+                        _outputs = redact(patch(result), self.ignore_outputs)
+                        span.set_attributes(
+                            attributes={"outputs": _outputs},
+                            namespace="data",
+                            max_depth=self.max_depth,
+                        )
+
+                        span.set_status("OK")
+
+                    with suppress():
+                        if hasattr(span, "parent") and span.parent is None:
+                            tracing_context.set(
+                                tracing_context.get()
+                                | {
+                                    "root": {
+                                        "trace_id": span.get_span_context().trace_id,
+                                        "span_id": span.get_span_context().span_id,
+                                    }
+                                }
+                            )
 
                     return result
 
-            return await _async_auto_instrumented(*args, **kwargs)
+            return await wrapped_func(*args, **kwargs)
 
         @wraps(func)
         def sync_wrapper(*args, **kwargs):
-            def _sync_auto_instrumented(*args, **kwargs):
-                self._parse_type_and_kind()
-
-                with ag.tracer.start_as_current_span(func.__name__, kind=self.kind):
-                    self._pre_instrument(func, *args, **kwargs)
+            def wrapped_func(*args, **kwargs):
+                if not ag.tracing.get_current_span().is_recording():
+                    self.kind = "workflow"
+
+                kind = parse_span_kind(self.kind)
+
+                with ag.tracer.start_as_current_span(func.__name__, kind=kind):
+                    span = ag.tracing.get_current_span()
+
+                    with suppress():
+                        span.set_attributes(
+                            attributes={"node": self.kind},
+                            namespace="type",
+                        )
+
+                        if span.parent is None:
+                            rctx = tracing_context.get()
+
+                            span.set_attributes(
+                                attributes={"configuration": rctx.get("config", {})},
+                                namespace="meta",
+                            )
+                            span.set_attributes(
+                                attributes={"environment": rctx.get("environment", {})},
+                                namespace="meta",
+                            )
+                            span.set_attributes(
+                                attributes={"version": rctx.get("version", {})},
+                                namespace="meta",
+                            )
+                            span.set_attributes(
+                                attributes={"variant": rctx.get("variant", {})},
+                                namespace="meta",
+                            )
+
+                        _inputs = redact(parse(*args, **kwargs), self.ignore_inputs)
+                        span.set_attributes(
+                            attributes={"inputs": _inputs},
+                            namespace="data",
+                            max_depth=self.max_depth,
+                        )
 
                     result = func(*args, **kwargs)
 
-                    self._post_instrument(result)
+                    with suppress():
+                        cost = None
+                        usage = {}
+                        if isinstance(result, dict):
+                            cost = result.get("cost", None)
+                            usage = result.get("usage", {})
+
+                        span.set_attributes(
+                            attributes={"total": cost},
+                            namespace="metrics.unit.costs",
+                        )
+                        span.set_attributes(
+                            attributes=(
+                                {
+                                    "prompt": usage.get("prompt_tokens", None),
+                                    "completion": usage.get("completion_tokens", None),
+                                    "total": usage.get("total_tokens", None),
+                                }
+                            ),
+                            namespace="metrics.unit.tokens",
+                        )
+
+                        _outputs = redact(patch(result), self.ignore_outputs)
+                        span.set_attributes(
+                            attributes={"outputs": _outputs},
+                            namespace="data",
+                            max_depth=self.max_depth,
+                        )
+
+                        span.set_status("OK")
+
+                    with suppress():
+                        if hasattr(span, "parent") and span.parent is None:
+                            tracing_context.set(
+                                tracing_context.get()
+                                | {
+                                    "root": {
+                                        "trace_id": span.get_span_context().trace_id,
+                                        "span_id": span.get_span_context().span_id,
+                                    }
+                                }
+                            )
 
                     return result
 
-            return _sync_auto_instrumented(*args, **kwargs)
+            return wrapped_func(*args, **kwargs)
 
         return async_wrapper if is_coroutine_function else sync_wrapper
-
-    def _parse_type_and_kind(self):
-        if not ag.tracing.get_current_span().is_recording():
-            self.type = "workflow"
-
-        self.kind = parse_span_kind(self.type)
-
-    def _pre_instrument(
-        self,
-        func,
-        *args,
-        **kwargs,
-    ):
-        span = ag.tracing.get_current_span()
-
-        with suppress():
-            span.set_attributes(
-                attributes={"node": self.type},
-                namespace="type",
-            )
-
-            if span.parent is None:
-                rctx = tracing_context.get()
-
-                span.set_attributes(
-                    attributes={"configuration": rctx.get("config", {})},
-                    namespace="meta",
-                )
-                span.set_attributes(
-                    attributes={"environment": rctx.get("environment", {})},
-                    namespace="meta",
-                )
-                span.set_attributes(
-                    attributes={"version": rctx.get("version", {})},
-                    namespace="meta",
-                )
-                span.set_attributes(
-                    attributes={"variant": rctx.get("variant", {})},
-                    namespace="meta",
-                )
-
-            _inputs = self._redact(
-                self._parse(
-                    func,
-                    *args,
-                    **kwargs,
-                ),
-                self.ignore_inputs,
-            )
-            span.set_attributes(
-                attributes={"inputs": _inputs},
-                namespace="data",
-                max_depth=self.max_depth,
-            )
-
-    def _post_instrument(
-        self,
-        result,
-    ):
-        span = ag.tracing.get_current_span()
-        with suppress():
-            cost = None
-            usage = {}
-
-            if isinstance(result, dict):
-                cost = result.get("cost", None)
-                usage = result.get("usage", {})
-
-            if isinstance(usage, (int, float)):
-                usage = {"total_tokens": usage}
-
-            span.set_attributes(
-                attributes={"total": cost},
-                namespace="metrics.unit.costs",
-            )
-            span.set_attributes(
-                attributes=(
-                    {
-                        "prompt": usage.get("prompt_tokens", None),
-                        "completion": usage.get("completion_tokens", None),
-                        "total": usage.get("total_tokens", None),
-                    }
-                ),
-                namespace="metrics.unit.tokens",
-            )
-
-            _outputs = self._redact(self._patch(result), self.ignore_outputs)
-            span.set_attributes(
-                attributes={"outputs": _outputs},
-                namespace="data",
-                max_depth=self.max_depth,
-            )
-
-            span.set_status("OK")
-
-        with suppress():
-            if hasattr(span, "parent") and span.parent is None:
-                tracing_context.set(
-                    tracing_context.get()
-                    | {
-                        "root": {
-                            "trace_id": span.get_span_context().trace_id,
-                            "span_id": span.get_span_context().span_id,
-                        }
-                    }
-                )
-
-    def _parse(
-        self,
-        func,
-        *args,
-        **kwargs,
-    ) -> Dict[str, Any]:
-        inputs = {
-            key: value
-            for key, value in chain(
-                zip(getfullargspec(func).args, args),
-                kwargs.items(),
-            )
-        }
-
-        return inputs
-
-    def _redact(
-        self,
-        io: Dict[str, Any],
-        ignore: Union[List[str], bool] = False,
-    ) -> Dict[str, Any]:
-        """
-        Redact user-defined sensitive information
-        from inputs and outputs as defined by the ignore list or boolean flag.
-
-        Example:
-        - ignore = ["password"] -> {"username": "admin", "password": "********"}
-            -> {"username": "admin"}
-        - ignore = True -> {"username": "admin", "password": "********"}
-            -> {}
-        - ignore = False -> {"username": "admin", "password": "********"}
-            -> {"username": "admin", "password": "********"}
-        """
-        io = {
-            key: value
-            for key, value in io.items()
-            if key
-            not in (
-                ignore
-                if isinstance(ignore, list)
-                else io.keys()
-                if ignore is True
-                else []
-            )
-        }
-
-        return io
-
-    def _patch(
-        self,
-        result: Any,
-    ) -> Dict[str, Any]:
-        """
-        Patch the result to ensure that it is a dictionary, with a default key when necessary.
-
-        Example:
-        - result = "Hello, World!"
-            -> {"__default__": "Hello, World!"}
-        - result = {"message": "Hello, World!", "cost": 0.0, "usage": {}}
-            -> {"__default__": "Hello, World!"}
-        - result = {"message": "Hello, World!"}
-            -> {"message": "Hello, World!"}
-        """
-        outputs = (
-            {instrument.DEFAULT_KEY: result}
-            if not isinstance(result, dict)
-            else (
-                {instrument.DEFAULT_KEY: result["message"]}
-                if all(key in result for key in ["message", "cost", "usage"])
-                else result
-            )
-        )
-
-        return outputs

agenta/sdk/litellm/litellm.py

@@ -1,15 +1,16 @@
-from opentelemetry.trace import SpanKind
-
 import agenta as ag
 
+from opentelemetry.trace import SpanKind
+
 from agenta.sdk.tracing.spans import CustomSpan
-from agenta.sdk.utils.exceptions import suppress  # TODO: use it !
+from agenta.sdk.utils.exceptions import suppress
 from agenta.sdk.utils.logging import log
 
 
 def litellm_handler():
     try:
-        from litellm.integrations.custom_logger import (  # pylint: disable=import-outside-toplevel
+        from litellm.utils import ModelResponse
+        from litellm.integrations.custom_logger import (
             CustomLogger as LitellmCustomLogger,
         )
     except ImportError as exc:
@@ -17,23 +18,18 @@ def litellm_handler():
             "The litellm SDK is not installed. Please install it using `pip install litellm`."
         ) from exc
     except Exception as exc:
-        raise Exception(  # pylint: disable=broad-exception-raised
-            f"Unexpected error occurred when importing litellm: {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.
+        """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.
+            LitellmCustomLogger (object): custom logger that allows us to override the events to capture.
         """
 
         def __init__(self):
-            super().__init__()
-
             self.span = None
 
         def log_pre_api_call(
@@ -42,7 +38,7 @@ def litellm_handler():
             messages,
             kwargs,
         ):
-            type = (  # pylint: disable=redefined-builtin
+            type = (
                 "chat"
                 if kwargs.get("call_type") in ["completion", "acompletion"]
                 else "embedding"
@@ -63,8 +59,10 @@ def litellm_handler():
                 log.error("LiteLLM callback error: span not found.")
                 return
 
+            log.info(f"log_pre_api_call({hex(self.span.context.span_id)[2:]})")
+
             self.span.set_attributes(
-                attributes={"inputs": {"prompt": kwargs["messages"]}},
+                attributes={"inputs": {"messages": kwargs["messages"]}},
                 namespace="data",
             )
 
@@ -89,14 +87,12 @@ def litellm_handler():
                 log.error("LiteLLM callback error: span not found.")
                 return
 
-            result = kwargs.get("complete_streaming_response")
-
-            outputs = (
-                {"__default__": result} if not isinstance(result, dict) else result
-            )
+            # log.info(f"log_stream({hex(self.span.context.span_id)[2:]})")
 
             self.span.set_attributes(
-                attributes={"outputs": outputs},
+                attributes={
+                    "output": {"__default__": kwargs.get("complete_streaming_response")}
+                },
                 namespace="data",
             )
 
@@ -131,20 +127,14 @@ def litellm_handler():
                 log.error("LiteLLM callback error: span not found.")
                 return
 
-            try:
-                result = []
-                for choice in response_obj.choices:
-                    message = choice.message.__dict__
-                    result.append(message)
+            # log.info(f"log_success({hex(self.span.context.span_id)[2:]})")
 
-                outputs = {"completion": result}
-                self.span.set_attributes(
-                    attributes={"outputs": outputs},
-                    namespace="data",
-                )
-
-            except Exception as e:
-                pass
+            self.span.set_attributes(
+                attributes={
+                    "output": {"__default__": response_obj.choices[0].message.content}
+                },
+                namespace="data",
+            )
 
             self.span.set_attributes(
                 attributes={"total": kwargs.get("response_cost")},
@@ -177,6 +167,8 @@ def litellm_handler():
                 log.error("LiteLLM callback error: span not found.")
                 return
 
+            # log.info(f"log_failure({hex(self.span.context.span_id)[2:]})")
+
             self.span.record_exception(kwargs["exception"])
 
             self.span.set_status(status="ERROR")
@@ -194,14 +186,12 @@ def litellm_handler():
                 log.error("LiteLLM callback error: span not found.")
                 return
 
-            result = kwargs.get("complete_streaming_response")
-
-            outputs = (
-                {"__default__": result} if not isinstance(result, dict) else result
-            )
+            # log.info(f"async_log_stream({hex(self.span.context.span_id)[2:]})")
 
             self.span.set_attributes(
-                attributes={"outputs": outputs},
+                attributes={
+                    "output": {"__default__": kwargs.get("complete_streaming_response")}
+                },
                 namespace="data",
             )
 
@@ -236,15 +226,12 @@ def litellm_handler():
                 log.error("LiteLLM callback error: span not found.")
                 return
 
-            # result = kwargs.get("complete_streaming_response")
-            result = response_obj.choices[0].message.content
-
-            outputs = (
-                {"__default__": result} if not isinstance(result, dict) else result
-            )
+            log.info(f"async_log_success({hex(self.span.context.span_id)[2:]})")
 
             self.span.set_attributes(
-                attributes={"outputs": outputs},
+                attributes={
+                    "output": {"__default__": kwargs.get("complete_streaming_response")}
+                },
                 namespace="data",
             )
 
@@ -279,6 +266,8 @@ def litellm_handler():
                 log.error("LiteLLM callback error: span not found.")
                 return
 
+            # log.info(f"async_log_failure({hex(self.span.context.span_id)[2:]})")
+
             self.span.record_exception(kwargs["exception"])
 
             self.span.set_status(status="ERROR")