agenta 0.12.3__py3-none-any.whl → 0.32.0a1__py3-none-any.whl

agenta/sdk/decorators/tracing.py

@@ -0,0 +1,296 @@
+from typing import Callable, Optional, Any, Dict, List, Union
+
+from functools import wraps
+from itertools import chain
+from inspect import iscoroutinefunction, getfullargspec
+
+from opentelemetry import baggage as baggage
+from opentelemetry.context import attach, detach
+
+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
+    DEFAULT_KEY = "__default__"
+
+    def __init__(
+        self,
+        type: str = "task",  # pylint: disable=redefined-builtin
+        config: Optional[Dict[str, Any]] = None,
+        ignore_inputs: Optional[bool] = None,
+        ignore_outputs: Optional[bool] = None,
+        redact: Optional[Callable[..., Any]] = None,
+        redact_on_error: Optional[bool] = True,
+        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.ignore_inputs = ignore_inputs
+        self.ignore_outputs = ignore_outputs
+        self.redact = redact
+        self.redact_on_error = redact_on_error
+        self.max_depth = max_depth
+
+    def __call__(self, func: Callable[..., Any]):
+        is_coroutine_function = iscoroutinefunction(func)
+
+        @wraps(func)
+        async def awrapper(*args, **kwargs):
+            async def aauto_instrumented(*args, **kwargs):
+                self._parse_type_and_kind()
+
+                token = self._attach_baggage()
+
+                with ag.tracer.start_as_current_span(func.__name__, kind=self.kind):
+                    self._pre_instrument(func, *args, **kwargs)
+
+                    result = await func(*args, **kwargs)
+
+                    self._post_instrument(result)
+
+                    return result
+
+                self._detach_baggage(token)
+
+            return await aauto_instrumented(*args, **kwargs)
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            def auto_instrumented(*args, **kwargs):
+                self._parse_type_and_kind()
+
+                token = self._attach_baggage()
+
+                with ag.tracer.start_as_current_span(func.__name__, kind=self.kind):
+                    self._pre_instrument(func, *args, **kwargs)
+
+                    result = func(*args, **kwargs)
+
+                    self._post_instrument(result)
+
+                    return result
+
+                self._detach_baggage(token)
+
+            return auto_instrumented(*args, **kwargs)
+
+        return awrapper if is_coroutine_function else 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 _attach_baggage(self):
+        context = tracing_context.get()
+
+        references = context.references
+
+        token = None
+        if references:
+            for k, v in references.items():
+                token = attach(baggage.set_baggage(f"ag.refs.{k}", v))
+
+        return token
+
+    def _detach_baggage(
+        self,
+        token,
+    ):
+        if token:
+            detach(token)
+
+    def _pre_instrument(
+        self,
+        func,
+        *args,
+        **kwargs,
+    ):
+        span = ag.tracing.get_current_span()
+
+        context = tracing_context.get()
+
+        with suppress():
+            trace_id = span.context.trace_id
+
+            ag.tracing.credentials[trace_id] = context.credentials
+
+            span.set_attributes(
+                attributes={"node": self.type},
+                namespace="type",
+            )
+
+            if span.parent is None:
+                span.set_attributes(
+                    attributes={"configuration": context.parameters or {}},
+                    namespace="meta",
+                )
+
+            _inputs = self._redact(
+                name=span.name,
+                field="inputs",
+                io=self._parse(func, *args, **kwargs),
+                ignore=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(
+                name=span.name,
+                field="outputs",
+                io=self._patch(result),
+                ignore=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:
+                context = tracing_context.get()
+                context.link = {
+                    "tree_id": span.get_span_context().trace_id,
+                    "node_id": span.get_span_context().span_id,
+                }
+                tracing_context.set(context)
+
+    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,
+        *,
+        name: str,
+        field: str,
+        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 []
+            )
+        }
+
+        if self.redact is not None:
+            try:
+                io = self.redact(name, field, io)
+            except:  # pylint: disable=bare-except
+                if self.redact_on_error:
+                    io = {}
+
+        if ag.tracing.redact is not None:
+            try:
+                io = ag.tracing.redact(name, field, io)
+            except:  # pylint: disable=bare-except
+                if ag.tracing.redact_on_error:
+                    io = {}
+
+        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,314 @@
+from typing import Dict
+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: Dict[str, CustomSpan] = dict()
+
+        def log_pre_api_call(
+            self,
+            model,
+            messages,
+            kwargs,
+        ):
+            litellm_call_id = kwargs.get("litellm_call_id")
+
+            if not litellm_call_id:
+                log.warning("Agenta SDK - litellm tracing failed")
+                return
+
+            type = (  # pylint: disable=redefined-builtin
+                "chat"
+                if kwargs.get("call_type") in ["completion", "acompletion"]
+                else "embedding"
+            )
+
+            kind = SpanKind.CLIENT
+
+            self.span[litellm_call_id] = CustomSpan(
+                ag.tracer.start_span(name=f"litellm_{kind.name.lower()}", kind=kind)
+            )
+
+            span = self.span[litellm_call_id]
+
+            if not span:
+                log.warning("Agenta SDK - litellm tracing failed")
+                return
+
+            if not span.is_recording():
+                log.error("Agenta SDK - litellm span not recording.")
+                return
+
+            span.set_attributes(
+                attributes={"node": type},
+                namespace="type",
+            )
+
+            span.set_attributes(
+                attributes={"inputs": {"prompt": kwargs["messages"]}},
+                namespace="data",
+            )
+
+            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,
+        ):
+            litellm_call_id = kwargs.get("litellm_call_id")
+
+            if not litellm_call_id:
+                log.warning("Agenta SDK - litellm tracing failed")
+                return
+
+            span = self.span[litellm_call_id]
+
+            if not span:
+                log.warning("Agenta SDK - litellm tracing failed")
+                return
+
+            if not span.is_recording():
+                return
+
+        def log_success_event(
+            self,
+            kwargs,
+            response_obj,
+            start_time,
+            end_time,
+        ):
+            if kwargs.get("stream"):
+                return
+
+            litellm_call_id = kwargs.get("litellm_call_id")
+
+            if not litellm_call_id:
+                log.warning("Agenta SDK - litellm tracing failed")
+                return
+
+            span = self.span[litellm_call_id]
+
+            if not span:
+                log.warning("Agenta SDK - litellm tracing failed")
+                return
+
+            if not span.is_recording():
+                return
+
+            try:
+                result = []
+                for choice in response_obj.choices:
+                    message = choice.message.__dict__
+                    result.append(message)
+
+                outputs = {"completion": result}
+                span.set_attributes(
+                    attributes={"outputs": outputs},
+                    namespace="data",
+                )
+
+            except Exception as e:
+                pass
+
+            span.set_attributes(
+                attributes={"total": kwargs.get("response_cost")},
+                namespace="metrics.unit.costs",
+            )
+
+            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",
+            )
+
+            span.set_status(status="OK")
+
+            span.end()
+
+        def log_failure_event(
+            self,
+            kwargs,
+            response_obj,
+            start_time,
+            end_time,
+        ):
+            litellm_call_id = kwargs.get("litellm_call_id")
+
+            if not litellm_call_id:
+                log.warning("Agenta SDK - litellm tracing failed")
+                return
+
+            span = self.span[litellm_call_id]
+
+            if not span:
+                log.warning("Agenta SDK - litellm tracing failed")
+                return
+
+            if not span.is_recording():
+                return
+
+            span.record_exception(kwargs["exception"])
+
+            span.set_status(status="ERROR")
+
+            span.end()
+
+        async def async_log_stream_event(
+            self,
+            kwargs,
+            response_obj,
+            start_time,
+            end_time,
+        ):
+            if kwargs.get("stream"):
+                return
+
+            litellm_call_id = kwargs.get("litellm_call_id")
+
+            if not litellm_call_id:
+                log.warning("Agenta SDK - litellm tracing failed")
+                return
+
+            span = self.span[litellm_call_id]
+
+            if not span:
+                log.warning("Agenta SDK - litellm tracing failed")
+                return
+
+            if not span.is_recording():
+                return
+
+        async def async_log_success_event(
+            self,
+            kwargs,
+            response_obj,
+            start_time,
+            end_time,
+        ):
+            litellm_call_id = kwargs.get("litellm_call_id")
+
+            if not litellm_call_id:
+                log.warning("Agenta SDK - litellm tracing failed")
+                return
+
+            span = self.span[litellm_call_id]
+
+            if not span:
+                log.warning("Agenta SDK - litellm tracing failed")
+                return
+
+            if not span.is_recording():
+                return
+
+            try:
+                result = []
+                for choice in response_obj.choices:
+                    message = choice.message.__dict__
+                    result.append(message)
+
+                outputs = {"completion": result}
+                span.set_attributes(
+                    attributes={"outputs": outputs},
+                    namespace="data",
+                )
+
+            except Exception as e:
+                pass
+
+            span.set_attributes(
+                attributes={"total": kwargs.get("response_cost")},
+                namespace="metrics.unit.costs",
+            )
+
+            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",
+            )
+
+            span.set_status(status="OK")
+
+            span.end()
+
+        async def async_log_failure_event(
+            self,
+            kwargs,
+            response_obj,
+            start_time,
+            end_time,
+        ):
+            litellm_call_id = kwargs.get("litellm_call_id")
+
+            if not litellm_call_id:
+                log.warning("Agenta SDK - litellm tracing failed")
+                return
+
+            span = self.span[litellm_call_id]
+
+            if not span:
+                log.warning("Agenta SDK - litellm tracing failed")
+                return
+
+            if not span.is_recording():
+                return
+
+            span.record_exception(kwargs["exception"])
+
+            span.set_status(status="ERROR")
+
+            span.end()
+
+    return LitellmHandler()

