agent-framework-devui 0.0.1a0__py3-none-any.whl → 1.0.0b251007__py3-none-any.whl

agent_framework_devui/_session.py

@@ -0,0 +1,191 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Session management for agent execution tracking."""
+
+import logging
+import uuid
+from datetime import datetime
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+# Type aliases for better readability
+SessionData = dict[str, Any]
+RequestRecord = dict[str, Any]
+SessionSummary = dict[str, Any]
+
+
+class SessionManager:
+    """Manages execution sessions for tracking requests and context."""
+
+    def __init__(self) -> None:
+        """Initialize the session manager."""
+        self.sessions: dict[str, SessionData] = {}
+
+    def create_session(self, session_id: str | None = None) -> str:
+        """Create a new execution session.
+
+        Args:
+            session_id: Optional session ID, if not provided a new one is generated
+
+        Returns:
+            Session ID
+        """
+        if not session_id:
+            session_id = str(uuid.uuid4())
+
+        self.sessions[session_id] = {
+            "id": session_id,
+            "created_at": datetime.now(),
+            "requests": [],
+            "context": {},
+            "active": True,
+        }
+
+        logger.debug(f"Created session: {session_id}")
+        return session_id
+
+    def get_session(self, session_id: str) -> SessionData | None:
+        """Get session information.
+
+        Args:
+            session_id: Session ID
+
+        Returns:
+            Session data or None if not found
+        """
+        return self.sessions.get(session_id)
+
+    def close_session(self, session_id: str) -> None:
+        """Close and cleanup a session.
+
+        Args:
+            session_id: Session ID to close
+        """
+        if session_id in self.sessions:
+            self.sessions[session_id]["active"] = False
+            logger.debug(f"Closed session: {session_id}")
+
+    def add_request_record(
+        self, session_id: str, entity_id: str, executor_name: str, request_input: Any, model_id: str
+    ) -> str:
+        """Add a request record to a session.
+
+        Args:
+            session_id: Session ID
+            entity_id: ID of the entity being executed
+            executor_name: Name of the executor
+            request_input: Input for the request
+            model_id: Model name
+
+        Returns:
+            Request ID
+        """
+        session = self.get_session(session_id)
+        if not session:
+            return ""
+
+        request_record: RequestRecord = {
+            "id": str(uuid.uuid4()),
+            "timestamp": datetime.now(),
+            "entity_id": entity_id,
+            "executor": executor_name,
+            "input": request_input,
+            "model_id": model_id,
+            "stream": True,
+        }
+        session["requests"].append(request_record)
+        return str(request_record["id"])
+
+    def update_request_record(self, session_id: str, request_id: str, updates: dict[str, Any]) -> None:
+        """Update a request record in a session.
+
+        Args:
+            session_id: Session ID
+            request_id: Request ID to update
+            updates: Dictionary of updates to apply
+        """
+        session = self.get_session(session_id)
+        if not session:
+            return
+
+        for request in session["requests"]:
+            if request["id"] == request_id:
+                request.update(updates)
+                break
+
+    def get_session_history(self, session_id: str) -> SessionSummary | None:
+        """Get session execution history.
+
+        Args:
+            session_id: Session ID
+
+        Returns:
+            Session history or None if not found
+        """
+        session = self.get_session(session_id)
+        if not session:
+            return None
+
+        return {
+            "session_id": session_id,
+            "created_at": session["created_at"].isoformat(),
+            "active": session["active"],
+            "request_count": len(session["requests"]),
+            "requests": [
+                {
+                    "id": req["id"],
+                    "timestamp": req["timestamp"].isoformat(),
+                    "entity_id": req["entity_id"],
+                    "executor": req["executor"],
+                    "model": req["model"],
+                    "input_length": len(str(req["input"])) if req["input"] else 0,
+                    "execution_time": req.get("execution_time"),
+                    "status": req.get("status", "unknown"),
+                }
+                for req in session["requests"]
+            ],
+        }
+
+    def get_active_sessions(self) -> list[SessionSummary]:
+        """Get list of active sessions.
+
+        Returns:
+            List of active session summaries
+        """
+        active_sessions = []
+
+        for session_id, session in self.sessions.items():
+            if session["active"]:
+                active_sessions.append({
+                    "session_id": session_id,
+                    "created_at": session["created_at"].isoformat(),
+                    "request_count": len(session["requests"]),
+                    "last_activity": (
+                        session["requests"][-1]["timestamp"].isoformat()
+                        if session["requests"]
+                        else session["created_at"].isoformat()
+                    ),
+                })
+
+        return active_sessions
+
+    def cleanup_old_sessions(self, max_age_hours: int = 24) -> None:
+        """Cleanup old sessions to prevent memory leaks.
+
+        Args:
+            max_age_hours: Maximum age of sessions to keep in hours
+        """
+        cutoff_time = datetime.now().timestamp() - (max_age_hours * 3600)
+
+        sessions_to_remove = []
+        for session_id, session in self.sessions.items():
+            if session["created_at"].timestamp() < cutoff_time:
+                sessions_to_remove.append(session_id)
+
+        for session_id in sessions_to_remove:
+            del self.sessions[session_id]
+            logger.debug(f"Cleaned up old session: {session_id}")
+
+        if sessions_to_remove:
+            logger.info(f"Cleaned up {len(sessions_to_remove)} old sessions")

agent_framework_devui/_tracing.py

@@ -0,0 +1,168 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Simplified tracing integration for Agent Framework Server."""
+
+import logging
+from collections.abc import Generator, Sequence
+from contextlib import contextmanager
+from datetime import datetime
+from typing import Any
+
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+
+from .models import ResponseTraceEvent
+
+logger = logging.getLogger(__name__)
+
+
+class SimpleTraceCollector(SpanExporter):
+    """Simple trace collector that captures spans for direct yielding."""
+
+    def __init__(self, session_id: str | None = None, entity_id: str | None = None) -> None:
+        """Initialize trace collector.
+
+        Args:
+            session_id: Session identifier for context
+            entity_id: Entity identifier for context
+        """
+        self.session_id = session_id
+        self.entity_id = entity_id
+        self.collected_events: list[ResponseTraceEvent] = []
+
+    def export(self, spans: Sequence[Any]) -> SpanExportResult:
+        """Collect spans as trace events.
+
+        Args:
+            spans: Sequence of OpenTelemetry spans
+
+        Returns:
+            SpanExportResult indicating success
+        """
+        logger.debug(f"SimpleTraceCollector received {len(spans)} spans")
+
+        try:
+            for span in spans:
+                trace_event = self._convert_span_to_trace_event(span)
+                if trace_event:
+                    self.collected_events.append(trace_event)
+                    logger.debug(f"Collected trace event: {span.name}")
+
+            return SpanExportResult.SUCCESS
+
+        except Exception as e:
+            logger.error(f"Error collecting trace spans: {e}")
+            return SpanExportResult.FAILURE
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        """Force flush spans (no-op for simple collection)."""
+        return True
+
+    def get_pending_events(self) -> list[ResponseTraceEvent]:
+        """Get and clear pending trace events.
+
+        Returns:
+            List of collected trace events, clearing the internal list
+        """
+        events = self.collected_events.copy()
+        self.collected_events.clear()
+        return events
+
+    def _convert_span_to_trace_event(self, span: Any) -> ResponseTraceEvent | None:
+        """Convert OpenTelemetry span to ResponseTraceEvent.
+
+        Args:
+            span: OpenTelemetry span
+
+        Returns:
+            ResponseTraceEvent or None if conversion fails
+        """
+        try:
+            start_time = span.start_time / 1_000_000_000  # Convert from nanoseconds
+            end_time = span.end_time / 1_000_000_000 if span.end_time else None
+            duration_ms = ((end_time - start_time) * 1000) if end_time else None
+
+            # Build trace data
+            trace_data = {
+                "type": "trace_span",
+                "span_id": str(span.context.span_id),
+                "trace_id": str(span.context.trace_id),
+                "parent_span_id": str(span.parent.span_id) if span.parent else None,
+                "operation_name": span.name,
+                "start_time": start_time,
+                "end_time": end_time,
+                "duration_ms": duration_ms,
+                "attributes": dict(span.attributes) if span.attributes else {},
+                "status": str(span.status.status_code) if hasattr(span, "status") else "OK",
+                "session_id": self.session_id,
+                "entity_id": self.entity_id,
+            }
+
+            # Add events if available
+            if hasattr(span, "events") and span.events:
+                trace_data["events"] = [
+                    {
+                        "name": event.name,
+                        "timestamp": event.timestamp / 1_000_000_000,
+                        "attributes": dict(event.attributes) if event.attributes else {},
+                    }
+                    for event in span.events
+                ]
+
+            # Add error information if span failed
+            if hasattr(span, "status") and span.status.status_code.name == "ERROR":
+                trace_data["error"] = span.status.description or "Unknown error"
+
+            return ResponseTraceEvent(type="trace_event", data=trace_data, timestamp=datetime.now().isoformat())
+
+        except Exception as e:
+            logger.warning(f"Failed to convert span {getattr(span, 'name', 'unknown')}: {e}")
+            return None
+
+
+@contextmanager
+def capture_traces(
+    session_id: str | None = None, entity_id: str | None = None
+) -> Generator[SimpleTraceCollector, None, None]:
+    """Context manager to capture traces during execution.
+
+    Args:
+        session_id: Session identifier for context
+        entity_id: Entity identifier for context
+
+    Yields:
+        SimpleTraceCollector instance to get trace events from
+    """
+    collector = SimpleTraceCollector(session_id, entity_id)
+
+    try:
+        from opentelemetry import trace
+        from opentelemetry.sdk.trace import TracerProvider
+        from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
+        # Get current tracer provider and add our collector
+        provider = trace.get_tracer_provider()
+        processor = SimpleSpanProcessor(collector)
+
+        # Check if this is a real TracerProvider (not the default NoOpTracerProvider)
+        if isinstance(provider, TracerProvider):
+            provider.add_span_processor(processor)
+            logger.debug(f"Added trace collector to TracerProvider for session: {session_id}, entity: {entity_id}")
+
+            try:
+                yield collector
+            finally:
+                # Clean up - shutdown processor
+                try:
+                    processor.shutdown()
+                except Exception as e:
+                    logger.debug(f"Error shutting down processor: {e}")
+        else:
+            logger.warning(f"No real TracerProvider available, got: {type(provider)}")
+            yield collector
+
+    except ImportError:
+        logger.debug("OpenTelemetry not available")
+        yield collector
+    except Exception as e:
+        logger.error(f"Error setting up trace capture: {e}")
+        yield collector