mseep-agentops 0.4.18__py3-none-any.whl

agentops/instrumentation/common/wrappers.py

@@ -0,0 +1,235 @@
+"""Common wrapper utilities for OpenTelemetry instrumentation.
+
+This module provides common utilities for creating and managing wrappers
+around functions and methods for OpenTelemetry instrumentation. It includes
+a configuration class for wrapping methods, helper functions for updating
+spans with attributes, and functions for creating and applying wrappers.
+"""
+
+from typing import Any, Optional, Tuple, Dict, Callable
+from dataclasses import dataclass
+import logging
+from wrapt import wrap_function_wrapper  # type: ignore
+from opentelemetry.instrumentation.utils import unwrap as _unwrap
+from opentelemetry.trace import Tracer
+from opentelemetry.trace import Span, SpanKind, Status, StatusCode
+from opentelemetry import context as context_api
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+
+from agentops.instrumentation.common.attributes import AttributeMap
+
+logger = logging.getLogger(__name__)
+
+AttributeHandler = Callable[[Optional[Tuple], Optional[Dict], Optional[Any]], AttributeMap]
+
+
+@dataclass
+class WrapConfig:
+    """Configuration for wrapping a method with OpenTelemetry instrumentation.
+
+    This class defines how a method should be wrapped for instrumentation,
+    including what package, class, and method to wrap, what span attributes
+    to set, and how to name the resulting trace spans.
+
+    Attributes:
+        trace_name: The name to use for the trace span
+        package: The package containing the target class
+        class_name: The name of the class containing the method
+        method_name: The name of the method to wrap
+        handler: A function that extracts attributes from args, kwargs, or return value
+        is_async: Whether the method is asynchronous (default: False)
+            We explicitly specify async methods since `asyncio.iscoroutinefunction`
+            is not reliable in this context.
+        span_kind: The kind of span to create (default: CLIENT)
+    """
+
+    trace_name: str
+    package: str
+    class_name: str
+    method_name: str
+    handler: AttributeHandler
+    is_async: bool = False
+    span_kind: SpanKind = SpanKind.CLIENT
+
+    def __repr__(self):
+        return f"{self.package}.{self.class_name}.{self.method_name}"
+
+
+def _update_span(span: Span, attributes: AttributeMap) -> None:
+    """Update a span with the provided attributes.
+
+    Args:
+        span: The OpenTelemetry span to update
+        attributes: A dictionary of attributes to set on the span
+    """
+    for key, value in attributes.items():
+        span.set_attribute(key, value)
+
+
+def _finish_span_success(span: Span) -> None:
+    """Mark a span as successful by setting its status to OK.
+
+    Args:
+        span: The OpenTelemetry span to update
+    """
+    span.set_status(Status(StatusCode.OK))
+
+
+def _finish_span_error(span: Span, exception: Exception) -> None:
+    """Mark a span as failed by recording the exception and setting error status.
+
+    Args:
+        span: The OpenTelemetry span to update
+        exception: The exception that caused the error
+    """
+    span.record_exception(exception)
+    span.set_status(Status(StatusCode.ERROR, str(exception)))
+
+
+def _create_wrapper(wrap_config: WrapConfig, tracer: Tracer) -> Callable:
+    """Create a wrapper function for the specified configuration.
+
+    This function creates a wrapper that:
+    1. Creates a new span for the wrapped method
+    2. Sets attributes on the span based on input arguments
+    3. Calls the wrapped method
+    4. Sets attributes on the span based on the return value
+    5. Handles exceptions by recording them on the span
+
+    Args:
+        wrap_config: Configuration for the wrapper
+        tracer: The OpenTelemetry tracer to use for creating spans
+
+    Returns:
+        A wrapper function compatible with wrapt.wrap_function_wrapper
+    """
+    handler = wrap_config.handler
+
+    async def awrapper(wrapped, instance, args, kwargs):
+        # Skip instrumentation if it's suppressed in the current context
+        # TODO I don't understand what this actually does
+        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+            return wrapped(*args, **kwargs)
+
+        return_value = None
+
+        with tracer.start_as_current_span(
+            wrap_config.trace_name,
+            kind=wrap_config.span_kind,
+        ) as span:
+            try:
+                # Add the input attributes to the span before execution
+                attributes = handler(args=args, kwargs=kwargs)
+                _update_span(span, attributes)
+
+                return_value = await wrapped(*args, **kwargs)
+
+                # Add the output attributes to the span after execution
+                attributes = handler(return_value=return_value)
+                _update_span(span, attributes)
+                _finish_span_success(span)
+            except Exception as e:
+                # Add everything we have in the case of an error
+                attributes = handler(args=args, kwargs=kwargs, return_value=return_value)
+                _update_span(span, attributes)
+                _finish_span_error(span, e)
+                raise
+
+        return return_value
+
+    def wrapper(wrapped, instance, args, kwargs):
+        # Skip instrumentation if it's suppressed in the current context
+        # TODO I don't understand what this actually does
+        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+            return wrapped(*args, **kwargs)
+
+        return_value = None
+
+        with tracer.start_as_current_span(
+            wrap_config.trace_name,
+            kind=wrap_config.span_kind,
+        ) as span:
+            try:
+                # Add the input attributes to the span before execution
+                attributes = handler(args=args, kwargs=kwargs)
+                _update_span(span, attributes)
+
+                return_value = wrapped(*args, **kwargs)
+
+                # Add the output attributes to the span after execution
+                attributes = handler(return_value=return_value)
+                _update_span(span, attributes)
+                _finish_span_success(span)
+            except Exception as e:
+                # Add everything we have in the case of an error
+                attributes = handler(args=args, kwargs=kwargs, return_value=return_value)
+                _update_span(span, attributes)
+                _finish_span_error(span, e)
+                raise
+
+        return return_value
+
+    if wrap_config.is_async:
+        return awrapper
+    else:
+        return wrapper
+
+
+def wrap(wrap_config: WrapConfig, tracer: Tracer) -> Callable:
+    """Wrap a method with OpenTelemetry instrumentation.
+
+    This function applies the wrapper created by _create_wrapper
+    to the method specified in the wrap_config.
+
+    Args:
+        wrap_config: Configuration specifying what to wrap and how
+        tracer: The OpenTelemetry tracer to use for creating spans
+
+    Returns:
+        The result of wrap_function_wrapper (typically None)
+    """
+    return wrap_function_wrapper(
+        wrap_config.package,
+        f"{wrap_config.class_name}.{wrap_config.method_name}",
+        _create_wrapper(wrap_config, tracer),
+    )
+
+
+def unwrap(wrap_config: WrapConfig):
+    """Remove instrumentation wrapper from a method.
+
+    This function removes the wrapper applied by wrap().
+
+    Args:
+        wrap_config: Configuration specifying what to unwrap
+
+    Returns:
+        The result of the unwrap operation (typically None)
+    """
+    return _unwrap(
+        f"{wrap_config.package}.{wrap_config.class_name}",
+        wrap_config.method_name,
+    )
+
+
+def _with_tracer_wrapper(func):
+    """Wrap a function with a tracer.
+
+    This decorator creates a higher-order function that takes a tracer as its first argument
+    and returns a function suitable for use with wrapt's wrap_function_wrapper. It's used
+    to consistently apply OpenTelemetry tracing to SDK functions.
+
+    Args:
+        func: The instrumentation function to wrap
+
+    Returns:
+        A decorator function that takes a tracer and returns a wrapt-compatible wrapper
+    """
+
+    def _with_tracer(tracer):
+        def wrapper(wrapped, instance, args, kwargs):
+            return func(tracer, wrapped, instance, args, kwargs)
+
+        return wrapper
+
+    return _with_tracer