agenta/sdk/litellm/mockllm.py

@@ -0,0 +1,27 @@
+from typing import Optional, Protocol, Any
+
+from agenta.sdk.context.routing import routing_context
+from agenta.sdk.litellm.mocks import MOCKS
+
+
+class LitellmProtocol(Protocol):
+    async def acompletion(self, *args: Any, **kwargs: Any) -> Any:
+        ...
+
+
+litellm: Optional[LitellmProtocol] = None  # pylint: disable=invalid-name
+
+
+async def acompletion(*args, **kwargs):
+    mock = routing_context.get().mock
+
+    if mock:
+        if mock not in MOCKS:
+            raise ValueError(f"Mock {mock} not found")
+
+        return MOCKS[mock]
+
+    if not litellm:
+        raise ValueError("litellm not found")
+
+    return await litellm.acompletion(*args, **kwargs)

agenta/sdk/litellm/mocks/__init__.py

@@ -0,0 +1,26 @@
+from pydantic import BaseModel
+
+
+class MockMessageModel(BaseModel):
+    content: str
+
+
+class MockChoiceModel(BaseModel):
+    message: MockMessageModel
+
+
+class MockResponseModel(BaseModel):
+    choices: list[MockChoiceModel]
+
+
+MOCKS = {
+    "hello": MockResponseModel(
+        choices=[
+            MockChoiceModel(
+                message=MockMessageModel(
+                    content="world",
+                )
+            )
+        ],
+    ),
+}

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"]