agenta 0.25.4a3__py3-none-any.whl → 0.25.4a4__py3-none-any.whl

agenta/sdk/decorators/tracing.py

@@ -1,196 +1,111 @@
+# Stdlib Imports
 import inspect
 import traceback
 from functools import wraps
-from itertools import chain
-from typing import Callable, Optional, Any, Dict, List
+from typing import Any, Callable, Optional, List, Union
 
+# 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
 
 
-from agenta.sdk.utils.exceptions import suppress
+logging.setLevel("DEBUG")
 
-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.
 
-class instrument:
-    DEFAULT_KEY = "__default__"
+    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 ...
+    ```
+    """
 
     def __init__(
         self,
-        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",
+        config: Optional[dict] = None,
+        spankind: str = "workflow",
+        ignore_inputs: Union[List[str], bool] = False,
+        ignore_outputs: Union[List[str], bool] = False,
     ) -> 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 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 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 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.
+            return input_dict
 
-            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()
+        def redact(io, blacklist):
+            return {
+                key: io[key]
+                for key in io.keys()
                 if key
                 not in (
-                    ignore
-                    if isinstance(ignore, list)
-                    else io.keys()
-                    if ignore is True
+                    blacklist
+                    if isinstance(blacklist, list)
                     else []
+                    if blacklist is False
+                    else io.keys()
                 )
             }
 
-            return io
+        def patch(result):
+            TRACE_DEFAULT_KEY = "__default__"
 
-        def patch(result: Any) -> Dict[str, Any]:
-            """
-            Patch the result to ensure that it is a dictionary, with a default key when necessary.
+            outputs = result
 
-            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
-                )
-            )
+            # 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
 
         @wraps(func)
         async def async_wrapper(*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,
-                        )
-
-                    try:
-                        result = await func(*args, **kwargs)
-                    except Exception as e:
-                        traceback.print_exc()
-
-                        span.record_exception(e)
-
-                        span.set_status("ERROR")
+                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)
 
-                        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,
-                                    }
-                                }
-                            )
+                    ag.tracing.store_outputs(redact(patch(result), self.ignore_outputs))
 
                     return result
 
@@ -199,98 +114,15 @@ class instrument:
         @wraps(func)
         def sync_wrapper(*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,
-                        )
-
-                    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,
-                                    }
-                                }
-                            )
+                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))
 
                     return result
 

agenta/sdk/router.py

@@ -1,8 +1,15 @@
 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

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

agenta/sdk/tracing/callbacks.py

@@ -0,0 +1,187 @@
+import agenta as ag
+
+from agenta.sdk.tracing.tracing_context import tracing_context, TracingContext
+from agenta.sdk.tracing.logger import llm_logger as logging
+
+from agenta.sdk.utils.debug import debug
+
+TRACE_DEFAULT_KEY = "__default__"
+
+
+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
+
+        @property
+        def _trace(self):
+            return ag.tracing
+
+        @debug()
+        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.span = ag.tracing.open_span(
+                name=f"{span_kind}_call",
+                input={"messages": kwargs["messages"]},
+                spankind=span_kind,
+                active=False,
+            )
+            logging.info(f"log_pre_api_call({self.span.id})")
+            ag.tracing.set_attributes(
+                {
+                    "model_config": {
+                        "model": kwargs.get("model"),
+                        **kwargs.get(
+                            "optional_params"
+                        ),  # model-specific params passed in
+                    },
+                }
+            )
+
+        @debug()
+        def log_stream_event(self, kwargs, response_obj, start_time, end_time):
+            logging.info(f"log_stream_event({self.span.id})")
+            ag.tracing.set_status(status="OK", span_id=self.span.id)
+            ag.tracing.store_cost(kwargs.get("response_cost"), span_id=self.span.id)
+            ag.tracing.store_usage(
+                response_obj.usage.dict() if hasattr(response_obj, "usage") else None,
+                span_id=self.span.id,
+            )
+            ag.tracing.store_outputs(
+                # the complete streamed response (only set if `completion(..stream=True)`
+                outputs={TRACE_DEFAULT_KEY: kwargs.get("complete_streaming_response")},
+                span_id=self.span.id,
+            )
+            ag.tracing.close_span(span_id=self.span.id)
+
+        @debug()
+        def log_success_event(self, kwargs, response_obj, start_time, end_time):
+            logging.info(f"log_success_event({self.span.id})")
+            ag.tracing.set_status(status="OK", span_id=self.span.id)
+            ag.tracing.store_cost(kwargs.get("response_cost"), span_id=self.span.id)
+            ag.tracing.store_usage(
+                response_obj.usage.dict() if hasattr(response_obj, "usage") else None,
+                span_id=self.span.id,
+            )
+            ag.tracing.store_outputs(
+                outputs={TRACE_DEFAULT_KEY: response_obj.choices[0].message.content},
+                span_id=self.span.id,
+            )
+            ag.tracing.close_span(span_id=self.span.id)
+
+        @debug()
+        def log_failure_event(self, kwargs, response_obj, start_time, end_time):
+            logging.info("log_failure_event()", self.span)
+            ag.tracing.set_status(status="ERROR", span_id=self.span.id)
+            ag.tracing.set_attributes(
+                {
+                    "traceback_exception": repr(
+                        kwargs["traceback_exception"]
+                    ),  # the traceback generated via `traceback.format_exc()`
+                    "call_end_time": kwargs[
+                        "end_time"
+                    ],  # datetime object of when call was completed
+                },
+                span_id=self.span.id,
+            )
+            ag.tracing.store_cost(kwargs.get("response_cost"), span_id=self.span.id)
+            ag.tracing.store_usage(
+                response_obj.usage.dict() if hasattr(response_obj, "usage") else None,
+                span_id=self.span.id,
+            )
+            ag.tracing.store_outputs(
+                # the Exception raised
+                outputs={TRACE_DEFAULT_KEY: repr(kwargs["exception"])},
+                span_id=self.span.id,
+            )
+            ag.tracing.close_span(span_id=self.span.id)
+
+        @debug()
+        async def async_log_stream_event(
+            self, kwargs, response_obj, start_time, end_time
+        ):
+            logging.info(f"async_log_stream_event({self.span.id})")
+            ag.tracing.set_status(status="OK", span_id=self.span.id)
+            ag.tracing.store_cost(kwargs.get("response_cost"), span_id=self.span.id)
+            ag.tracing.store_usage(
+                response_obj.usage.dict() if hasattr(response_obj, "usage") else None,
+                span_id=self.span.id,
+            )
+            ag.tracing.store_outputs(
+                # the complete streamed response (only set if `completion(..stream=True)`)
+                outputs={TRACE_DEFAULT_KEY: kwargs.get("complete_streaming_response")},
+                span_id=self.span.id,
+            )
+            ag.tracing.close_span(span_id=self.span.id)
+
+        @debug()
+        async def async_log_success_event(
+            self, kwargs, response_obj, start_time, end_time
+        ):
+            logging.info(f"async_log_success_event({self.span.id})")
+            ag.tracing.set_status(status="OK", span_id=self.span.id)
+            ag.tracing.store_cost(kwargs.get("response_cost"), span_id=self.span.id)
+            ag.tracing.store_usage(
+                response_obj.usage.dict() if hasattr(response_obj, "usage") else None,
+                span_id=self.span.id,
+            )
+            ag.tracing.store_outputs(
+                outputs={TRACE_DEFAULT_KEY: response_obj.choices[0].message.content},
+                span_id=self.span.id,
+            )
+            ag.tracing.close_span(span_id=self.span.id)
+
+        @debug()
+        async def async_log_failure_event(
+            self, kwargs, response_obj, start_time, end_time
+        ):
+            logging.info(f"async_log_failure_event({self.span.id})")
+            ag.tracing.set_status(status="ERROR", span_id=self.span.id)
+            ag.tracing.set_attributes(
+                {
+                    "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
+                },
+                span_id=self.span.id,
+            )
+            ag.tracing.store_cost(kwargs.get("response_cost"), span_id=self.span.id)
+            ag.tracing.store_usage(
+                response_obj.usage.dict() if hasattr(response_obj, "usage") else None,
+                span_id=self.span.id,
+            )
+            ag.tracing.store_outputs(
+                # the Exception raised
+                outputs={TRACE_DEFAULT_KEY: repr(kwargs["exception"])},
+                span_id=self.span.id,
+            )
+            ag.tracing.close_span(span_id=self.span.id)
+
+    return LitellmHandler()