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

agenta/sdk/litellm/litellm.py

@@ -1,275 +0,0 @@
-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/tracing/attributes.py

@@ -1,181 +0,0 @@
-from json import loads, dumps
-from typing import Optional, Union, Sequence, Any, Dict
-
-Primitive = Union[str, int, float, bool, bytes]
-PrimitivesSequence = Sequence[Primitive]
-Attribute = Union[Primitive, PrimitivesSequence]
-
-
-def _marshal(
-    unmarshalled: Dict[str, Any],
-    *,
-    parent_key: Optional[str] = "",
-    depth: Optional[int] = 0,
-    max_depth: Optional[int] = None,
-) -> Dict[str, Any]:
-    """
-    Marshals a dictionary of unmarshalled attributes into a flat dictionary
-
-    Example:
-    unmarshalled = {
-        "ag": {
-            "type": "tree",
-            "node": {
-                "name": "root",
-                "children": [
-                    {
-                        "name": "child1",
-                    },
-                    {
-                        "name": "child2",
-                    }
-                ]
-            }
-        }
-    }
-    marshalled = {
-        "ag.type": "tree",
-        "ag.node.name": "root",
-        "ag.node.children.0.name": "child1",
-        "ag.node.children.1.name": "child2"
-    }
-    """
-    marshalled = {}
-
-    # If max_depth is set and we've reached it,
-    # just return the unmarshalled attributes
-    if max_depth is not None and depth >= max_depth:
-        marshalled[parent_key] = unmarshalled
-        # MISSING ENCODING TO JSON IF NOT PRIMITIVE
-
-        return marshalled
-
-    # Otherwise,
-    # iterate over the unmarshalled attributes and marshall them
-    for key, value in unmarshalled.items():
-        child_key = f"{parent_key}.{key}" if parent_key else key
-
-        if isinstance(value, dict):
-            dict_key = child_key
-
-            marshalled.update(
-                _marshal(
-                    value,
-                    parent_key=dict_key,
-                    depth=depth + 1,
-                    max_depth=max_depth,
-                )
-            )
-        elif isinstance(value, list):
-            if max_depth is not None and depth + 1 >= max_depth:
-                marshalled[child_key] = value
-                # MISSING ENCODING TO JSON IF NOT PRIMITIVE
-            else:
-                for i, item in enumerate(value):
-                    list_key = f"{child_key}.{i}"
-
-                    if isinstance(item, (dict, list)):
-                        marshalled.update(
-                            _marshal(
-                                item,
-                                parent_key=list_key,
-                                depth=depth + 1,
-                                max_depth=max_depth,
-                            )
-                        )
-                    else:
-                        marshalled[list_key] = item
-                        # MISSING ENCODING TO JSON IF NOT PRIMITIVE
-        else:
-            marshalled[child_key] = value
-            # MISSING ENCODING TO JSON IF NOT PRIMITIVE
-
-    return marshalled
-
-
-def _encode_key(namespace: Optional[str] = None, key: str = "") -> str:
-    if namespace is None:
-        return key
-
-    return f"ag.{namespace}.{key}"
-
-
-def _encode_value(value: Any) -> Optional[Attribute]:
-    if value is None:
-        return None
-
-    if isinstance(value, (str, int, float, bool, bytes)):
-        return value
-
-    if isinstance(value, dict) or isinstance(value, list):
-        encoded = dumps(value)
-        value = "@ag.type=json:" + encoded
-        return value
-
-    return repr(value)
-
-
-def serialize(
-    *,
-    namespace: str,
-    attributes: Dict[str, Any],
-    max_depth: Optional[int] = None,
-) -> Dict[str, str]:
-    if not isinstance(attributes, dict):
-        return {}
-
-    _attributes = {
-        k: v
-        for k, v in {
-            _encode_key(namespace, key): _encode_value(value)
-            for key, value in _marshal(attributes, max_depth=max_depth).items()
-        }.items()
-        if v is not None
-    }
-
-    return _attributes
-
-
-"""
-def _decode_key(namespace: Optional[str] = None, key: str = "") -> str:
-    if namespace is None:
-        return key
-
-    return key.replace(f"ag.{namespace}.", "")
-"""
-
-"""
-def _decode_value(value: Attribute) -> Any:
-    if isinstance(value, (int, float, bool, bytes)):
-        return value
-
-    if isinstance(value, str):
-        if value == "@ag.type=none:":
-            return None
-
-        if value.startswith("@ag.type=json:"):
-            encoded = value[len("@ag.type=json:") :]
-            value = loads(encoded)
-            return value
-
-        return value
-
-    return value
-"""
-
-"""
-def deserialize(
-    *,
-    namespace: str,
-    attributes: Dict[str, Any],
-    max_depth: Optional[int] = None,
-) -> Dict[str, Any]:
-    if not isinstance(attributes, dict):
-        return {}
-
-    return {
-        _decode_key(namespace, key): _decode_value(value)
-        for key, value in attributes
-        if key != _decode_key(namespace, key)
-    }
-"""

