Swapnil J.

        unwanted_prefixes = [
            # 'autogen create',
            # 'autogen send',
            # 'autogen process',
            # 'autogen publish',
            'autogen ack',
            'create_agent',
            # 'autogen consume',
            # 'autogen receive',
        ]

"""
The `tracing_piece.py` module provides tracing capabilities for the chatbot application.

This module uses:
1. **OpenTelemetry**: Core tracing infrastructure.
2. **OpenInference**: Extends OpenTelemetry with AI/ML specific conventions.
3. **Arize Phoenix**: Open-source observability platform for AI applications.

The tracing uses default auto-instrumentation for OpenAI calls with custom filtering.

Key classes and methods:
- `FilteringSpanProcessor`: A custom span processor that filters out unwanted AutoGen internal spans to reduce noise in the tracing data.
- `PhoenixTracing`: The main tracing class that configures and manages the entire tracing system.
    - `__init__(self)`: Initializes the tracing instance.
    - `initialize(self)`: Configures and starts the Phoenix tracing system. It requires `PHOENIX_COLLECTOR_ENDPOINT`, `PHOENIX_PROJECT_NAME` and `PHOENIX_API_KEY` environment variables.
    - It registers with Phoenix OTEL, adds custom filtering, instruments OpenAI, and sets up a global tracer.
"""

import os

from openinference.instrumentation.autogen_agentchat import AutogenAgentChatInstrumentor
from openinference.instrumentation.openai import OpenAIInstrumentor
from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from phoenix.otel import register


class FilteringSpanProcessor(BatchSpanProcessor):
    """Custom span processor that filters out unwanted AutoGen spans."""

    def _filter_condition(self, span) -> bool:
        """Check if span should be filtered out."""
        # Filter out verbose AutoGen internal spans that don't add value
        unwanted_prefixes = [
            'autogen create',
            'autogen send',
            'autogen process',
            'autogen publish',
            'autogen ack',
            'create_agent',
            'autogen consume',
            'autogen receive',
        ]

        # Also filter out very short-lived spans that are noise
        unwanted_exact_names = ['_process_message', '_handle_response', '_format_message']

        span_name = getattr(span, 'name', '')

        return any(span_name.startswith(prefix) for prefix in unwanted_prefixes) or span_name in unwanted_exact_names

    def on_start(self, span, parent_context=None) -> None:
        """Override on_start to filter spans."""
        if self._filter_condition(span):
            return  # Skip this span
        super().on_start(span, parent_context)

    def on_end(self, span) -> None:
        """Override on_end to filter spans."""
        if self._filter_condition(span):
            return  # Skip this span
        super().on_end(span)


class PhoenixTracing:
    """Phoenix tracing configuration with auto-instrumentation for AutoGen agents."""

    def __init__(self):
        self.tracer_provider: TracerProvider | None = None
        self.tracer: trace.Tracer | None = None
        self._is_initialized = False

    def initialize(self):
        """Initialize Phoenix tracing with auto-instrumentation for AutoGen agents."""
        if self._is_initialized:
            return

        # Get environment variables
        phoenix_collector_endpoint = os.getenv('PHOENIX_COLLECTOR_ENDPOINT')
        phoenix_project_name = os.getenv('PHOENIX_PROJECT_NAME', 'mutb-manual-poc4')
        phoenix_api_key = os.getenv('PHOENIX_API_KEY')

        if not phoenix_collector_endpoint:
            return

        try:
            # Register with Phoenix OTEL with auto-instrumentation
            self.tracer_provider = register(
                project_name=phoenix_project_name, endpoint=phoenix_collector_endpoint, auto_instrument=True
            )

            # Add custom filtering processor
            if phoenix_api_key:
                # Use API key for authentication
                self.tracer_provider.add_span_processor(
                    FilteringSpanProcessor(
                        OTLPSpanExporter(
                            endpoint=phoenix_collector_endpoint, headers={'Authorization': f'Bearer {phoenix_api_key}'}
                        )
                    )
                )
            else:
                # Use without authentication
                self.tracer_provider.add_span_processor(
                    FilteringSpanProcessor(OTLPSpanExporter(endpoint=phoenix_collector_endpoint))
                )

            # # Instrument OpenAI for chat completion tracing
            # OpenAIInstrumentor().instrument(tracer_provider=self.tracer_provider)

            # Instrument AutoGen for chat completion tracing
            AutogenAgentChatInstrumentor().instrument(tracer_provider=self.tracer_provider)

            # Get the global tracer
            self.tracer = trace.get_tracer(__name__)
            self._is_initialized = True

        except Exception:
            self._is_initialized = False


# Global instance
tracing_piece = PhoenixTracing()

{
  "messages": [
    {
      "id": "e07eeae7-9fe4-4c05-adac-5b318017baa2",
      "source": "user",
      "models_usage": null,
      "metadata": {},
      "created_at": "2025-09-05T08:48:07.043629+00:00",
      "content": "User. Question",
      "type": "TextMessage"
    },
    {
      "id": "8770aaf0-0437-41f8-82ff-8cf2643d1447",
      "source": "PlanningAgent",
      "models_usage": {
        "prompt_tokens": 275,
        "completion_tokens": 68
      },
      "metadata": {},
      // ...
    }
  ],
  "stop_reason": "Text 'TERMINATE' mentioned"
}