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

agenta/sdk/decorators/tracing.py

@@ -1,131 +1,250 @@
-# Stdlib Imports
-import inspect
-import traceback
+from typing import Callable, Optional, Any, Dict, List, Union
 from functools import wraps
-from typing import Any, Callable, Optional, List, Union
+from itertools import chain
+from inspect import iscoroutinefunction, getfullargspec
 
-# Own Imports
-import agenta as ag
-from agenta.sdk.decorators.base import BaseDecorator
-from agenta.sdk.tracing.logger import llm_logger as logging
-from agenta.sdk.utils.debug import debug, DEBUG, SHIFT
-
-
-logging.setLevel("DEBUG")
-
-
-class instrument(BaseDecorator):
-    """Decorator class for monitoring llm apps functions.
+from agenta.sdk.utils.exceptions import suppress
+from agenta.sdk.context.tracing import tracing_context
+from agenta.sdk.tracing.conventions import parse_span_kind
 
-    Args:
-        BaseDecorator (object): base decorator class
-
-    Example:
-    ```python
-        import agenta as ag
-
-        prompt_config = {"system_prompt": ..., "temperature": 0.5, "max_tokens": ...}
+import agenta as ag
 
-        @ag.instrument(spankind="llm")
-        async def litellm_openai_call(prompt:str) -> str:
-            return "do something"
 
-        @ag.instrument(config=prompt_config) # spankind for parent span defaults to workflow
-        async def generate(prompt: str):
-            return ...
-    ```
-    """
+class instrument:  # pylint: disable=invalid-name
+    DEFAULT_KEY = "__default__"
 
     def __init__(
         self,
-        config: Optional[dict] = None,
-        spankind: str = "workflow",
-        ignore_inputs: Union[List[str], bool] = False,
-        ignore_outputs: Union[List[str], bool] = False,
+        type: str = "task",  # pylint: disable=redefined-builtin
+        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",
+        spankind: Optional[str] = "TASK",
     ) -> None:
+        self.type = spankind or kind or type
+        self.kind = None
         self.config = config
-        self.spankind = spankind
         self.ignore_inputs = ignore_inputs
         self.ignore_outputs = ignore_outputs
+        self.max_depth = max_depth
 
     def __call__(self, func: Callable[..., Any]):
-        is_coroutine_function = inspect.iscoroutinefunction(func)
-
-        def get_inputs(*args, **kwargs):
-            func_args = inspect.getfullargspec(func).args
-            input_dict = {name: value for name, value in zip(func_args, args)}
-            input_dict.update(kwargs)
-
-            return input_dict
-
-        def redact(io, blacklist):
-            return {
-                key: io[key]
-                for key in io.keys()
-                if key
-                not in (
-                    blacklist
-                    if isinstance(blacklist, list)
-                    else []
-                    if blacklist is False
-                    else io.keys()
-                )
-            }
-
-        def patch(result):
-            TRACE_DEFAULT_KEY = "__default__"
-
-            outputs = result
-
-            # PATCH : if result is not a dict, make it a dict
-            if not isinstance(result, dict):
-                outputs = {TRACE_DEFAULT_KEY: result}
-            else:
-                # PATCH : if result is a legacy dict, clean it up
-                if (
-                    "message" in result.keys()
-                    and "cost" in result.keys()
-                    and "usage" in result.keys()
-                ):
-                    outputs = {TRACE_DEFAULT_KEY: result["message"]}
-
-                    ag.tracing.store_cost(result["cost"])
-                    ag.tracing.store_usage(result["usage"])
-
-            return outputs
+        is_coroutine_function = iscoroutinefunction(func)
 
         @wraps(func)
         async def async_wrapper(*args, **kwargs):
-            async def wrapped_func(*args, **kwargs):
-                with ag.tracing.Context(
-                    name=func.__name__,
-                    input=redact(get_inputs(*args, **kwargs), self.ignore_inputs),
-                    spankind=self.spankind,
-                    config=self.config,
-                ):
+            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)
+
                     result = await func(*args, **kwargs)
 
-                    ag.tracing.store_outputs(redact(patch(result), self.ignore_outputs))
+                    self._post_instrument(result)
 
                     return result
 
-            return await wrapped_func(*args, **kwargs)
+            return await _async_auto_instrumented(*args, **kwargs)
 
         @wraps(func)
         def sync_wrapper(*args, **kwargs):
-            def wrapped_func(*args, **kwargs):
-                with ag.tracing.Context(
-                    name=func.__name__,
-                    input=redact(get_inputs(*args, **kwargs), self.ignore_inputs),
-                    spankind=self.spankind,
-                    config=self.config,
-                ):
+            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)
+
                     result = func(*args, **kwargs)
 
-                    ag.tracing.store_outputs(redact(patch(result), self.ignore_outputs))
+                    self._post_instrument(result)
 
                     return result
 
