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

agentops/instrumentation/agentic/agno/instrumentor.py

@@ -0,0 +1,1313 @@
+"""Agno Agent Instrumentation for AgentOps
+
+This module provides instrumentation for the Agno Agent library, implementing OpenTelemetry
+instrumentation for agent workflows and LLM model calls.
+
+This provides clean visibility into agent workflows and actual tool usage with proper
+parent-child span relationships.
+"""
+
+from typing import List, Any, Optional, Dict
+from opentelemetry import trace, context as otel_context
+from opentelemetry.trace import Status, StatusCode
+from opentelemetry.metrics import Meter
+import threading
+import json
+
+from agentops.logging import logger
+from agentops.instrumentation.common import (
+    CommonInstrumentor,
+    StandardMetrics,
+    InstrumentorConfig,
+)
+from agentops.instrumentation.common.wrappers import WrapConfig
+
+# Import attribute handlers
+from agentops.instrumentation.agentic.agno.attributes import (
+    get_agent_run_attributes,
+    get_metrics_attributes,
+    get_team_run_attributes,
+    get_tool_execution_attributes,
+    get_workflow_run_attributes,
+    get_workflow_session_attributes,
+    get_storage_read_attributes,
+    get_storage_write_attributes,
+)
+
+
+class StreamingContextManager:
+    """Manages span contexts for streaming agent and workflow executions."""
+
+    def __init__(self):
+        self._contexts = {}  # context_id -> (span_context, span)
+        self._agent_sessions = {}  # session_id -> agent_id mapping for context lookup
+        self._lock = threading.Lock()
+
+    def store_context(self, context_id: str, span_context: Any, span: Any) -> None:
+        """Store span context for streaming execution."""
+        with self._lock:
+            self._contexts[context_id] = (span_context, span)
+
+    def get_context(self, context_id: str) -> Optional[tuple]:
+        """Retrieve stored span context."""
+        with self._lock:
+            return self._contexts.get(context_id)
+
+    def remove_context(self, context_id: str) -> None:
+        """Remove stored context (when streaming completes)."""
+        with self._lock:
+            self._contexts.pop(context_id, None)
+
+    def store_agent_session_mapping(self, session_id: str, agent_id: str) -> None:
+        """Store mapping between session and agent for context lookup."""
+        with self._lock:
+            self._agent_sessions[session_id] = agent_id
+
+    def get_agent_context_by_session(self, session_id: str) -> Optional[tuple]:
+        """Get agent context using session ID."""
+        with self._lock:
+            agent_id = self._agent_sessions.get(session_id)
+            if agent_id:
+                return self._contexts.get(agent_id)
+            return None
+
+    def clear_all(self) -> None:
+        """Clear all stored contexts."""
+        with self._lock:
+            self._contexts.clear()
+            self._agent_sessions.clear()
+
+
+# Methods to wrap for instrumentation
+# Empty list - all wrapping will be done in _custom_wrap to avoid circular imports
+WRAPPED_METHODS: List[WrapConfig] = []
+
+
+class StreamingResultWrapper:
+    """Wrapper for streaming results that maintains agent span as active throughout iteration."""
+
+    def __init__(self, original_result, span, agent_id, agent_context, streaming_context_manager):
+        self.original_result = original_result
+        self.span = span
+        self.agent_id = agent_id
+        self.agent_context = agent_context
+        self.streaming_context_manager = streaming_context_manager
+        self._consumed = False
+
+    def __iter__(self):
+        """Return iterator that keeps agent span active during iteration."""
+        context_token = otel_context.attach(self.agent_context)
+        try:
+            # Execute iteration within agent context
+            for item in self.original_result:
+                # Each item is yielded within the agent span context
+                yield item
+        finally:
+            # Clean up when iteration is complete
+            otel_context.detach(context_token)
+            if not self._consumed:
+                self._consumed = True
+                self.span.end()
+                self.streaming_context_manager.remove_context(self.agent_id)
+
+    def __getattr__(self, name):
+        """Delegate attribute access to the original result."""
+        return getattr(self.original_result, name)
+
+
+def create_streaming_workflow_wrapper(tracer, streaming_context_manager):
+    """Create a streaming-aware wrapper for workflow run methods."""
+
+    def wrapper(wrapped, instance, args, kwargs):
+        # Get workflow ID for context storage
+        workflow_id = getattr(instance, "workflow_id", None) or getattr(instance, "id", None) or id(instance)
+        workflow_id = str(workflow_id)
+
+        # Get workflow name for span naming
+        workflow_name = getattr(instance, "name", None) or type(instance).__name__
+        span_name = f"{workflow_name}.agno.workflow.run.workflow" if workflow_name else "agno.workflow.run.workflow"
+
+        # Check if streaming is enabled
+        is_streaming = kwargs.get("stream", getattr(instance, "stream", False))
+
+        # For streaming, manually manage span lifecycle
+        if is_streaming:
+            span = tracer.start_span(span_name)
+
+            try:
+                # Set workflow attributes
+                attributes = get_workflow_run_attributes(args=(instance,) + args, kwargs=kwargs)
+                for key, value in attributes.items():
+                    span.set_attribute(key, value)
+
+                # Store context for streaming - capture current context with active span
+                current_context = trace.set_span_in_context(span, otel_context.get_current())
+                streaming_context_manager.store_context(workflow_id, current_context, span)
+
+                # Execute the original function within workflow context
+                context_token = otel_context.attach(current_context)
+                try:
+                    result = wrapped(*args, **kwargs)
+                finally:
+                    otel_context.detach(context_token)
+
+                # Set result attributes
+                result_attributes = get_workflow_run_attributes(
+                    args=(instance,) + args, kwargs=kwargs, return_value=result
+                )
+                for key, value in result_attributes.items():
+                    if key not in attributes:  # Avoid duplicates
+                        span.set_attribute(key, value)
+
+                span.set_status(Status(StatusCode.OK))
+
+                # For streaming results, we need to keep the span open
+                # The span will be closed when streaming completes
+                return result
+
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                span.end()
+                streaming_context_manager.remove_context(workflow_id)
+                raise
+        else:
+            # For non-streaming, use normal context manager
+            with tracer.start_as_current_span(span_name) as span:
+                try:
+                    # Set workflow attributes
+                    attributes = get_workflow_run_attributes(args=(instance,) + args, kwargs=kwargs)
+                    for key, value in attributes.items():
+                        span.set_attribute(key, value)
+
+                    # Execute the original function
+                    result = wrapped(*args, **kwargs)
+
+                    # Set result attributes
+                    result_attributes = get_workflow_run_attributes(
+                        args=(instance,) + args, kwargs=kwargs, return_value=result
+                    )
+                    for key, value in result_attributes.items():
+                        if key not in attributes:  # Avoid duplicates
+                            span.set_attribute(key, value)
+
+                    span.set_status(Status(StatusCode.OK))
+                    return result
+
+                except Exception as e:
+                    span.set_status(Status(StatusCode.ERROR, str(e)))
+                    span.record_exception(e)
+                    raise
+
+    return wrapper
+
+
+def create_streaming_workflow_async_wrapper(tracer, streaming_context_manager):
+    """Create a streaming-aware async wrapper for workflow run methods."""
+
+    async def wrapper(wrapped, instance, args, kwargs):
+        # Get workflow ID for context storage
+        workflow_id = getattr(instance, "workflow_id", None) or getattr(instance, "id", None) or id(instance)
+        workflow_id = str(workflow_id)
+
+        # Get workflow name for span naming
+        workflow_name = getattr(instance, "name", None) or type(instance).__name__
+        span_name = f"{workflow_name}.agno.workflow.run.workflow" if workflow_name else "agno.workflow.run.workflow"
+
+        # Check if streaming is enabled
+        is_streaming = kwargs.get("stream", getattr(instance, "stream", False))
+
+        # For streaming, manually manage span lifecycle
+        if is_streaming:
+            span = tracer.start_span(span_name)
+
+            try:
+                # Set workflow attributes
+                attributes = get_workflow_run_attributes(args=(instance,) + args, kwargs=kwargs)
+                for key, value in attributes.items():
+                    span.set_attribute(key, value)
+
+                # Store context for streaming - capture current context with active span
+                current_context = trace.set_span_in_context(span, otel_context.get_current())
+                streaming_context_manager.store_context(workflow_id, current_context, span)
+
+                # Execute the original function within workflow context
+                context_token = otel_context.attach(current_context)
+                try:
+                    result = await wrapped(*args, **kwargs)
+                finally:
+                    otel_context.detach(context_token)
+
+                # Set result attributes
+                result_attributes = get_workflow_run_attributes(
+                    args=(instance,) + args, kwargs=kwargs, return_value=result
+                )
+                for key, value in result_attributes.items():
+                    if key not in attributes:  # Avoid duplicates
+                        span.set_attribute(key, value)
+
+                span.set_status(Status(StatusCode.OK))
+
+                # For streaming results, we need to keep the span open
+                # The span will be closed when streaming completes
+                return result
+
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                span.end()
+                streaming_context_manager.remove_context(workflow_id)
+                raise
+        else:
+            # For non-streaming, use normal context manager
+            with tracer.start_as_current_span(span_name) as span:
+                try:
+                    # Set workflow attributes
+                    attributes = get_workflow_run_attributes(args=(instance,) + args, kwargs=kwargs)
+                    for key, value in attributes.items():
+                        span.set_attribute(key, value)
+
+                    # Execute the original function
+                    result = await wrapped(*args, **kwargs)
+
+                    # Set result attributes
+                    result_attributes = get_workflow_run_attributes(
+                        args=(instance,) + args, kwargs=kwargs, return_value=result
+                    )
+                    for key, value in result_attributes.items():
+                        if key not in attributes:  # Avoid duplicates
+                            span.set_attribute(key, value)
+
+                    span.set_status(Status(StatusCode.OK))
+                    return result
+
+                except Exception as e:
+                    span.set_status(Status(StatusCode.ERROR, str(e)))
+                    span.record_exception(e)
+                    raise
+
+    return wrapper
+
+
+def create_streaming_agent_wrapper(tracer, streaming_context_manager):
+    """Create a streaming-aware wrapper for agent run methods."""
+
+    def wrapper(wrapped, instance, args, kwargs):
+        # Get agent ID for context storage
+        agent_id = getattr(instance, "agent_id", None) or getattr(instance, "id", None) or id(instance)
+        agent_id = str(agent_id)
+
+        # Get session ID for context mapping
+        session_id = getattr(instance, "session_id", None)
+
+        # Get agent name for span naming
+        agent_name = getattr(instance, "name", None)
+        span_name = f"{agent_name}.agno.agent.run.agent" if agent_name else "agno.agent.run.agent"
+
+        # Check if streaming is enabled
+        is_streaming = kwargs.get("stream", getattr(instance, "stream", False))
+
+        # For streaming, manually manage span lifecycle
+        if is_streaming:
+            span = tracer.start_span(span_name)
+
+            try:
+                # Set agent attributes
+                attributes = get_agent_run_attributes(args=(instance,) + args, kwargs=kwargs)
+                for key, value in attributes.items():
+                    span.set_attribute(key, value)
+
+                # Store context for streaming - capture current context with active span
+                current_context = trace.set_span_in_context(span, otel_context.get_current())
+                streaming_context_manager.store_context(agent_id, current_context, span)
+
+                # Store session-to-agent mapping for LLM context lookup
+                if session_id:
+                    streaming_context_manager.store_agent_session_mapping(session_id, agent_id)
+
+                # Execute the original function within agent context
+                context_token = otel_context.attach(current_context)
+                try:
+                    result = wrapped(*args, **kwargs)
+                finally:
+                    otel_context.detach(context_token)
+
+                # Set result attributes
+                result_attributes = get_agent_run_attributes(
+                    args=(instance,) + args, kwargs=kwargs, return_value=result
+                )
+                for key, value in result_attributes.items():
+                    if key not in attributes:  # Avoid duplicates
+                        span.set_attribute(key, value)
+
+                span.set_status(Status(StatusCode.OK))
+
+                # Wrap the result to maintain context and end span when complete
+                if hasattr(result, "__iter__"):
+                    return StreamingResultWrapper(result, span, agent_id, current_context, streaming_context_manager)
+                else:
+                    # Not actually streaming, clean up immediately
+                    span.end()
+                    streaming_context_manager.remove_context(agent_id)
+                    return result
+
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                span.end()
+                streaming_context_manager.remove_context(agent_id)
+                raise
+        else:
+            # For non-streaming, use normal context manager
+            with tracer.start_as_current_span(span_name) as span:
+                try:
+                    # Set agent attributes
+                    attributes = get_agent_run_attributes(args=(instance,) + args, kwargs=kwargs)
+                    for key, value in attributes.items():
+                        span.set_attribute(key, value)
+
+                    # Execute the original function
+                    result = wrapped(*args, **kwargs)
+
+                    # Set result attributes
+                    result_attributes = get_agent_run_attributes(
+                        args=(instance,) + args, kwargs=kwargs, return_value=result
+                    )
+                    for key, value in result_attributes.items():
+                        if key not in attributes:  # Avoid duplicates
+                            span.set_attribute(key, value)
+
+                    span.set_status(Status(StatusCode.OK))
+                    return result
+
+                except Exception as e:
+                    span.set_status(Status(StatusCode.ERROR, str(e)))
+                    span.record_exception(e)
+                    raise
+
+    return wrapper
+
+
+def create_streaming_agent_async_wrapper(tracer, streaming_context_manager):
+    """Create a streaming-aware async wrapper for agent run methods."""
+
+    async def wrapper(wrapped, instance, args, kwargs):
+        # Get agent ID for context storage
+        agent_id = getattr(instance, "agent_id", None) or getattr(instance, "id", None) or id(instance)
+        agent_id = str(agent_id)
+
+        # Get session ID for context mapping
+        session_id = getattr(instance, "session_id", None)
+
+        # Get agent name for span naming
+        agent_name = getattr(instance, "name", None)
+        span_name = f"{agent_name}.agno.agent.run.agent" if agent_name else "agno.agent.run.agent"
+
+        # Check if streaming is enabled
+        is_streaming = kwargs.get("stream", getattr(instance, "stream", False))
+
+        # For streaming, manually manage span lifecycle
+        if is_streaming:
+            span = tracer.start_span(span_name)
+
+            try:
+                # Set agent attributes
+                attributes = get_agent_run_attributes(args=(instance,) + args, kwargs=kwargs)
+                for key, value in attributes.items():
+                    span.set_attribute(key, value)
+
+                # Store context for streaming - capture current context with active span
+                current_context = trace.set_span_in_context(span, otel_context.get_current())
+                streaming_context_manager.store_context(agent_id, current_context, span)
+
+                # Store session-to-agent mapping for LLM context lookup
+                if session_id:
+                    streaming_context_manager.store_agent_session_mapping(session_id, agent_id)
+
+                # Execute the original function within agent context
+                context_token = otel_context.attach(current_context)
+                try:
+                    result = await wrapped(*args, **kwargs)
+                finally:
+                    otel_context.detach(context_token)
+
+                # Set result attributes
+                result_attributes = get_agent_run_attributes(
+                    args=(instance,) + args, kwargs=kwargs, return_value=result
+                )
+                for key, value in result_attributes.items():
+                    if key not in attributes:  # Avoid duplicates
+                        span.set_attribute(key, value)
+
+                span.set_status(Status(StatusCode.OK))
+
+                # Wrap the result to maintain context and end span when complete
+                if hasattr(result, "__iter__"):
+                    return StreamingResultWrapper(result, span, agent_id, current_context, streaming_context_manager)
+                else:
+                    # Not actually streaming, clean up immediately
+                    span.end()
+                    streaming_context_manager.remove_context(agent_id)
+                    return result
+
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                span.end()
+                streaming_context_manager.remove_context(agent_id)
+                raise
+        else:
+            # For non-streaming, use normal context manager
+            with tracer.start_as_current_span(span_name) as span:
+                try:
+                    # Set agent attributes
+                    attributes = get_agent_run_attributes(args=(instance,) + args, kwargs=kwargs)
+                    for key, value in attributes.items():
+                        span.set_attribute(key, value)
+
+                    # Execute the original function
+                    result = await wrapped(*args, **kwargs)
+
+                    # Set result attributes
+                    result_attributes = get_agent_run_attributes(
+                        args=(instance,) + args, kwargs=kwargs, return_value=result
+                    )
+                    for key, value in result_attributes.items():
+                        if key not in attributes:  # Avoid duplicates
+                            span.set_attribute(key, value)
+
+                    span.set_status(Status(StatusCode.OK))
+                    return result
+
+                except Exception as e:
+                    span.set_status(Status(StatusCode.ERROR, str(e)))
+                    span.record_exception(e)
+                    raise
+
+    return wrapper
+
+
+def create_streaming_tool_wrapper(tracer, streaming_context_manager):
+    """Create a streaming-aware wrapper for tool execution methods."""
+
+    def wrapper(wrapped, instance, args, kwargs):
+        # Try to find the agent or workflow context for proper span hierarchy
+        parent_context = None
+        parent_span = None
+
+        # Try to get context from agent
+        try:
+            if hasattr(instance, "_agent"):
+                agent = instance._agent
+                agent_id = getattr(agent, "agent_id", None) or getattr(agent, "id", None) or id(agent)
+                agent_id = str(agent_id)
+                context_info = streaming_context_manager.get_context(agent_id)
+                if context_info:
+                    parent_context, parent_span = context_info
+        except Exception:
+            pass  # Continue without agent context if not found
+
+        # Try to get context from workflow if agent context not found
+        if not parent_context:
+            try:
+                if hasattr(instance, "_workflow"):
+                    workflow = instance._workflow
+                    workflow_id = (
+                        getattr(workflow, "workflow_id", None) or getattr(workflow, "id", None) or id(workflow)
+                    )
+                    workflow_id = str(workflow_id)
+                    context_info = streaming_context_manager.get_context(workflow_id)
+                    if context_info:
+                        parent_context, parent_span = context_info
+            except Exception:
+                pass  # Continue without workflow context if not found
+
+        # Use parent context if available, otherwise use current context
+        if parent_context:
+            context_token = otel_context.attach(parent_context)
+            try:
+                with tracer.start_as_current_span("agno.tool.execute.tool_usage") as span:
+                    try:
+                        # Set tool attributes
+                        attributes = get_tool_execution_attributes(args=(instance,) + args, kwargs=kwargs)
+                        for key, value in attributes.items():
+                            span.set_attribute(key, value)
+
+                        # Execute the original function
+                        result = wrapped(*args, **kwargs)
+
+                        # Set result attributes
+                        result_attributes = get_tool_execution_attributes(
+                            args=(instance,) + args, kwargs=kwargs, return_value=result
+                        )
+                        for key, value in result_attributes.items():
+                            if key not in attributes:  # Avoid duplicates
+                                span.set_attribute(key, value)
+
+                        span.set_status(Status(StatusCode.OK))
+                        return result
+
+                    except Exception as e:
+                        span.set_status(Status(StatusCode.ERROR, str(e)))
+                        span.record_exception(e)
+                        raise
+            finally:
+                otel_context.detach(context_token)
+        else:
+            # Fallback to normal span creation
+            with tracer.start_as_current_span("agno.tool.execute.tool_usage") as span:
+                try:
+                    # Set tool attributes
+                    attributes = get_tool_execution_attributes(args=(instance,) + args, kwargs=kwargs)
+                    for key, value in attributes.items():
+                        span.set_attribute(key, value)
+
+                    # Execute the original function
+                    result = wrapped(*args, **kwargs)
+
+                    # Set result attributes
+                    result_attributes = get_tool_execution_attributes(
+                        args=(instance,) + args, kwargs=kwargs, return_value=result
+                    )
+                    for key, value in result_attributes.items():
+                        if key not in attributes:  # Avoid duplicates
+                            span.set_attribute(key, value)
+
+                    span.set_status(Status(StatusCode.OK))
+                    return result
+
+                except Exception as e:
+                    span.set_status(Status(StatusCode.ERROR, str(e)))
+                    span.record_exception(e)
+                    raise
+
+    return wrapper
+
+
+def create_metrics_wrapper(tracer, streaming_context_manager):
+    """Create a wrapper for metrics methods with dynamic span naming."""
+
+    def wrapper(wrapped, instance, args, kwargs):
+        # Extract model ID for dynamic span naming
+        span_name = "agno.agent.metrics"  # fallback
+        if hasattr(instance, "model") and instance.model and hasattr(instance.model, "id"):
+            model_id = str(instance.model.id)
+            span_name = f"{model_id}.llm"
+
+        with tracer.start_as_current_span(span_name) as span:
+            try:
+                # Set attributes
+                attributes = get_metrics_attributes(args=(instance,) + args, kwargs=kwargs)
+                for key, value in attributes.items():
+                    span.set_attribute(key, value)
+
+                span.set_status(Status(StatusCode.OK))
+
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                raise
+
+    return wrapper
+
+
+def create_team_internal_wrapper(tracer, streaming_context_manager):
+    """Create a wrapper for Team internal methods (_run/_arun) that manages team span lifecycle."""
+
+    def wrapper(wrapped, instance, args, kwargs):
+        # Get team ID for context storage
+        team_id = getattr(instance, "team_id", None) or getattr(instance, "id", None) or id(instance)
+        team_id = str(team_id)
+
+        # Get team name for span naming
+        team_name = getattr(instance, "name", None)
+        span_name = f"{team_name}.agno.team.run.workflow" if team_name else "agno.team.run.workflow"
+
+        # Check if we already have a team context (from print_response)
+        existing_context = streaming_context_manager.get_context(team_id)
+
+        if existing_context:
+            # We're being called from print_response, use existing context
+            parent_context, parent_span = existing_context
+
+            # Execute within the existing team context
+            context_token = otel_context.attach(parent_context)
+            try:
+                with tracer.start_as_current_span(span_name) as span:
+                    try:
+                        # Set workflow attributes
+                        attributes = get_team_run_attributes(args=(instance,) + args, kwargs=kwargs)
+                        for key, value in attributes.items():
+                            span.set_attribute(key, value)
+
+                        # Execute the original function
+                        result = wrapped(*args, **kwargs)
+
+                        span.set_status(Status(StatusCode.OK))
+                        return result
+
+                    except Exception as e:
+                        span.set_status(Status(StatusCode.ERROR, str(e)))
+                        span.record_exception(e)
+                        raise
+                    finally:
+                        # Close the parent team span when workflow completes
+                        if parent_span:
+                            parent_span.end()
+                            streaming_context_manager.remove_context(team_id)
+            finally:
+                otel_context.detach(context_token)
+        else:
+            # Direct call to _run, create new team span
+            with tracer.start_as_current_span(span_name) as span:
+                try:
+                    # Set workflow attributes
+                    attributes = get_team_run_attributes(args=(instance,) + args, kwargs=kwargs)
+                    for key, value in attributes.items():
+                        span.set_attribute(key, value)
+
+                    # Execute the original function
+                    result = wrapped(*args, **kwargs)
+
+                    span.set_status(Status(StatusCode.OK))
+                    return result
+
+                except Exception as e:
+                    span.set_status(Status(StatusCode.ERROR, str(e)))
+                    span.record_exception(e)
+                    raise
+
+    return wrapper
+
+
+def create_team_internal_async_wrapper(tracer, streaming_context_manager):
+    """Create an async wrapper for Team internal methods (_arun) that manages team span lifecycle."""
+
+    async def wrapper(wrapped, instance, args, kwargs):
+        # Get team ID for context storage
+        team_id = getattr(instance, "team_id", None) or getattr(instance, "id", None) or id(instance)
+        team_id = str(team_id)
+
+        # Get team name for span naming
+        team_name = getattr(instance, "name", None)
+        span_name = f"{team_name}.agno.team.run.workflow" if team_name else "agno.team.run.workflow"
+
+        # Check if we already have a team context (from print_response)
+        existing_context = streaming_context_manager.get_context(team_id)
+
+        if existing_context:
+            # We're being called from print_response, use existing context
+            parent_context, parent_span = existing_context
+
+            # Execute within the existing team context
+            context_token = otel_context.attach(parent_context)
+            try:
+                with tracer.start_as_current_span(span_name) as span:
+                    try:
+                        # Set workflow attributes
+                        attributes = get_team_run_attributes(args=(instance,) + args, kwargs=kwargs)
+                        for key, value in attributes.items():
+                            span.set_attribute(key, value)
+
+                        # Execute the original function
+                        result = await wrapped(*args, **kwargs)
+
+                        span.set_status(Status(StatusCode.OK))
+                        return result
+
+                    except Exception as e:
+                        span.set_status(Status(StatusCode.ERROR, str(e)))
+                        span.record_exception(e)
+                        raise
+                    finally:
+                        # Close the parent team span when workflow completes
+                        if parent_span:
+                            parent_span.end()
+                            streaming_context_manager.remove_context(team_id)
+            finally:
+                otel_context.detach(context_token)
+        else:
+            # Direct call to _arun, create new team span
+            with tracer.start_as_current_span(span_name) as span:
+                try:
+                    # Set workflow attributes
+                    attributes = get_team_run_attributes(args=(instance,) + args, kwargs=kwargs)
+                    for key, value in attributes.items():
+                        span.set_attribute(key, value)
+
+                    # Execute the original function
+                    result = await wrapped(*args, **kwargs)
+
+                    span.set_status(Status(StatusCode.OK))
+                    return result
+
+                except Exception as e:
+                    span.set_status(Status(StatusCode.ERROR, str(e)))
+                    span.record_exception(e)
+                    raise
+
+    return wrapper
+
+
+def create_team_wrapper(tracer, streaming_context_manager):
+    """Create a wrapper for Team methods that establishes the team context."""
+
+    def wrapper(wrapped, instance, args, kwargs):
+        # Get team ID for context storage
+        team_id = getattr(instance, "team_id", None) or getattr(instance, "id", None) or id(instance)
+        team_id = str(team_id)
+
+        # Check if streaming is enabled
+        is_streaming = kwargs.get("stream", getattr(instance, "stream", False))
+
+        # Get team name for span naming
+        team_name = getattr(instance, "name", None)
+        base_span_name = f"{team_name}.agno.team.run.workflow" if team_name else "agno.team.run.workflow"
+
+        # For print_response, we need to wrap the internal _run method instead
+        # because print_response returns immediately
+        if wrapped.__name__ == "print_response":
+            # Create team span but don't manage it here
+            span = tracer.start_span(base_span_name)
+
+            try:
+                # Set team attributes
+                attributes = get_team_run_attributes(args=(instance,) + args, kwargs=kwargs)
+                for key, value in attributes.items():
+                    span.set_attribute(key, value)
+
+                # Store context for child spans
+                current_context = trace.set_span_in_context(span, otel_context.get_current())
+                streaming_context_manager.store_context(team_id, current_context, span)
+
+                # The span will be closed by the internal _run method
+                # Just execute print_response normally
+                result = wrapped(*args, **kwargs)
+                return result
+
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                span.end()
+                streaming_context_manager.remove_context(team_id)
+                raise
+        else:
+            # For run/arun methods, use standard span management
+            span = tracer.start_span(base_span_name)
+
+            try:
+                # Set team attributes
+                attributes = get_team_run_attributes(args=(instance,) + args, kwargs=kwargs)
+                for key, value in attributes.items():
+                    span.set_attribute(key, value)
+
+                # Store context for child spans
+                current_context = trace.set_span_in_context(span, otel_context.get_current())
+                streaming_context_manager.store_context(team_id, current_context, span)
+
+                # Execute the original function within team context
+                context_token = otel_context.attach(current_context)
+                try:
+                    result = wrapped(*args, **kwargs)
+
+                    # For streaming results, wrap them to keep span alive
+                    if is_streaming and hasattr(result, "__iter__"):
+                        return StreamingResultWrapper(result, span, team_id, current_context, streaming_context_manager)
+                    else:
+                        # Non-streaming, close span
+                        span.end()
+                        streaming_context_manager.remove_context(team_id)
+                        return result
+
+                finally:
+                    otel_context.detach(context_token)
+
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                span.end()
+                streaming_context_manager.remove_context(team_id)
+                raise
+
+    return wrapper
+
+
+def create_team_async_wrapper(tracer, streaming_context_manager):
+    """Create an async wrapper for Team methods that establishes the team context."""
+
+    async def wrapper(wrapped, instance, args, kwargs):
+        # Get team ID for context storage
+        team_id = getattr(instance, "team_id", None) or getattr(instance, "id", None) or id(instance)
+        team_id = str(team_id)
+
+        # Check if streaming is enabled
+        is_streaming = kwargs.get("stream", getattr(instance, "stream", False))
+
+        # Get team name for span naming
+        team_name = getattr(instance, "name", None)
+        span_name = f"{team_name}.agno.team.run.workflow" if team_name else "agno.team.run.workflow"
+
+        # Create team span
+        span = tracer.start_span(span_name)
+
+        try:
+            # Set team attributes
+            attributes = get_team_run_attributes(args=(instance,) + args, kwargs=kwargs)
+            for key, value in attributes.items():
+                span.set_attribute(key, value)
+
+            # Store context for child spans - capture current context with active span
+            current_context = trace.set_span_in_context(span, otel_context.get_current())
+            streaming_context_manager.store_context(team_id, current_context, span)
+
+            # Execute the original function within team context
+            context_token = otel_context.attach(current_context)
+            try:
+                result = await wrapped(*args, **kwargs)
+
+                # For non-streaming, close the span
+                if not is_streaming:
+                    span.end()
+                    streaming_context_manager.remove_context(team_id)
+
+                return result
+            finally:
+                otel_context.detach(context_token)
+
+        except Exception as e:
+            span.set_status(Status(StatusCode.ERROR, str(e)))
+            span.record_exception(e)
+            span.end()
+            streaming_context_manager.remove_context(team_id)
+            raise
+
+    return wrapper
+
+
+def create_storage_read_wrapper(tracer, streaming_context_manager):
+    """Create a wrapper for storage read operations with cache-aware span naming."""
+
+    def wrapper(wrapped, instance, args, kwargs):
+        # Start with a basic span name
+        span_name = "agno.workflow.storage.read"
+
+        with tracer.start_as_current_span(span_name) as span:
+            try:
+                # Set flag to indicate we're in a storage operation
+                SessionStateProxy._set_storage_operation(True)
+
+                # Set initial attributes
+                attributes = get_storage_read_attributes(args=(instance,) + args, kwargs=kwargs)
+                for key, value in attributes.items():
+                    span.set_attribute(key, value)
+
+                # Execute the original function
+                result = wrapped(*args, **kwargs)
+
+                # Set result attributes including cache hit/miss
+                result_attributes = get_storage_read_attributes(
+                    args=(instance,) + args, kwargs=kwargs, return_value=result
+                )
+                for key, value in result_attributes.items():
+                    if key not in attributes:  # Avoid duplicates
+                        span.set_attribute(key, value)
+
+                # Update span name based on result and cache state
+                if hasattr(instance, "session_state") and isinstance(instance.session_state, dict):
+                    cache_size = len(instance.session_state)
+                    if result is not None:
+                        span.update_name(f"Storage.Read.Hit[cache:{cache_size}]")
+                    else:
+                        span.update_name(f"Storage.Read.Miss[cache:{cache_size}]")
+                else:
+                    # No cache info available
+                    if result is not None:
+                        span.update_name("Storage.Read.Hit")
+                    else:
+                        span.update_name("Storage.Read.Miss")
+
+                span.set_status(Status(StatusCode.OK))
+                return result
+
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                raise
+            finally:
+                # Clear the flag when done
+                SessionStateProxy._set_storage_operation(False)
+
+    return wrapper
+
+
+def create_storage_write_wrapper(tracer, streaming_context_manager):
+    """Create a wrapper for storage write operations with descriptive span naming."""
+
+    def wrapper(wrapped, instance, args, kwargs):
+        # Start with a basic span name
+        span_name = "agno.workflow.storage.write"
+
+        with tracer.start_as_current_span(span_name) as span:
+            try:
+                # Set flag to indicate we're in a storage operation
+                SessionStateProxy._set_storage_operation(True)
+
+                # Set initial attributes
+                attributes = get_storage_write_attributes(args=(instance,) + args, kwargs=kwargs)
+                for key, value in attributes.items():
+                    span.set_attribute(key, value)
+
+                # Execute the original function
+                result = wrapped(*args, **kwargs)
+
+                # Set result attributes
+                result_attributes = get_storage_write_attributes(
+                    args=(instance,) + args, kwargs=kwargs, return_value=result
+                )
+                for key, value in result_attributes.items():
+                    if key not in attributes:  # Avoid duplicates
+                        span.set_attribute(key, value)
+
+                # Update span name to show cache state after write
+                if hasattr(instance, "session_state") and isinstance(instance.session_state, dict):
+                    cache_size = len(instance.session_state)
+                    span.update_name(f"Storage.Write[cache:{cache_size}]")
+                else:
+                    span.update_name("Storage.Write")
+
+                span.set_status(Status(StatusCode.OK))
+                return result
+
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                raise
+            finally:
+                # Clear the flag when done
+                SessionStateProxy._set_storage_operation(False)
+
+    return wrapper
+
+
+class SessionStateProxy(dict):
+    """Proxy class for session_state that instruments cache operations."""
+
+    # Thread-local storage to track if we're in a storage operation
+    _thread_local = threading.local()
+
+    def __init__(self, original_dict, workflow, tracer):
+        super().__init__(original_dict)
+        self._workflow = workflow
+        self._tracer = tracer
+
+    @classmethod
+    def _in_storage_operation(cls):
+        """Check if we're currently in a storage operation."""
+        return getattr(cls._thread_local, "in_storage_operation", False)
+
+    @classmethod
+    def _set_storage_operation(cls, value):
+        """Set whether we're in a storage operation."""
+        cls._thread_local.in_storage_operation = value
+
+    def get(self, key, default=None):
+        """Instrumented get method for cache checking."""
+        # Check if we're already in a storage operation to avoid nested spans
+        if self._in_storage_operation():
+            # We're inside a storage operation, skip instrumentation
+            return super().get(key, default)
+
+        span_name = "Cache.Check"
+
+        with self._tracer.start_as_current_span(span_name) as span:
+            # Set cache attributes
+            span.set_attribute("cache.key", str(key))
+            span.set_attribute("cache.size", len(self))
+            span.set_attribute("cache.keys", json.dumps(list(self.keys())))
+
+            # Get workflow info
+            if hasattr(self._workflow, "workflow_id") and self._workflow.workflow_id:
+                span.set_attribute("cache.workflow_id", str(self._workflow.workflow_id))
+            if hasattr(self._workflow, "session_id") and self._workflow.session_id:
+                span.set_attribute("cache.session_id", str(self._workflow.session_id))
+
+            # Call the original method
+            result = super().get(key, default)
+
+            # Update span based on result
+            if result is not None and result != default:
+                span.set_attribute("cache.hit", True)
+                span.set_attribute("cache.result", "hit")
+                span.update_name(f"Cache.Hit[{len(self)} entries]")
+
+                # Add value info
+                if isinstance(result, str):
+                    span.set_attribute("cache.value_size", len(result))
+                    if len(result) <= 100:
+                        span.set_attribute("cache.value", result)
+                    else:
+                        span.set_attribute("cache.value_preview", result[:100] + "...")
+            else:
+                span.set_attribute("cache.hit", False)
+                span.set_attribute("cache.result", "miss")
+                span.update_name(f"Cache.Miss[{len(self)} entries]")
+
+            span.set_status(Status(StatusCode.OK))
+            return result
+
+    def __setitem__(self, key, value):
+        """Instrumented setitem method for cache storing."""
+        # Check if we're already in a storage operation to avoid nested spans
+        if self._in_storage_operation():
+            # We're inside a storage operation, skip instrumentation
+            return super().__setitem__(key, value)
+
+        span_name = "Cache.Store"
+
+        with self._tracer.start_as_current_span(span_name) as span:
+            # Set cache attributes
+            span.set_attribute("cache.key", str(key))
+
+            # Get workflow info
+            if hasattr(self._workflow, "workflow_id") and self._workflow.workflow_id:
+                span.set_attribute("cache.workflow_id", str(self._workflow.workflow_id))
+            if hasattr(self._workflow, "session_id") and self._workflow.session_id:
+                span.set_attribute("cache.session_id", str(self._workflow.session_id))
+
+            # Call the original method
+            super().__setitem__(key, value)
+
+            # Set post-store attributes
+            span.set_attribute("cache.size", len(self))
+            span.set_attribute("cache.keys", json.dumps(list(self.keys())))
+
+            # Add value info
+            if isinstance(value, str):
+                span.set_attribute("cache.value_size", len(value))
+                if len(value) <= 100:
+                    span.set_attribute("cache.value", value)
+                else:
+                    span.set_attribute("cache.value_preview", value[:100] + "...")
+
+            span.update_name(f"Cache.Store[{len(self)} entries]")
+            span.set_status(Status(StatusCode.OK))
+
+
+def create_workflow_init_wrapper(tracer):
+    """Wrapper to instrument workflow initialization and wrap session_state."""
+
+    def wrapper(wrapped, instance, args, kwargs):
+        # Call the original __init__
+        result = wrapped(*args, **kwargs)
+
+        # Wrap session_state if it exists
+        if hasattr(instance, "session_state") and isinstance(instance.session_state, dict):
+            # Replace session_state with our proxy
+            original_state = instance.session_state
+            instance.session_state = SessionStateProxy(original_state, instance, tracer)
+
+        return result
+
+    return wrapper
+
+
+def get_agent_context_for_llm():
+    """Helper function for LLM instrumentation to get current agent context."""
+    current_context = otel_context.get_current()
+    current_span = trace.get_current_span(current_context)
+
+    # Check if we're already in an agent span
+    if current_span and hasattr(current_span, "name") and "agent" in current_span.name:
+        return current_context, current_span
+
+    # Try to find stored agent context by checking active contexts
+    # This is a fallback for cases where context isn't properly propagated
+    return None, None
+
+
+class AgnoInstrumentor(CommonInstrumentor):
+    """Agno instrumentation class."""
+
+    def __init__(self):
+        """Initialize the Agno instrumentor."""
+        self._streaming_context_manager = StreamingContextManager()
+
+        # Create instrumentor config with populated wrapped methods
+        config = InstrumentorConfig(
+            library_name="agentops.instrumentation.agno",
+            library_version="0.1.0",
+            wrapped_methods=self._get_initial_wrapped_methods(),
+            metrics_enabled=True,
+            dependencies=["agno >= 0.1.0"],
+        )
+
+        super().__init__(config)
+
+    def _get_initial_wrapped_methods(self) -> List[WrapConfig]:
+        """Return list of methods to be wrapped during initialization."""
+        # Only return the standard wrapped methods that don't need custom wrappers
+        return WRAPPED_METHODS.copy()
+
+    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 _initialize(self, **kwargs):
+        """Perform custom initialization."""
+        logger.info("Agno instrumentation: Beginning immediate instrumentation")
+        # Perform wrapping immediately instead of with a delay
+        try:
+            self._perform_wrapping()
+            logger.info("Agno instrumentation: Immediate instrumentation completed successfully")
+        except Exception as e:
+            logger.error(f"Failed to perform immediate wrapping: {e}")
+
+    def _custom_wrap(self, **kwargs):
+        """Skip custom wrapping during initialization - it's done in _initialize."""
+        pass
+
+    def _perform_wrapping(self):
+        """Actually perform the wrapping - called after imports are complete."""
+        if not self._tracer:
+            logger.debug("No tracer available for Agno wrapping")
+            return
+
+        from agentops.instrumentation.common.wrappers import wrap_function_wrapper, WrapConfig, wrap
+
+        # Import Agno modules now that they should be fully loaded
+        try:
+            import agno.agent
+            import agno.workflow.workflow
+            import agno.tools.function
+            import agno.team.team  # Noqa: F401
+        except ImportError as e:
+            logger.error(f"Failed to import Agno modules for wrapping: {e}")
+            return
+
+        # First wrap the standard workflow session methods using the standard wrapper
+        session_methods = [
+            WrapConfig(
+                trace_name="agno.workflow.session.load_session",
+                package="agno.workflow.workflow",
+                class_name="Workflow",
+                method_name="load_session",
+                handler=get_workflow_session_attributes,
+            ),
+            WrapConfig(
+                trace_name="agno.workflow.session.new_session",
+                package="agno.workflow.workflow",
+                class_name="Workflow",
+                method_name="new_session",
+                handler=get_workflow_session_attributes,
+            ),
+            # Note: read_from_storage and write_to_storage use custom wrappers below
+        ]
+
+        wrapped_count = 0
+        for wrap_config in session_methods:
+            try:
+                wrap(wrap_config, self._tracer)
+                wrapped_count += 1
+            except Exception as e:
+                logger.debug(f"Failed to wrap {wrap_config}: {e}")
+
+        # Now wrap the streaming methods that need custom wrappers
+        streaming_methods = [
+            # Streaming agent methods
+            ("agno.agent", "Agent.run", self._create_streaming_agent_wrapper()),
+            ("agno.agent", "Agent.arun", self._create_streaming_agent_async_wrapper()),
+            # Streaming workflow methods
+            ("agno.workflow.workflow", "Workflow.run_workflow", self._create_streaming_workflow_wrapper()),
+            ("agno.workflow.workflow", "Workflow.arun_workflow", self._create_streaming_workflow_async_wrapper()),
+            # Streaming tool execution
+            ("agno.tools.function", "FunctionCall.execute", self._create_streaming_tool_wrapper()),
+            # Metrics wrapper
+            ("agno.agent", "Agent._set_session_metrics", self._create_metrics_wrapper()),
+            # Team methods - wrap all public and internal methods
+            ("agno.team.team", "Team.print_response", self._create_team_wrapper()),
+            ("agno.team.team", "Team.run", self._create_team_wrapper()),
+            ("agno.team.team", "Team.arun", self._create_team_async_wrapper()),
+            # Team internal methods with special handling
+            ("agno.team.team", "Team._run", self._create_team_internal_wrapper()),
+            ("agno.team.team", "Team._arun", self._create_team_internal_async_wrapper()),
+            # Storage methods with custom wrappers for cache-aware naming
+            ("agno.workflow.workflow", "Workflow.read_from_storage", self._create_storage_read_wrapper()),
+            ("agno.workflow.workflow", "Workflow.write_to_storage", self._create_storage_write_wrapper()),
+            # Workflow init wrapper to instrument session_state
+            ("agno.workflow.workflow", "Workflow.__init__", self._create_workflow_init_wrapper()),
+        ]
+
+        for package, method, wrapper in streaming_methods:
+            try:
+                wrap_function_wrapper(package, method, wrapper)
+                wrapped_count += 1
+            except Exception as e:
+                logger.debug(f"Failed to wrap {package}.{method}: {e}")
+
+        if wrapped_count > 0:
+            logger.info(f"Agno instrumentation: Successfully wrapped {wrapped_count} methods")
+        else:
+            logger.warning("Agno instrumentation: No methods were successfully wrapped")
+
+    def _custom_unwrap(self, **kwargs):
+        """Perform custom unwrapping."""
+        # Clear streaming contexts
+        self._streaming_context_manager.clear_all()
+        logger.info("Agno instrumentation removed successfully")
+
+    # Method wrappers converted to instance methods
+    def _create_streaming_agent_wrapper(self, args=None, kwargs=None, return_value=None):
+        """Wrapper function for streaming agent methods."""
+        return create_streaming_agent_wrapper(self._tracer, self._streaming_context_manager)
+
+    def _create_streaming_agent_async_wrapper(self, args=None, kwargs=None, return_value=None):
+        """Wrapper function for async streaming agent methods."""
+        return create_streaming_agent_async_wrapper(self._tracer, self._streaming_context_manager)
+
+    def _create_streaming_workflow_wrapper(self, args=None, kwargs=None, return_value=None):
+        """Wrapper function for streaming workflow methods."""
+        return create_streaming_workflow_wrapper(self._tracer, self._streaming_context_manager)
+
+    def _create_streaming_workflow_async_wrapper(self, args=None, kwargs=None, return_value=None):
+        """Wrapper function for async streaming workflow methods."""
+        return create_streaming_workflow_async_wrapper(self._tracer, self._streaming_context_manager)
+
+    def _create_streaming_tool_wrapper(self, args=None, kwargs=None, return_value=None):
+        """Wrapper function for streaming tool methods."""
+        return create_streaming_tool_wrapper(self._tracer, self._streaming_context_manager)
+
+    def _create_metrics_wrapper(self, args=None, kwargs=None, return_value=None):
+        """Wrapper function for metrics methods."""
+        return create_metrics_wrapper(self._tracer, self._streaming_context_manager)
+
+    def _create_team_wrapper(self, args=None, kwargs=None, return_value=None):
+        """Wrapper function for team methods."""
+        return create_team_wrapper(self._tracer, self._streaming_context_manager)
+
+    def _create_team_async_wrapper(self, args=None, kwargs=None, return_value=None):
+        """Wrapper function for async team methods."""
+        return create_team_async_wrapper(self._tracer, self._streaming_context_manager)
+
+    def _create_team_internal_wrapper(self, args=None, kwargs=None, return_value=None):
+        """Wrapper function for team internal methods."""
+        return create_team_internal_wrapper(self._tracer, self._streaming_context_manager)
+
+    def _create_team_internal_async_wrapper(self, args=None, kwargs=None, return_value=None):
+        """Wrapper function for async team internal methods."""
+        return create_team_internal_async_wrapper(self._tracer, self._streaming_context_manager)
+
+    def _create_storage_read_wrapper(self, args=None, kwargs=None, return_value=None):
+        """Wrapper function for storage read operations."""
+        return create_storage_read_wrapper(self._tracer, self._streaming_context_manager)
+
+    def _create_storage_write_wrapper(self, args=None, kwargs=None, return_value=None):
+        """Wrapper function for storage write operations."""
+        return create_storage_write_wrapper(self._tracer, self._streaming_context_manager)
+
+    def _create_workflow_init_wrapper(self, args=None, kwargs=None, return_value=None):
+        """Wrapper function for workflow initialization to instrument session_state."""
+        return create_workflow_init_wrapper(self._tracer)