agentops/legacy/__init__.py

@@ -0,0 +1,277 @@
+"""
+Compatibility layer for deprecated functions and classes.
+
+CrewAI contains direct integrations with AgentOps across multiple versions.
+These integrations use different patterns:
+- CrewAI < 0.105.0: Direct calls to agentops.end_session() with kwargs
+- CrewAI >= 0.105.0: Event-based integration using Session objects
+
+This module maintains backward compatibility with all these API patterns.
+"""
+
+from typing import Optional, Any, Dict, List, Union
+
+from agentops.logging import logger
+from agentops.sdk.core import TraceContext, tracer
+from agentops.helpers.deprecation import deprecated
+
+_current_session: Optional["Session"] = None
+_current_trace_context: Optional[TraceContext] = None
+
+
+class Session:
+    """
+    This class provides compatibility with CrewAI >= 0.105.0, which uses an event-based
+    integration pattern where it calls methods directly on the Session object:
+
+    - create_agent(): Called when a CrewAI agent is created
+    - record(): Called when a CrewAI tool is used
+    - end_session(): Called when a CrewAI run completes
+    """
+
+    def __init__(self, trace_context: Optional[TraceContext]):
+        self.trace_context = trace_context
+
+    @property
+    def span(self) -> Optional[Any]:
+        return self.trace_context.span if self.trace_context else None
+
+    @property
+    def token(self) -> Optional[Any]:
+        return self.trace_context.token if self.trace_context else None
+
+    def __del__(self):
+        if self.trace_context and self.trace_context.span and self.trace_context.span.is_recording():
+            if not self.trace_context.is_init_trace:
+                logger.warning(
+                    f"Legacy Session (trace ID: {self.trace_context.span.get_span_context().span_id}) \
+was garbage collected but its trace might still be recording. Ensure legacy sessions are ended with end_session()."
+                )
+
+    def create_agent(self, name: Optional[str] = None, agent_id: Optional[str] = None, **kwargs: Any):
+        """Method for CrewAI >= 0.105.0 compatibility. Currently a no-op."""
+        pass
+
+    def record(self, event: Any = None):
+        """Method for CrewAI >= 0.105.0 compatibility. Currently a no-op."""
+        pass
+
+    def end_session(self, **kwargs: Any):
+        """Ends the session for CrewAI >= 0.105.0 compatibility. Calls the global legacy end_session."""
+        end_session(session_or_status=self, **kwargs)
+
+
+@deprecated("Use agentops.start_trace() instead.")
+def start_session(
+    tags: Union[Dict[str, Any], List[str], None] = None,
+) -> Session:
+    """
+    @deprecated Use agentops.start_trace() instead.
+    Starts a legacy AgentOps session. Calls tracer.start_trace internally.
+    """
+    global _current_session, _current_trace_context
+
+    if not tracer.initialized:
+        from agentops import Client
+
+        try:
+            Client().init(auto_start_session=False)
+            if not tracer.initialized:
+                logger.warning("AgentOps client init failed during legacy start_session. Creating dummy session.")
+                dummy_session = Session(None)
+                _current_session = dummy_session
+                _current_trace_context = None
+                return dummy_session
+        except Exception as e:
+            logger.warning(f"AgentOps client init failed: {str(e)}. Creating dummy session.")
+            dummy_session = Session(None)
+            _current_session = dummy_session
+            _current_trace_context = None
+            return dummy_session
+
+    trace_context = tracer.start_trace(trace_name="session", tags=tags)
+    if trace_context is None:
+        logger.error("Failed to start trace via global tracer. Returning dummy session.")
+        dummy_session = Session(None)
+        _current_session = dummy_session
+        _current_trace_context = None
+        return dummy_session
+
+    session_obj = Session(trace_context)
+    _current_session = session_obj
+    _current_trace_context = trace_context
+
+    try:
+        import agentops.client.client
+
+        agentops.client.client._active_session = session_obj  # type: ignore
+        if hasattr(agentops.client.client, "_active_trace_context"):
+            agentops.client.client._active_trace_context = trace_context  # type: ignore
+    except (ImportError, AttributeError):
+        pass
+    return session_obj
+
+
+def _set_span_attributes(span: Any, attributes: Dict[str, Any]) -> None:
+    """Helper to set attributes on a span for legacy purposes."""
+    if span is None or not attributes:
+        return
+    for key, value in attributes.items():
+        if key.lower() == "end_state" and "end_state" in attributes:
+            pass
+        else:
+            span.set_attribute(f"agentops.legacy.{key}", str(value))
+
+
+@deprecated("Use agentops.end_trace() instead.")
+def end_session(session_or_status: Any = None, **kwargs: Any) -> None:
+    """
+    @deprecated Use agentops.end_trace() instead.
+    Ends a legacy AgentOps session. Calls tracer.end_trace internally.
+    Supports multiple calling patterns for backward compatibility.
+    """
+    global _current_session, _current_trace_context
+
+    if not tracer.initialized:
+        logger.debug("Ignoring end_session: global tracer not initialized.")
+        return
+
+    target_trace_context: Optional[TraceContext] = None
+    end_state_from_args = "Success"
+    extra_attributes = kwargs.copy()
+
+    if isinstance(session_or_status, Session):
+        target_trace_context = session_or_status.trace_context
+        if "end_state" in extra_attributes:
+            end_state_from_args = str(extra_attributes.pop("end_state"))
+    elif isinstance(session_or_status, str):
+        end_state_from_args = session_or_status
+        target_trace_context = _current_trace_context
+        if "end_state" in extra_attributes:
+            end_state_from_args = str(extra_attributes.pop("end_state"))
+    elif session_or_status is None and kwargs:
+        target_trace_context = _current_trace_context
+        if "end_state" in extra_attributes:
+            end_state_from_args = str(extra_attributes.pop("end_state"))
+    else:
+        target_trace_context = _current_trace_context
+        if "end_state" in extra_attributes:
+            end_state_from_args = str(extra_attributes.pop("end_state"))
+
+    if not target_trace_context:
+        logger.warning("end_session called but no active trace context found.")
+        return
+
+    if target_trace_context.span and extra_attributes:
+        _set_span_attributes(target_trace_context.span, extra_attributes)
+
+    tracer.end_trace(target_trace_context, end_state=end_state_from_args)
+
+    if target_trace_context is _current_trace_context:
+        _current_session = None
+        _current_trace_context = None
+
+    try:
+        import agentops.client.client
+
+        if (
+            hasattr(agentops.client.client, "_active_trace_context")
+            and agentops.client.client._active_trace_context is target_trace_context
+        ):  # type: ignore
+            agentops.client.client._active_trace_context = None  # type: ignore
+            agentops.client.client._active_session = None  # type: ignore
+        elif (
+            hasattr(agentops.client.client, "_init_trace_context")
+            and agentops.client.client._init_trace_context is target_trace_context
+        ):  # type: ignore
+            logger.debug("Legacy end_session called on client's auto-init trace. This is unusual.")
+    except (ImportError, AttributeError):
+        pass
+
+
+@deprecated("Use agentops.end_trace() instead.")
+def end_all_sessions() -> None:
+    """@deprecated Ends all active sessions/traces."""
+    if not tracer.initialized:
+        logger.debug("Ignoring end_all_sessions: global tracer not initialized.")
+        return
+
+    # Use the new end_trace functionality to end all active traces
+    tracer.end_trace(trace_context=None, end_state="Success")
+
+    # Clear legacy global state
+    global _current_session, _current_trace_context
+    _current_session = None
+    _current_trace_context = None
+
+
+@deprecated("Automatically tracked in v4.")
+def ToolEvent(*args: Any, **kwargs: Any) -> None:
+    """@deprecated Automatically tracked in v4."""
+    return None
+
+
+@deprecated("Automatically tracked in v4.")
+def ErrorEvent(*args: Any, **kwargs: Any) -> Any:
+    """@deprecated Automatically tracked in v4. Returns minimal object for test compatibility."""
+    from agentops.helpers.time import get_ISO_time
+
+    class LegacyErrorEvent:
+        def __init__(self):
+            self.init_timestamp = get_ISO_time()
+            self.end_timestamp = None
+
+    return LegacyErrorEvent()
+
+
+@deprecated("Automatically tracked in v4.")
+def ActionEvent(*args: Any, **kwargs: Any) -> Any:
+    """@deprecated Automatically tracked in v4. Returns minimal object for test compatibility."""
+    from agentops.helpers.time import get_ISO_time
+
+    class LegacyActionEvent:
+        def __init__(self):
+            self.init_timestamp = get_ISO_time()
+            self.end_timestamp = None
+
+    return LegacyActionEvent()
+
+
+@deprecated("Automatically tracked in v4.")
+def LLMEvent(*args: Any, **kwargs: Any) -> None:
+    """@deprecated Automatically tracked in v4."""
+    return None
+
+
+@deprecated("Use @agent decorator instead.")
+def track_agent(*args: Any, **kwargs: Any) -> Any:
+    """@deprecated No-op decorator."""
+
+    def noop(f: Any) -> Any:
+        return f
+
+    return noop
+
+
+@deprecated("Use @tool decorator instead.")
+def track_tool(*args: Any, **kwargs: Any) -> Any:
+    """@deprecated No-op decorator."""
+
+    def noop(f: Any) -> Any:
+        return f
+
+    return noop
+
+
+__all__ = [
+    "start_session",
+    "end_session",
+    "ToolEvent",
+    "ErrorEvent",
+    "ActionEvent",
+    "track_agent",
+    "track_tool",
+    "end_all_sessions",
+    "Session",
+    "LLMEvent",
+]