-            return wrapped_func(*args, **kwargs)
+            return _sync_auto_instrumented(*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/__init__.py

@@ -0,0 +1 @@
+from .litellm import litellm_handler

agenta/sdk/litellm/litellm.py

@@ -0,0 +1,288 @@
+from opentelemetry.trace import SpanKind
+
+import agenta as ag
+
+from agenta.sdk.tracing.spans import CustomSpan
+from agenta.sdk.utils.exceptions import suppress  # TODO: use it !
+from agenta.sdk.utils.logging import log
+
+
+def litellm_handler():
+    try:
+        from litellm.integrations.custom_logger import (  # pylint: disable=import-outside-toplevel
+            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(  # pylint: disable=broad-exception-raised
+            f"Unexpected error occurred when importing litellm: {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.
+        """
+
+        def __init__(self):
+            super().__init__()
+
+            self.span = None
+
+        def log_pre_api_call(
+            self,
+            model,
+            messages,
+            kwargs,
+        ):
+            type = (  # pylint: disable=redefined-builtin
+                "chat"
+                if kwargs.get("call_type") in ["completion", "acompletion"]
+                else "embedding"
+            )
+
+            kind = SpanKind.CLIENT
+
+            self.span = CustomSpan(
+                ag.tracer.start_span(name=f"litellm_{kind.name.lower()}", kind=kind)
+            )
+
+            self.span.set_attributes(
+                attributes={"node": type},
+                namespace="type",
+            )
+
+            if not self.span:
+                log.error("LiteLLM callback error: span not found.")
+                return
+
+            self.span.set_attributes(
+                attributes={"inputs": {"prompt": kwargs["messages"]}},
+                namespace="data",
+            )
+
+            self.span.set_attributes(
+                attributes={
+                    "configuration": {
+                        "model": kwargs.get("model"),
+                        **kwargs.get("optional_params"),
+                    }
+                },
+                namespace="meta",
+            )
+
+        def log_stream_event(
+            self,
+            kwargs,
+            response_obj,
+            start_time,
+            end_time,
+        ):
+            if not self.span:
+                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
+            )
+
+            self.span.set_attributes(
+                attributes={"outputs": outputs},
+                namespace="data",
+            )
+
+            self.span.set_attributes(
+                attributes={"total": kwargs.get("response_cost")},
+                namespace="metrics.unit.costs",
+            )
+
+            self.span.set_attributes(
+                attributes=(
+                    {
+                        "prompt": response_obj.usage.prompt_tokens,
+                        "completion": response_obj.usage.completion_tokens,
+                        "total": response_obj.usage.total_tokens,
+                    }
+                ),
+                namespace="metrics.unit.tokens",
+            )
+
+            self.span.set_status(status="OK")
+
+            self.span.end()
+
+        def log_success_event(
+            self,
+            kwargs,
+            response_obj,
+            start_time,
+            end_time,
+        ):
+            if not self.span:
+                log.error("LiteLLM callback error: span not found.")
+                return
+
+            try:
+                result = []
+                for choice in response_obj.choices:
+                    message = choice.message.__dict__
+                    result.append(message)
+
+                outputs = {"completion": result}
+                self.span.set_attributes(
+                    attributes={"outputs": outputs},
+                    namespace="data",
+                )
+
+            except Exception as e:
+                pass
+
+            self.span.set_attributes(
+                attributes={"total": kwargs.get("response_cost")},
+                namespace="metrics.unit.costs",
+            )
+
+            self.span.set_attributes(
+                attributes=(
+                    {
+                        "prompt": response_obj.usage.prompt_tokens,
+                        "completion": response_obj.usage.completion_tokens,
+                        "total": response_obj.usage.total_tokens,
+                    }
+                ),
+                namespace="metrics.unit.tokens",
+            )
+
+            self.span.set_status(status="OK")
+
+            self.span.end()
+
+        def log_failure_event(
+            self,
+            kwargs,
+            response_obj,
+            start_time,
+            end_time,
+        ):
+            if not self.span:
+                log.error("LiteLLM callback error: span not found.")
+                return
+
+            self.span.record_exception(kwargs["exception"])
+
+            self.span.set_status(status="ERROR")
+
+            self.span.end()
+
+        async def async_log_stream_event(
+            self,
+            kwargs,
+            response_obj,
+            start_time,
+            end_time,
+        ):
+            if not self.span:
+                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
+            )
+
+            self.span.set_attributes(
+                attributes={"outputs": outputs},
+                namespace="data",
+            )
+
+            self.span.set_attributes(
+                attributes={"total": kwargs.get("response_cost")},
+                namespace="metrics.unit.costs",
+            )
+
+            self.span.set_attributes(
+                attributes=(
+                    {
+                        "prompt": response_obj.usage.prompt_tokens,
+                        "completion": response_obj.usage.completion_tokens,
+                        "total": response_obj.usage.total_tokens,
+                    }
+                ),
+                namespace="metrics.unit.tokens",
+            )
+
+            self.span.set_status(status="OK")
+
+            self.span.end()
+
+        async def async_log_success_event(
+            self,
+            kwargs,
+            response_obj,
+            start_time,
+            end_time,
+        ):
+            if not self.span:
+                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
+            )
+
+            self.span.set_attributes(
+                attributes={"outputs": outputs},
+                namespace="data",
+            )
+
+            self.span.set_attributes(
+                attributes={"total": kwargs.get("response_cost")},
+                namespace="metrics.unit.costs",
+            )
+
+            self.span.set_attributes(
+                attributes=(
+                    {
+                        "prompt": response_obj.usage.prompt_tokens,
+                        "completion": response_obj.usage.completion_tokens,
+                        "total": response_obj.usage.total_tokens,
+                    }
+                ),
+                namespace="metrics.unit.tokens",
+            )
+
+            self.span.set_status(status="OK")
+
+            self.span.end()
+
+        async def async_log_failure_event(
+            self,
+            kwargs,
+            response_obj,
+            start_time,
+            end_time,
+        ):
+            if not self.span:
+                log.error("LiteLLM callback error: span not found.")
+                return
+
+            self.span.record_exception(kwargs["exception"])
+
+            self.span.set_status(status="ERROR")
+
+            self.span.end()
+
+    return LitellmHandler()

agenta/sdk/managers/__init__.py

@@ -0,0 +1,6 @@
+from agenta.sdk.managers.config import ConfigManager
+from agenta.sdk.managers.variant import VariantManager
+from agenta.sdk.managers.deployment import DeploymentManager
+
+
+__all__ = ["ConfigManager", "VariantManager", "DeploymentManager"]