agenta/sdk/tracing/context.py

@@ -1,21 +0,0 @@
-from contextvars import ContextVar
-from contextlib import contextmanager
-from traceback import format_exc
-
-from agenta.sdk.utils.logging import log
-
-tracing_context = ContextVar("tracing_context", default={})
-
-
-@contextmanager
-def tracing_context_manager():
-    _tracing_context = {"health": {"status": "ok"}}
-
-    token = tracing_context.set(_tracing_context)
-    try:
-        yield
-    except Exception as e:
-        log.error(f"Error with tracing context: {_tracing_context}")
-        log.error(f"Exception: {format_exc()}")
-    finally:
-        tracing_context.reset(token)

agenta/sdk/tracing/conventions.py

@@ -1,43 +0,0 @@
-from typing import Literal
-
-from opentelemetry.trace import SpanKind
-
-Namespace = Literal[
-    "data.inputs",
-    "data.internals",
-    "data.outputs",
-    "metrics.scores",
-    "metrics.unit.costs",
-    "metrics.unit.tokens",
-    "meta.configuration",
-    "meta.version",
-    "tags",
-    "refs",
-]
-
-Code = Literal[
-    "OK",
-    "UNSET",
-    "ERROR",
-]
-
-
-def parse_span_kind(type: str) -> SpanKind:
-    kind = SpanKind.INTERNAL
-    if type in [
-        "agent",
-        "chain",
-        "workflow",
-    ]:
-        kind = SpanKind.SERVER
-    elif type in [
-        "tool",
-        "embedding",
-        "query",
-        "completion",
-        "chat",
-        "rerank",
-    ]:
-        kind = SpanKind.CLIENT
-
-    return kind

agenta/sdk/tracing/exporters.py

@@ -1,53 +0,0 @@
-from typing import Sequence, Dict, List
-
-from opentelemetry.sdk.trace.export import (
-    ConsoleSpanExporter,
-    SpanExporter,
-    SpanExportResult,
-    ReadableSpan,
-)
-from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
-    OTLPSpanExporter,
-)
-
-from agenta.sdk.utils.exceptions import suppress
-
-
-class InlineTraceExporter(SpanExporter):
-    def __init__(self, registry: Dict[str, List[ReadableSpan]]):
-        self._shutdown = False
-        self._registry = registry
-
-    def export(
-        self,
-        spans: Sequence[ReadableSpan],
-    ) -> SpanExportResult:
-        if self._shutdown:
-            return
-
-        with suppress():
-            for span in spans:
-                trace_id = span.get_span_context().trace_id
-
-                if trace_id not in self._registry:
-                    self._registry[trace_id] = []
-
-                self._registry[trace_id].append(span)
-
-    def shutdown(self) -> None:
-        self._shutdown = True
-
-    def force_flush(self, timeout_millis: int = 30000) -> bool:
-        return True
-
-    def fetch(self, trace_id: int) -> List[ReadableSpan]:
-        trace = self._registry.get(trace_id, [])
-
-        del self._registry[trace_id]
-
-        return trace
-
-
-ConsoleExporter = ConsoleSpanExporter
-InlineExporter = InlineTraceExporter
-OTLPExporter = OTLPSpanExporter