mseep-agentops 0.4.18__py3-none-any.whl → 0.4.23__py3-none-any.whl

agentops/instrumentation/agentic/smolagents/instrumentor.py

@@ -0,0 +1,286 @@
+"""SmoLAgents instrumentation for AgentOps."""
+
+from typing import Dict, Any
+from opentelemetry.trace import SpanKind
+from opentelemetry.metrics import Meter
+from wrapt import wrap_function_wrapper
+
+from agentops.instrumentation.common import CommonInstrumentor, StandardMetrics, InstrumentorConfig
+from agentops.logging import logger
+
+# Library info for tracer/meter
+LIBRARY_NAME = "agentops.instrumentation.smolagents"
+LIBRARY_VERSION = "0.1.0"
+
+# Import attribute handlers
+try:
+    from agentops.instrumentation.agentic.smolagents.attributes.agent import (
+        get_agent_attributes,
+        get_tool_call_attributes,
+        get_planning_step_attributes,
+        get_agent_step_attributes,
+        get_agent_stream_attributes,
+        get_managed_agent_attributes,
+    )
+    from agentops.instrumentation.agentic.smolagents.attributes.model import (
+        get_model_attributes,
+        get_stream_attributes,
+    )
+except ImportError:
+    # Fallback functions if imports fail
+    def get_agent_attributes(*args, **kwargs):
+        return {}
+
+    def get_tool_call_attributes(*args, **kwargs):
+        return {}
+
+    def get_planning_step_attributes(*args, **kwargs):
+        return {}
+
+    def get_agent_step_attributes(*args, **kwargs):
+        return {}
+
+    def get_agent_stream_attributes(*args, **kwargs):
+        return {}
+
+    def get_managed_agent_attributes(*args, **kwargs):
+        return {}
+
+    def get_model_attributes(*args, **kwargs):
+        return {}
+
+    def get_stream_attributes(*args, **kwargs):
+        return {}
+
+
+class SmolagentsInstrumentor(CommonInstrumentor):
+    """Instrumentor for SmoLAgents library."""
+
+    def __init__(self):
+        """Initialize the SmoLAgents instrumentor."""
+        # Create instrumentor config
+        config = InstrumentorConfig(
+            library_name=LIBRARY_NAME,
+            library_version=LIBRARY_VERSION,
+            wrapped_methods=[],  # We use custom wrapping
+            metrics_enabled=True,
+            dependencies=["smolagents >= 1.0.0", "litellm"],
+        )
+
+        super().__init__(config)
+
+    def _create_metrics(self, meter: Meter) -> Dict[str, Any]:
+        """Create metrics for the instrumentor.
+
+        Returns a dictionary of metric name to metric instance.
+        """
+        # Create standard metrics for LLM operations
+        return StandardMetrics.create_standard_metrics(meter)
+
+    def _custom_wrap(self, **kwargs):
+        """Apply custom wrapping for SmoLAgents.
+
+        This is called after normal wrapping, but we use it for all wrapping
+        since we don't have normal wrapped methods.
+        """
+        # Core agent operations
+        wrap_function_wrapper("smolagents.agents", "CodeAgent.run", self._agent_run_wrapper(self._tracer))
+        wrap_function_wrapper("smolagents.agents", "ToolCallingAgent.run", self._agent_run_wrapper(self._tracer))
+
+        # Tool calling operations
+        wrap_function_wrapper(
+            "smolagents.agents", "ToolCallingAgent.execute_tool_call", self._tool_execution_wrapper(self._tracer)
+        )
+
+        # Model operations with proper model name extraction
+        wrap_function_wrapper("smolagents.models", "LiteLLMModel.generate", self._llm_wrapper(self._tracer))
+        wrap_function_wrapper("smolagents.models", "LiteLLMModel.generate_stream", self._llm_wrapper(self._tracer))
+
+        logger.info("SmoLAgents instrumentation enabled")
+
+    def _agent_run_wrapper(self, tracer):
+        """Wrapper for agent run methods."""
+
+        def wrapper(wrapped, instance, args, kwargs):
+            # Get proper agent name - handle None case
+            agent_name = getattr(instance, "name", None)
+            if not agent_name:  # Handle None, empty string, or missing attribute
+                agent_name = instance.__class__.__name__
+
+            span_name = f"{agent_name}.run"
+
+            with tracer.start_as_current_span(
+                span_name,
+                kind=SpanKind.CLIENT,
+            ) as span:
+                # Extract attributes
+                attributes = get_agent_attributes(args=(instance,) + args, kwargs=kwargs)
+
+                # Fix managed agents attribute
+                if hasattr(instance, "managed_agents") and instance.managed_agents:
+                    managed_agent_names = []
+                    for agent in instance.managed_agents:
+                        name = getattr(agent, "name", None)
+                        if not name:  # Handle None case for managed agents too
+                            name = agent.__class__.__name__
+                        managed_agent_names.append(name)
+                    attributes["agent.managed_agents"] = str(managed_agent_names)
+                else:
+                    attributes["agent.managed_agents"] = "[]"
+
+                for key, value in attributes.items():
+                    if value is not None:
+                        span.set_attribute(key, value)
+
+                try:
+                    result = wrapped(*args, **kwargs)
+
+                    # Set output attribute
+                    if result is not None:
+                        span.set_attribute("agentops.entity.output", str(result))
+
+                    return result
+                except Exception as e:
+                    span.record_exception(e)
+                    raise
+
+        return wrapper
+
+    def _tool_execution_wrapper(self, tracer):
+        """Wrapper for tool execution methods."""
+
+        def wrapper(wrapped, instance, args, kwargs):
+            # Extract tool name for better span naming
+            tool_name = "unknown"
+            if args and len(args) > 0:
+                tool_call = args[0]
+                if hasattr(tool_call, "function"):
+                    tool_name = tool_call.function.name
+
+            span_name = f"tool.{tool_name}" if tool_name != "unknown" else "tool.execute"
+
+            with tracer.start_as_current_span(
+                span_name,
+                kind=SpanKind.CLIENT,
+            ) as span:
+                # Extract tool information from kwargs or args
+                tool_params = "{}"
+
+                # Try to extract tool call information
+                if args and len(args) > 0:
+                    tool_call = args[0]
+                    if hasattr(tool_call, "function"):
+                        if hasattr(tool_call.function, "arguments"):
+                            tool_params = str(tool_call.function.arguments)
+
+                # Extract attributes
+                attributes = get_tool_call_attributes(args=(instance,) + args, kwargs=kwargs)
+
+                # Override with better tool information if available
+                if tool_name != "unknown":
+                    attributes["tool.name"] = tool_name
+                    attributes["tool.parameters"] = tool_params
+
+                for key, value in attributes.items():
+                    if value is not None:
+                        span.set_attribute(key, value)
+
+                try:
+                    result = wrapped(*args, **kwargs)
+
+                    # Set success status and result
+                    span.set_attribute("tool.status", "success")
+                    if result is not None:
+                        span.set_attribute("tool.result", str(result))
+
+                    return result
+                except Exception as e:
+                    span.set_attribute("tool.status", "error")
+                    span.record_exception(e)
+                    raise
+
+        return wrapper
+
+    def _llm_wrapper(self, tracer):
+        """Wrapper for LLM generation methods with proper model name extraction."""
+
+        def wrapper(wrapped, instance, args, kwargs):
+            # Extract model name from instance
+            model_name = getattr(instance, "model_id", "unknown")
+
+            # Determine if this is streaming
+            is_streaming = "generate_stream" in wrapped.__name__
+            operation = "generate_stream" if is_streaming else "generate"
+            span_name = f"litellm.{operation} ({model_name})" if model_name != "unknown" else f"litellm.{operation}"
+
+            with tracer.start_as_current_span(
+                span_name,
+                kind=SpanKind.CLIENT,
+            ) as span:
+                # Extract attributes
+                if is_streaming:
+                    attributes = get_stream_attributes(args=(instance,) + args, kwargs=kwargs)
+                else:
+                    attributes = get_model_attributes(args=(instance,) + args, kwargs=kwargs)
+
+                # Ensure model name is properly set
+                attributes["gen_ai.request.model"] = model_name
+
+                for key, value in attributes.items():
+                    if value is not None:
+                        span.set_attribute(key, value)
+
+                try:
+                    result = wrapped(*args, **kwargs)
+
+                    # Extract response attributes if available
+                    if result and hasattr(result, "content"):
+                        span.set_attribute("gen_ai.completion.0.content", str(result.content))
+                    if result and hasattr(result, "token_usage"):
+                        token_usage = result.token_usage
+                        if hasattr(token_usage, "input_tokens"):
+                            span.set_attribute("gen_ai.usage.prompt_tokens", token_usage.input_tokens)
+                        if hasattr(token_usage, "output_tokens"):
+                            span.set_attribute("gen_ai.usage.completion_tokens", token_usage.output_tokens)
+
+                    return result
+                except Exception as e:
+                    span.record_exception(e)
+                    raise
+
+        return wrapper
+
+    def _custom_unwrap(self, **kwargs):
+        """Remove custom wrapping from SmoLAgents.
+
+        This method removes all custom wrapping we applied.
+        """
+        # Unwrap all instrumented methods
+        from opentelemetry.instrumentation.utils import unwrap
+
+        try:
+            unwrap("smolagents.agents", "CodeAgent.run")
+        except Exception as e:
+            logger.debug(f"Failed to unwrap CodeAgent.run: {e}")
+
+        try:
+            unwrap("smolagents.agents", "ToolCallingAgent.run")
+        except Exception as e:
+            logger.debug(f"Failed to unwrap ToolCallingAgent.run: {e}")
+
+        try:
+            unwrap("smolagents.agents", "ToolCallingAgent.execute_tool_call")
+        except Exception as e:
+            logger.debug(f"Failed to unwrap ToolCallingAgent.execute_tool_call: {e}")
+
+        try:
+            unwrap("smolagents.models", "LiteLLMModel.generate")
+        except Exception as e:
+            logger.debug(f"Failed to unwrap LiteLLMModel.generate: {e}")
+
+        try:
+            unwrap("smolagents.models", "LiteLLMModel.generate_stream")
+        except Exception as e:
+            logger.debug(f"Failed to unwrap LiteLLMModel.generate_stream: {e}")
+
+        logger.info("SmoLAgents instrumentation disabled")