agentops/legacy/event.py

@@ -0,0 +1,156 @@
+"""
+AgentOps events.
+
+Data Class:
+    Event: Represents discrete events to be recorded.
+"""
+
+import traceback
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List, Optional, Sequence, Union
+from uuid import UUID, uuid4
+
+from agentops.helpers import get_ISO_time
+
+
+class EventType(Enum):
+    LLM = "llms"
+    ACTION = "actions"
+    API = "apis"
+    TOOL = "tools"
+    ERROR = "errors"
+
+
+@dataclass
+class Event:
+    """
+    Abstract base class for events that will be recorded. Should not be instantiated directly.
+
+    event_type(str): The type of event. Defined in events.EventType. Some values are 'llm', 'action', 'api', 'tool', 'error'.
+    params(dict, optional): The parameters of the function containing the triggered event, e.g. {'x': 1} in example below
+    returns(str, optional): The return value of the function containing the triggered event, e.g. 2 in example below
+    init_timestamp(str): A timestamp indicating when the event began. Defaults to the time when this Event was instantiated.
+    end_timestamp(str): A timestamp indicating when the event ended. Defaults to the time when this Event was instantiated.
+    agent_id(UUID, optional): The unique identifier of the agent that triggered the event.
+    id(UUID): A unique identifier for the event. Defaults to a new UUID.
+    session_id(UUID, optional): The unique identifier of the session that the event belongs to.
+
+    foo(x=1) {
+        ...
+        // params equals {'x': 1}
+        record(ActionEvent(params=**kwargs, ...))
+        ...
+        // returns equals 2
+        return x+1
+    }
+    """
+
+    event_type: str
+    params: Optional[dict] = None
+    returns: Optional[Union[str, List[str]]] = None
+    init_timestamp: str = field(default_factory=get_ISO_time)
+    end_timestamp: Optional[str] = None
+    agent_id: Optional[UUID] = None
+    id: UUID = field(default_factory=uuid4)
+    session_id: Optional[UUID] = None
+
+
+@dataclass
+class ActionEvent(Event):
+    """
+    For generic events
+
+    action_type(str, optional): High level name describing the action
+    logs(str, optional): For detailed information/logging related to the action
+    screenshot(str, optional): url to snapshot if agent interacts with UI
+    """
+
+    event_type: str = EventType.ACTION.value
+    # TODO: Should not be optional, but non-default argument 'agent_id' follows default argument error
+    action_type: Optional[str] = None
+    logs: Optional[Union[str, Sequence[Any]]] = None
+    screenshot: Optional[str] = None
+
+
+@dataclass
+class LLMEvent(Event):
+    """
+    For recording calls to LLMs. AgentOps auto-instruments calls to the most popular LLMs e.g. GPT, Claude, Gemini, etc.
+
+    thread_id(UUID, optional): The unique identifier of the contextual thread that a message pertains to.
+    prompt(str, list, optional): The message or messages that were used to prompt the LLM. Preferably in ChatML format which is more fully supported by AgentOps.
+    prompt_tokens(int, optional): The number of tokens in the prompt message.
+    completion(str, object, optional): The message or messages returned by the LLM. Preferably in ChatML format which is more fully supported by AgentOps.
+    completion_tokens(int, optional): The number of tokens in the completion message.
+    model(str, optional): LLM model e.g. "gpt-4", "gpt-3.5-turbo".
+
+    """
+
+    event_type: str = EventType.LLM.value
+    thread_id: Optional[UUID] = None
+    prompt: Optional[Union[str, List]] = None
+    prompt_tokens: Optional[int] = None
+    completion: Union[str, object] = None
+    completion_tokens: Optional[int] = None
+    cost: Optional[float] = None
+    model: Optional[str] = None
+
+
+@dataclass
+class ToolEvent(Event):
+    """
+    For recording calls to tools e.g. searchWeb, fetchFromDB
+
+    name(str, optional): A name describing the tool or the actual function name if applicable e.g. searchWeb, fetchFromDB.
+    logs(str, dict, optional): For detailed information/logging related to the tool.
+
+    """
+
+    event_type: str = EventType.TOOL.value
+    name: Optional[str] = None
+    logs: Optional[Union[str, dict]] = None
+
+
+# Does not inherit from Event because error will (optionally) be linked to an ActionEvent, LLMEvent, etc that will have the details
+
+
+@dataclass
+class ErrorEvent(Event):
+    """
+    For recording any errors e.g. ones related to agent execution
+
+    trigger_event(Event, optional): The event object that triggered the error if applicable.
+    exception(BaseException, optional): The thrown exception. We will automatically parse the error_type and details from this.
+    error_type(str, optional): The type of error e.g. "ValueError".
+    code(str, optional): A code that can be used to identify the error e.g. 501.
+    details(str, optional): Detailed information about the error.
+    logs(str, optional): For detailed information/logging related to the error.
+    """
+
+    # Inherit common Event fields
+    event_type: str = field(default=EventType.ERROR.value)
+
+    # Error-specific fields
+    trigger_event: Optional[Event] = None
+    exception: Optional[BaseException] = None
+    error_type: Optional[str] = None
+    code: Optional[str] = None
+    details: Optional[Union[str, Dict[str, str]]] = None
+    logs: Optional[str] = field(default_factory=traceback.format_exc)
+
+    def __post_init__(self):
+        """Process exception if provided"""
+        if self.exception:
+            self.error_type = self.error_type or type(self.exception).__name__
+            self.details = self.details or str(self.exception)
+            self.exception = None  # removes exception from serialization
+
+        # Ensure end timestamp is set
+        if not self.end_timestamp:
+            self.end_timestamp = get_ISO_time()
+
+    @property
+    def timestamp(self) -> str:
+        """Maintain backward compatibility with old code expecting timestamp"""
+        return self.init_timestamp

agentops/logging/__init__.py

@@ -0,0 +1,4 @@
+from agentops.logging.config import configure_logging, logger
+from agentops.logging.instrument_logging import setup_print_logger, upload_logfile
+
+__all__ = ["logger", "configure_logging", "setup_print_logger", "upload_logfile"]