agenta 0.25.4__py3-none-any.whl → 0.25.4a1__py3-none-any.whl

agenta/sdk/decorators/tracing.py

@@ -1,111 +1,196 @@
-# Stdlib Imports
 import inspect
 import traceback
 from functools import wraps
-from typing import Any, Callable, Optional, List, Union
+from itertools import chain
+from typing import Callable, Optional, Any, Dict, List
 
-# 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")
+from agenta.sdk.utils.exceptions import suppress
 
+from agenta.sdk.context.tracing import tracing_context
+from agenta.sdk.tracing.conventions import parse_span_kind
 
-class instrument(BaseDecorator):
-    """Decorator class for monitoring llm apps functions.
 
-    Args:
-        BaseDecorator (object): base decorator class
-
-    Example:
-    ```python
-        import agenta as ag
-
-        prompt_config = {"system_prompt": ..., "temperature": 0.5, "max_tokens": ...}
-
-        @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:
+    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,
+        kind: str = "task",
+        config: Optional[Dict[str, Any]] = None,
+        ignore_inputs: Optional[bool] = None,
+        ignore_outputs: Optional[bool] = None,
+        max_depth: Optional[int] = 2,
+        # DEPRECATED
+        spankind: Optional[str] = "TASK",
     ) -> None:
+        self.kind = spankind if spankind is not None else kind
         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)
+        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
 
-            return input_dict
+        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.
 
-        def redact(io, blacklist):
-            return {
-                key: io[key]
-                for key in io.keys()
+            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 (
-                    blacklist
-                    if isinstance(blacklist, list)
-                    else []
-                    if blacklist is False
+                    ignore
+                    if isinstance(ignore, list)
                     else io.keys()
+                    if ignore is True
+                    else []
                 )
             }
 
-        def patch(result):
-            TRACE_DEFAULT_KEY = "__default__"
-
-            outputs = result
+            return io
 
-            # 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"]}
+        def patch(result: Any) -> Dict[str, Any]:
+            """
+            Patch the result to ensure that it is a dictionary, with a default key when necessary.
 
-                    ag.tracing.store_cost(result["cost"])
-                    ag.tracing.store_usage(result["usage"])
+            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 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,
-                ):
-                    result = await 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,
+                        )
+
+                    try:
+                        result = await func(*args, **kwargs)
+                    except Exception as e:
+                        traceback.print_exc()
+
+                        span.record_exception(e)
+
+                        span.set_status("ERROR")
 
-                    ag.tracing.store_outputs(redact(patch(result), self.ignore_outputs))
+                        raise e
+
+                    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
 
@@ -114,15 +199,98 @@ class instrument(BaseDecorator):
         @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,
-                ):
-                    result = func(*args, **kwargs)
-
-                    ag.tracing.store_outputs(redact(patch(result), self.ignore_outputs))
+                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,
+                        )
+
+                    try:
+                        result = func(*args, **kwargs)
+                    except Exception as e:
+                        span.record_exception(e)
+
+                        span.set_status("ERROR")
+
+                        raise e
+
+                    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
 

agenta/sdk/litellm/__init__.py

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

agenta/sdk/litellm/litellm.py

@@ -0,0 +1,275 @@
+import agenta as ag
+
+from agenta.sdk.tracing.spans import CustomSpan
+from agenta.sdk.utils.exceptions import suppress
+from agenta.sdk.utils.logging import log
+
+
+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.
+        """
+
+        def __init__(self):
+            self.span = None
+
+        def log_pre_api_call(
+            self,
+            model,
+            messages,
+            kwargs,
+        ):
+            type = (
+                "chat"
+                if kwargs.get("call_type") in ["completion", "acompletion"]
+                else "embedding"
+            )
+
+            kind = "CLIENT"
+
+            self.span = CustomSpan(
+                ag.tracer.start_span(name=f"litellm_{kind.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
+
+            log.info(f"log_pre_api_call({hex(self.span.context.span_id)[2:]})")
+
+            self.span.set_attributes(
+                attributes={"inputs": {"messages": 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
+
+            # log.info(f"log_stream({hex(self.span.context.span_id)[2:]})")
+
+            self.span.set_attributes(
+                attributes={
+                    "output": {"__default__": kwargs.get("complete_streaming_response")}
+                },
+                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
+
+            # log.info(f"log_success({hex(self.span.context.span_id)[2:]})")
+
+            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")},
+                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
+
+            # log.info(f"log_failure({hex(self.span.context.span_id)[2:]})")
+
+            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
+
+            # log.info(f"async_log_stream({hex(self.span.context.span_id)[2:]})")
+
+            self.span.set_attributes(
+                attributes={
+                    "output": {"__default__": kwargs.get("complete_streaming_response")}
+                },
+                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
+
+            log.info(f"async_log_success({hex(self.span.context.span_id)[2:]})")
+
+            self.span.set_attributes(
+                attributes={
+                    "output": {"__default__": kwargs.get("complete_streaming_response")}
+                },
+                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
+
+            # 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")
+
+            self.span.end()
+
+    return LitellmHandler()

agenta/sdk/router.py

@@ -1,15 +1,8 @@
 from fastapi import APIRouter
-from .context import get_contexts
 
 router = APIRouter()
 
 
-@router.get("/contexts/")
-def get_all_contexts():
-    contexts = get_contexts()
-    return {"contexts": contexts}
-
-
 @router.get("/health")
 def health():
     return {"status": "ok"}

agenta/sdk/tracing/__init__.py

@@ -0,0 +1 @@
+from .tracing import Tracing, get_tracer