agentops/instrumentation/agentic/smolagents/stream_wrapper.py

@@ -0,0 +1,258 @@
+"""Stream wrapper for SmoLAgents model streaming responses."""
+
+import time
+import uuid
+from typing import Any, Generator, Optional
+from opentelemetry.trace import Status, StatusCode, Span
+
+from agentops.semconv.message import MessageAttributes
+from agentops.semconv.agent import AgentAttributes
+from agentops.semconv.tool import ToolAttributes
+from .attributes.model import get_stream_attributes
+from agentops.semconv.span_attributes import SpanAttributes
+
+
+def model_stream_wrapper(tracer):
+    """Wrapper for model streaming methods.
+
+    Args:
+        tracer: OpenTelemetry tracer
+
+    Returns:
+        Wrapped function
+    """
+
+    def wrapper(wrapped, instance, args, kwargs):
+        messages = kwargs.get("messages", [])
+        model_id = instance.model_id if hasattr(instance, "model_id") else "unknown"
+
+        with tracer.start_as_current_span(
+            name=f"{model_id}.generate_stream", attributes=get_stream_attributes(model_id=model_id, messages=messages)
+        ) as span:
+            try:
+                # Start streaming
+                stream = wrapped(*args, **kwargs)
+                first_token_received = False
+                start_time = time.time()
+                accumulated_text = ""
+
+                # Process stream
+                for chunk in stream:
+                    if not first_token_received:
+                        first_token_received = True
+                        span.set_attribute("gen_ai.time_to_first_token", time.time() - start_time)
+
+                    # Accumulate text and update attributes
+                    if hasattr(chunk, "content") and chunk.content:
+                        accumulated_text += chunk.content
+                        span.set_attribute(MessageAttributes.COMPLETION_CONTENT.format(i=0), accumulated_text)
+                        span.set_attribute(MessageAttributes.COMPLETION_TYPE.format(i=0), "text")
+
+                    yield chunk
+
+                # Set final attributes
+                span.set_attribute("gen_ai.streaming_duration", time.time() - start_time)
+                span.set_status(Status(StatusCode.OK))
+
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                raise
+
+    return wrapper
+
+
+def agent_stream_wrapper(tracer):
+    """Wrapper for agent streaming methods.
+
+    Args:
+        tracer: OpenTelemetry tracer
+
+    Returns:
+        Wrapped function
+    """
+
+    def wrapper(wrapped, instance, args, kwargs):
+        task = kwargs.get("task", args[0] if args else "unknown")
+        agent_type = instance.__class__.__name__
+        agent_id = str(uuid.uuid4())
+
+        with tracer.start_as_current_span(
+            name=f"{agent_type}.run_stream",
+            attributes={
+                AgentAttributes.AGENT_ID: agent_id,
+                AgentAttributes.AGENT_NAME: agent_type,
+                AgentAttributes.AGENT_ROLE: "executor",
+                AgentAttributes.AGENT_REASONING: task,
+            },
+        ) as span:
+            try:
+                # Initialize counters
+                step_count = 0
+                planning_steps = 0
+                tools_used = set()
+                start_time = time.time()
+
+                # Process stream
+                stream = wrapped(*args, **kwargs)
+                for step in stream:
+                    step_count += 1
+
+                    # Track step types
+                    if hasattr(step, "type"):
+                        if step.type == "planning":
+                            planning_steps += 1
+                        elif step.type == "tool_call":
+                            tools_used.add(step.tool_name)
+                            # Add tool-specific attributes
+                            span.set_attribute(ToolAttributes.TOOL_NAME, step.tool_name)
+                            if hasattr(step, "arguments"):
+                                span.set_attribute(ToolAttributes.TOOL_PARAMETERS, step.arguments)
+
+                    # Update span attributes
+                    span.set_attribute("agent.step_count", step_count)
+                    span.set_attribute("agent.planning_steps", planning_steps)
+                    span.set_attribute(AgentAttributes.AGENT_TOOLS, list(tools_used))
+
+                    yield step
+
+                # Set final attributes
+                span.set_attribute("agent.execution_time", time.time() - start_time)
+                span.set_status(Status(StatusCode.OK))
+
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                raise
+
+    return wrapper
+
+
+class SmoLAgentsStreamWrapper:
+    """Wrapper for streaming responses from SmoLAgents models."""
+
+    def __init__(
+        self,
+        stream: Generator,
+        span: Span,
+        model_id: Optional[str] = None,
+    ):
+        """Initialize the stream wrapper.
+
+        Args:
+            stream: The original generator from the model
+            span: The OpenTelemetry span to track the stream
+            model_id: Optional model identifier
+        """
+        self._stream = stream
+        self._span = span
+        self._model_id = model_id
+        self._chunks_received = 0
+        self._full_content = []
+        self._tool_calls = []
+        self._current_tool_call = None
+        self._token_count = 0
+
+    def __iter__(self):
+        """Iterate over the stream."""
+        return self
+
+    def __next__(self):
+        """Get the next chunk from the stream."""
+        try:
+            chunk = next(self._stream)
+            self._process_chunk(chunk)
+            return chunk
+        except StopIteration:
+            self._finalize_stream()
+            raise
+
+    def _process_chunk(self, chunk: Any) -> None:
+        """Process a chunk from the stream.
+
+        Args:
+            chunk: The chunk to process
+        """
+        self._chunks_received += 1
+
+        # Handle ChatMessageStreamDelta objects
+        if hasattr(chunk, "content") and chunk.content:
+            self._full_content.append(chunk.content)
+
+        # Handle tool calls in chunks
+        if hasattr(chunk, "tool_calls") and chunk.tool_calls:
+            for tool_call in chunk.tool_calls:
+                if tool_call.id not in [tc["id"] for tc in self._tool_calls]:
+                    self._tool_calls.append(
+                        {
+                            "id": tool_call.id,
+                            "type": tool_call.type,
+                            "name": tool_call.function.name,
+                            "arguments": tool_call.function.arguments,
+                        }
+                    )
+
+        # Track token usage if available
+        if hasattr(chunk, "token_usage") and chunk.token_usage:
+            if hasattr(chunk.token_usage, "output_tokens"):
+                self._token_count += chunk.token_usage.output_tokens
+
+        # Update span with chunk information
+        self._span.add_event(
+            "stream_chunk_received",
+            {
+                "chunk_number": self._chunks_received,
+                "chunk_content_length": len(chunk.content) if hasattr(chunk, "content") and chunk.content else 0,
+            },
+        )
+
+    def _finalize_stream(self) -> None:
+        """Finalize the stream and update span attributes."""
+        # Combine all content chunks
+        full_content = "".join(self._full_content)
+
+        # Set final attributes on the span
+        attributes = {
+            MessageAttributes.COMPLETION_CONTENT.format(i=0): full_content,
+            "stream.chunks_received": self._chunks_received,
+            "stream.total_content_length": len(full_content),
+        }
+
+        # Add tool calls if any
+        if self._tool_calls:
+            for j, tool_call in enumerate(self._tool_calls):
+                attributes.update(
+                    {
+                        MessageAttributes.COMPLETION_TOOL_CALL_ID.format(i=0, j=j): tool_call["id"],
+                        MessageAttributes.COMPLETION_TOOL_CALL_TYPE.format(i=0, j=j): tool_call["type"],
+                        MessageAttributes.COMPLETION_TOOL_CALL_NAME.format(i=0, j=j): tool_call["name"],
+                        MessageAttributes.COMPLETION_TOOL_CALL_ARGUMENTS.format(i=0, j=j): str(tool_call["arguments"]),
+                    }
+                )
+
+        # Add token usage if tracked
+        if self._token_count > 0:
+            attributes[SpanAttributes.LLM_USAGE_STREAMING_TOKENS] = self._token_count
+
+        self._span.set_attributes(attributes)
+
+    def close(self) -> None:
+        """Close the stream wrapper."""
+        if hasattr(self._stream, "close"):
+            self._stream.close()
+
+
+def wrap_stream(
+    stream: Generator,
+    span: Span,
+    model_id: Optional[str] = None,
+) -> SmoLAgentsStreamWrapper:
+    """Wrap a streaming response from a SmoLAgents model.
+
+    Args:
+        stream: The original generator from the model
+        span: The OpenTelemetry span to track the stream
+        model_id: Optional model identifier
+
+    Returns:
+        SmoLAgentsStreamWrapper: The wrapped stream
+    """
+    return SmoLAgentsStreamWrapper(stream, span, model_id)

agentops/instrumentation/agentic/xpander/__init__.py

@@ -0,0 +1,15 @@
+"""Xpander SDK instrumentation for AgentOps."""
+
+from agentops.instrumentation.agentic.xpander.instrumentor import XpanderInstrumentor
+from agentops.instrumentation.agentic.xpander.trace_probe import (
+    wrap_openai_call_for_xpander,
+    is_xpander_session_active,
+    get_active_xpander_session,
+)
+
+__all__ = [
+    "XpanderInstrumentor",
+    "wrap_openai_call_for_xpander",
+    "is_xpander_session_active",
+    "get_active_xpander_session",
+]

agentops/instrumentation/agentic/xpander/context.py

@@ -0,0 +1,112 @@
+"""Xpander context management for session tracking."""
+
+import time
+import threading
+from typing import Any, Dict, Optional
+
+
+class XpanderContext:
+    """Context manager for Xpander sessions with nested conversation spans."""
+
+    def __init__(self):
+        self._sessions = {}  # session_id -> session_data
+        self._workflow_spans = {}  # session_id -> active workflow span
+        self._agent_spans = {}  # session_id -> active agent span
+        self._conversation_spans = {}  # session_id -> active conversation span
+        self._conversation_counters = {}  # session_id -> conversation counter
+        self._lock = threading.Lock()
+
+    def start_session(self, session_id: str, agent_info: Dict[str, Any], workflow_span=None, agent_span=None) -> None:
+        """Start a new session with agent info."""
+        with self._lock:
+            self._sessions[session_id] = {
+                "agent_name": agent_info.get("agent_name", "unknown"),
+                "agent_id": agent_info.get("agent_id", "unknown"),
+                "task_input": agent_info.get("task_input"),
+                "phase": "planning",
+                "step_count": 0,
+                "total_tokens": 0,
+                "tools_executed": [],
+                "start_time": time.time(),
+            }
+            if workflow_span:
+                self._workflow_spans[session_id] = workflow_span
+            if agent_span:
+                self._agent_spans[session_id] = agent_span
+
+            # Initialize conversation counter
+            self._conversation_counters[session_id] = 0
+
+    def start_conversation(self, session_id: str, conversation_span) -> None:
+        """Start a new conversation within the session."""
+        with self._lock:
+            self._conversation_spans[session_id] = conversation_span
+            self._conversation_counters[session_id] = self._conversation_counters.get(session_id, 0) + 1
+
+    def end_conversation(self, session_id: str) -> None:
+        """End the current conversation."""
+        with self._lock:
+            if session_id in self._conversation_spans:
+                del self._conversation_spans[session_id]
+
+    def has_active_conversation(self, session_id: str) -> bool:
+        """Check if there's an active conversation for this session."""
+        with self._lock:
+            return session_id in self._conversation_spans
+
+    def get_conversation_counter(self, session_id: str) -> int:
+        """Get the current conversation counter."""
+        with self._lock:
+            return self._conversation_counters.get(session_id, 0)
+
+    def get_session(self, session_id: str) -> Optional[Dict[str, Any]]:
+        """Get session data."""
+        with self._lock:
+            return self._sessions.get(session_id)
+
+    def update_session(self, session_id: str, updates: Dict[str, Any]) -> None:
+        """Update session data."""
+        with self._lock:
+            if session_id in self._sessions:
+                self._sessions[session_id].update(updates)
+
+    def end_session(self, session_id: str) -> None:
+        """End a session."""
+        with self._lock:
+            if session_id in self._sessions:
+                del self._sessions[session_id]
+            if session_id in self._workflow_spans:
+                del self._workflow_spans[session_id]
+            if session_id in self._agent_spans:
+                del self._agent_spans[session_id]
+            if session_id in self._conversation_spans:
+                del self._conversation_spans[session_id]
+            if session_id in self._conversation_counters:
+                del self._conversation_counters[session_id]
+
+    def get_workflow_phase(self, session_id: str) -> str:
+        """Detect current workflow phase based on state."""
+        with self._lock:
+            session = self._sessions.get(session_id, {})
+
+            if session.get("tools_executed", []):
+                return "executing"
+            elif session.get("step_count", 0) > 0:
+                return "executing"
+            else:
+                return "planning"
+
+    def get_workflow_span(self, session_id: str):
+        """Get the active workflow span for a session."""
+        with self._lock:
+            return self._workflow_spans.get(session_id)
+
+    def get_agent_span(self, session_id: str):
+        """Get the active agent span for a session."""
+        with self._lock:
+            return self._agent_spans.get(session_id)
+
+    def get_conversation_span(self, session_id: str):
+        """Get the active conversation span for a session."""
+        with self._lock:
+            return self._conversation_spans.get(session_id)