django-agent-runtime 0.3.6__py3-none-any.whl

django_agent_runtime/runtime/runner.py

@@ -0,0 +1,540 @@
+"""
+Core runner for executing agent runs.
+
+Handles:
+- Claiming runs from queue
+- Executing agent runtimes
+- Heartbeats and lease management
+- Retries and error handling
+- Cancellation
+- Event emission
+"""
+
+import asyncio
+import logging
+import traceback
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Optional
+from uuid import UUID
+
+from django.conf import settings as django_settings
+
+from django_agent_runtime.conf import runtime_settings, get_event_visibility
+from django_agent_runtime.runtime.interfaces import (
+    AgentRuntime,
+    EventType,
+    Message,
+    RunContext,
+    RunResult,
+    ToolRegistry,
+    ErrorInfo,
+)
+from django_agent_runtime.runtime.registry import get_runtime
+from django_agent_runtime.runtime.queue.base import RunQueue, QueuedRun
+from django_agent_runtime.runtime.events.base import EventBus, Event
+
+logger = logging.getLogger(__name__)
+
+# Check DEBUG mode
+DEBUG = getattr(django_settings, 'DEBUG', False)
+
+
+def debug_print(msg: str):
+    """Print debug message if Django DEBUG is True."""
+    if DEBUG:
+        print(f"[agent-runner] {msg}", flush=True)
+
+
+@dataclass
+class RunContextImpl:
+    """
+    Concrete implementation of RunContext.
+
+    Provided to agent runtimes during execution.
+    """
+
+    run_id: UUID
+    conversation_id: Optional[UUID]
+    input_messages: list[Message]
+    params: dict
+    metadata: dict
+    tool_registry: ToolRegistry
+
+    # Internal state
+    _event_bus: EventBus = field(repr=False)
+    _queue: RunQueue = field(repr=False)
+    _worker_id: str = field(repr=False)
+    _seq: int = field(default=0, repr=False)
+    _state: Optional[dict] = field(default=None, repr=False)
+    _cancel_check_interval: float = field(default=1.0, repr=False)
+    _last_cancel_check: float = field(default=0.0, repr=False)
+    _is_cancelled: bool = field(default=False, repr=False)
+
+    async def emit(self, event_type: EventType | str, payload: dict) -> None:
+        """Emit an event to the event bus."""
+        event_type_str = event_type.value if isinstance(event_type, EventType) else event_type
+
+        # Get visibility for this event type
+        visibility_level, ui_visible = get_event_visibility(event_type_str)
+
+        event = Event(
+            run_id=self.run_id,
+            seq=self._seq,
+            event_type=event_type_str,
+            payload=payload,
+            timestamp=datetime.now(timezone.utc),
+            visibility_level=visibility_level,
+            ui_visible=ui_visible,
+        )
+
+        # Add detail for specific event types
+        detail = ""
+        if event_type == EventType.TOOL_CALL:
+            tool_name = payload.get("name", "unknown")
+            tool_args = str(payload.get("arguments", {}))[:80]
+            detail = f" -> {tool_name}({tool_args})"
+        elif event_type == EventType.ASSISTANT_MESSAGE:
+            content = str(payload.get("content", ""))[:80]
+            detail = f" -> {content}{'...' if len(str(payload.get('content', ''))) > 80 else ''}"
+
+        debug_print(f"Emitting event: type={event_type_str}, seq={self._seq}, visible={ui_visible}{detail}")
+        await self._event_bus.publish(event)
+        self._seq += 1
+
+    async def emit_user_message(self, content: str) -> None:
+        """
+        Emit a message that will always be shown to the user.
+
+        This is a convenience method for emitting assistant messages.
+
+        Args:
+            content: The message content to display
+        """
+        await self.emit(EventType.ASSISTANT_MESSAGE, {"content": content})
+
+    async def emit_error(self, error: str, details: dict = None) -> None:
+        """
+        Emit an error that will be shown to the user.
+
+        This is for runtime errors that should be displayed to users,
+        distinct from run.failed which is the final failure event.
+
+        Args:
+            error: The error message
+            details: Optional additional error details
+        """
+        await self.emit(EventType.ERROR, {
+            "message": error,
+            "details": details or {},
+        })
+
+    async def checkpoint(self, state: dict) -> None:
+        """Save a state checkpoint."""
+        from asgiref.sync import sync_to_async
+        from django_agent_runtime.models import AgentCheckpoint
+
+        self._state = state
+
+        @sync_to_async
+        def _save():
+            # Get next checkpoint seq
+            last = AgentCheckpoint.objects.filter(run_id=self.run_id).order_by("-seq").first()
+            next_seq = (last.seq + 1) if last else 0
+
+            AgentCheckpoint.objects.create(
+                run_id=self.run_id,
+                seq=next_seq,
+                state=state,
+            )
+
+        await _save()
+
+        # Also emit checkpoint event
+        await self.emit(EventType.STATE_CHECKPOINT, {"seq": self._seq - 1})
+
+    async def get_state(self) -> Optional[dict]:
+        """Get the last checkpointed state."""
+        if self._state is not None:
+            return self._state
+
+        from asgiref.sync import sync_to_async
+        from django_agent_runtime.models import AgentCheckpoint
+
+        @sync_to_async
+        def _get():
+            checkpoint = (
+                AgentCheckpoint.objects.filter(run_id=self.run_id)
+                .order_by("-seq")
+                .first()
+            )
+            return checkpoint.state if checkpoint else None
+
+        self._state = await _get()
+        return self._state
+
+    def cancelled(self) -> bool:
+        """Check if cancellation has been requested."""
+        return self._is_cancelled
+
+    async def check_cancelled(self) -> bool:
+        """
+        Async check for cancellation (queries database).
+
+        Call this periodically in long-running operations.
+        """
+        now = asyncio.get_event_loop().time()
+        if now - self._last_cancel_check < self._cancel_check_interval:
+            return self._is_cancelled
+
+        self._last_cancel_check = now
+
+        self._is_cancelled = await self._queue.is_cancelled(self.run_id)
+        return self._is_cancelled
+
+
+class AgentRunner:
+    """
+    Main runner for executing agent runs.
+
+    Manages the lifecycle of runs including:
+    - Claiming from queue
+    - Executing with timeout
+    - Heartbeat management
+    - Error handling and retries
+    - Cancellation
+    """
+
+    def __init__(
+        self,
+        worker_id: str,
+        queue: RunQueue,
+        event_bus: EventBus,
+        trace_sink: Optional["TraceSink"] = None,
+    ):
+        self.worker_id = worker_id
+        self.queue = queue
+        self.event_bus = event_bus
+        self.trace_sink = trace_sink
+        self.settings = runtime_settings()
+
+        self._running = False
+        self._current_runs: dict[UUID, asyncio.Task] = {}
+
+    async def run_once(self, queued_run: QueuedRun) -> None:
+        """Execute a single run."""
+        run_id = queued_run.run_id
+        agent_key = queued_run.agent_key
+
+        print(f"[agent-runner] Starting run {run_id} (agent={agent_key}, attempt={queued_run.attempt})", flush=True)
+
+        # Start tracing
+        if self.trace_sink:
+            self.trace_sink.start_run(run_id, {"agent_key": agent_key})
+
+        try:
+            # Get the runtime
+            debug_print(f"Getting runtime for agent_key={agent_key}")
+            runtime = get_runtime(agent_key)
+            debug_print(f"Got runtime: {runtime.__class__.__name__}")
+
+            # Build context
+            ctx = await self._build_context(queued_run, runtime)
+            debug_print(f"Context built: {len(ctx.input_messages)} messages")
+            for i, msg in enumerate(ctx.input_messages):
+                role = msg.get("role", "unknown")
+                content = msg.get("content", "")[:100]  # Truncate for readability
+                debug_print(f"  [{i}] {role}: {content}{'...' if len(msg.get('content', '')) > 100 else ''}")
+
+            # Emit started event
+            await ctx.emit(EventType.RUN_STARTED, {
+                "agent_key": agent_key,
+                "attempt": queued_run.attempt,
+            })
+
+            # Start heartbeat task
+            heartbeat_task = asyncio.create_task(
+                self._heartbeat_loop(run_id, ctx)
+            )
+
+            try:
+                # Execute with timeout
+                debug_print(f"Calling runtime.run() with timeout={self.settings.RUN_TIMEOUT_SECONDS}s")
+                result = await asyncio.wait_for(
+                    runtime.run(ctx),
+                    timeout=self.settings.RUN_TIMEOUT_SECONDS,
+                )
+
+                # Check for cancellation
+                if ctx.cancelled():
+                    await self._handle_cancellation(run_id, ctx)
+                    return
+
+                # Success!
+                await self._handle_success(run_id, ctx, result)
+
+            except asyncio.TimeoutError:
+                await self._handle_timeout(run_id, ctx)
+
+            except asyncio.CancelledError:
+                await self._handle_cancellation(run_id, ctx)
+
+            except Exception as e:
+                print(f"[agent-runner] Runtime error in run {run_id}: {e}", flush=True)
+                traceback.print_exc()
+                await self._handle_error(
+                    run_id, ctx, runtime, e,
+                    attempt=queued_run.attempt,
+                    max_attempts=self.settings.DEFAULT_MAX_ATTEMPTS,
+                )
+
+            finally:
+                heartbeat_task.cancel()
+                try:
+                    await heartbeat_task
+                except asyncio.CancelledError:
+                    pass
+
+        except Exception as e:
+            # Error before run started (e.g., runtime not found)
+            print(f"[agent-runner] Failed to start run {run_id}: {e}", flush=True)
+            traceback.print_exc()
+            await self.queue.release(
+                run_id,
+                self.worker_id,
+                success=False,
+                error={
+                    "type": type(e).__name__,
+                    "message": str(e),
+                    "stack": traceback.format_exc(),
+                    "retriable": False,
+                },
+            )
+
+        finally:
+            if self.trace_sink:
+                self.trace_sink.end_run(run_id, "completed")
+
+    async def _build_context(
+        self, queued_run: QueuedRun, runtime: AgentRuntime
+    ) -> RunContextImpl:
+        """Build the run context."""
+        input_data = queued_run.input
+        messages = input_data.get("messages", [])
+        params = input_data.get("params", {})
+
+        # Get conversation_id from metadata
+        conversation_id = queued_run.metadata.get("conversation_id")
+        if conversation_id:
+            conversation_id = UUID(conversation_id) if isinstance(conversation_id, str) else conversation_id
+
+        # Build tool registry (could be customized per agent)
+        tool_registry = ToolRegistry()
+
+        # Get next sequence number
+        seq = await self.event_bus.get_next_seq(queued_run.run_id)
+
+        return RunContextImpl(
+            run_id=queued_run.run_id,
+            conversation_id=conversation_id,
+            input_messages=messages,
+            params=params,
+            metadata=queued_run.metadata,
+            tool_registry=tool_registry,
+            _event_bus=self.event_bus,
+            _queue=self.queue,
+            _worker_id=self.worker_id,
+            _seq=seq,
+        )
+
+    async def _heartbeat_loop(self, run_id: UUID, ctx: RunContextImpl) -> None:
+        """Send periodic heartbeats to extend lease."""
+        while True:
+            await asyncio.sleep(self.settings.HEARTBEAT_INTERVAL_SECONDS)
+
+            # Extend lease
+            extended = await self.queue.extend_lease(
+                run_id,
+                self.worker_id,
+                self.settings.LEASE_TTL_SECONDS,
+            )
+
+            if not extended:
+                print(f"[agent-runner] Lost lease on run {run_id}", flush=True)
+                break
+
+            # Emit heartbeat event
+            await ctx.emit(EventType.RUN_HEARTBEAT, {})
+
+            # Check for cancellation
+            await ctx.check_cancelled()
+
+    async def _handle_success(
+        self, run_id: UUID, ctx: RunContextImpl, result: RunResult
+    ) -> None:
+        """Handle successful run completion."""
+        print(f"[agent-runner] Run {run_id} succeeded", flush=True)
+
+        output = {
+            "final_output": result.final_output,
+            "final_messages": result.final_messages,
+            "usage": result.usage,
+            "artifacts": result.artifacts,
+        }
+
+        # Emit success event
+        await ctx.emit(EventType.RUN_SUCCEEDED, {
+            "output": result.final_output,
+            "usage": result.usage,
+        })
+
+        # Release with success
+        await self.queue.release(
+            run_id,
+            self.worker_id,
+            success=True,
+            output=output,
+        )
+
+        # Call completion hook if configured
+        await self._call_completion_hook(run_id, output)
+
+    async def _call_completion_hook(self, run_id: UUID, output: dict) -> None:
+        """Call the configured completion hook if any."""
+        from django_agent_runtime.conf import get_hook
+
+        hook = get_hook(self.settings.RUN_COMPLETED_HOOK)
+        if not hook:
+            return
+
+        try:
+            # Run hook in thread pool since it may do sync I/O
+            from asgiref.sync import sync_to_async
+            await sync_to_async(hook)(str(run_id), output)
+        except Exception as e:
+            print(f"[agent-runner] Error in completion hook for run {run_id}: {e}", flush=True)
+
+    async def _handle_timeout(self, run_id: UUID, ctx: RunContextImpl) -> None:
+        """Handle run timeout."""
+        print(f"[agent-runner] Run {run_id} timed out after {self.settings.RUN_TIMEOUT_SECONDS}s", flush=True)
+
+        await ctx.emit(EventType.RUN_TIMED_OUT, {
+            "timeout_seconds": self.settings.RUN_TIMEOUT_SECONDS,
+        })
+
+        await self.queue.release(
+            run_id,
+            self.worker_id,
+            success=False,
+            error={
+                "type": "TimeoutError",
+                "message": f"Run exceeded {self.settings.RUN_TIMEOUT_SECONDS}s timeout",
+                "retriable": False,
+            },
+        )
+
+    async def _handle_cancellation(self, run_id: UUID, ctx: RunContextImpl) -> None:
+        """Handle run cancellation."""
+        print(f"[agent-runner] Run {run_id} cancelled", flush=True)
+
+        await ctx.emit(EventType.RUN_CANCELLED, {})
+
+        # Update status directly (not through queue.release)
+        from asgiref.sync import sync_to_async
+        from django_agent_runtime.models import AgentRun
+        from django_agent_runtime.models.base import RunStatus
+
+        @sync_to_async
+        def _update():
+            AgentRun.objects.filter(id=run_id).update(
+                status=RunStatus.CANCELLED,
+                finished_at=datetime.now(timezone.utc),
+                lease_owner="",
+                lease_expires_at=None,
+            )
+
+        await _update()
+
+    async def _handle_error(
+        self,
+        run_id: UUID,
+        ctx: RunContextImpl,
+        runtime: AgentRuntime,
+        error: Exception,
+        attempt: int = 1,
+        max_attempts: int = None,
+    ) -> None:
+        """Handle run error with retry logic."""
+        if max_attempts is None:
+            max_attempts = self.settings.DEFAULT_MAX_ATTEMPTS
+
+        print(f"[agent-runner] Run {run_id} failed (attempt {attempt}/{max_attempts}): {error}", flush=True)
+
+        # Let runtime classify the error
+        error_info = await runtime.on_error(ctx, error)
+        if error_info is None:
+            error_info = ErrorInfo(
+                type=type(error).__name__,
+                message=str(error),
+                stack=traceback.format_exc(),
+                retriable=True,
+            )
+
+        # Build comprehensive error dict for events and storage
+        error_dict = {
+            "type": error_info.type,
+            "message": error_info.message,
+            "stack": error_info.stack,
+            "retriable": error_info.retriable,
+            "details": error_info.details,
+        }
+
+        # Check if we should retry
+        can_retry = error_info.retriable and attempt < max_attempts
+
+        if can_retry:
+            # Try to requeue
+            requeued = await self.queue.requeue_for_retry(
+                run_id,
+                self.worker_id,
+                error_dict,
+                delay_seconds=self._calculate_backoff(ctx, attempt),
+            )
+
+            if requeued:
+                print(f"[agent-runner] Run {run_id} requeued for retry (attempt {attempt + 1})", flush=True)
+                # Emit an error event so UI knows about the retry
+                await ctx.emit(EventType.ERROR, {
+                    "message": f"Error occurred, retrying... (attempt {attempt}/{max_attempts})",
+                    "error": error_info.message,
+                    "error_type": error_info.type,
+                    "attempt": attempt,
+                    "max_attempts": max_attempts,
+                    "retriable": True,
+                })
+                return
+
+        # Final failure - emit detailed run.failed event
+        await ctx.emit(EventType.RUN_FAILED, {
+            "error": error_dict["message"],
+            "error_type": error_dict["type"],
+            "error_details": error_dict,
+            "attempt": attempt,
+            "max_attempts": max_attempts,
+            "retriable": False,  # No more retries
+        })
+
+        await self.queue.release(
+            run_id,
+            self.worker_id,
+            success=False,
+            error=error_dict,
+        )
+
+    def _calculate_backoff(self, ctx: RunContextImpl, attempt: int = 1) -> int:
+        """Calculate exponential backoff delay."""
+        base = self.settings.RETRY_BACKOFF_BASE
+        max_backoff = self.settings.RETRY_BACKOFF_MAX
+
+        delay = min(base ** attempt, max_backoff)
+        return int(delay)

django_agent_runtime/runtime/tracing/__init__.py

@@ -0,0 +1,48 @@
+"""
+Tracing/observability layer for agent runs.
+
+Provides:
+- TraceSink: Abstract interface (from interfaces.py)
+- NoopTraceSink: Default no-op implementation
+- LangfuseTraceSink: Langfuse integration (optional)
+"""
+
+from django_agent_runtime.runtime.interfaces import TraceSink
+from django_agent_runtime.runtime.tracing.noop import NoopTraceSink
+
+__all__ = [
+    "TraceSink",
+    "NoopTraceSink",
+    "get_trace_sink",
+]
+
+
+def get_trace_sink() -> TraceSink:
+    """
+    Factory function to get a trace sink based on settings.
+
+    Returns:
+        TraceSink instance (NoopTraceSink if tracing disabled)
+    """
+    from django_agent_runtime.conf import runtime_settings
+
+    settings = runtime_settings()
+
+    if settings.LANGFUSE_ENABLED:
+        try:
+            from django_agent_runtime.runtime.tracing.langfuse import LangfuseTraceSink
+
+            return LangfuseTraceSink(
+                public_key=settings.LANGFUSE_PUBLIC_KEY,
+                secret_key=settings.LANGFUSE_SECRET_KEY,
+                host=settings.LANGFUSE_HOST,
+            )
+        except ImportError:
+            import logging
+
+            logging.getLogger(__name__).warning(
+                "Langfuse enabled but langfuse package not installed. Using NoopTraceSink."
+            )
+
+    return NoopTraceSink()
+

django_agent_runtime/runtime/tracing/langfuse.py

@@ -0,0 +1,117 @@
+"""
+Langfuse trace sink implementation.
+
+Langfuse is an open-source LLM observability platform.
+This is an OPTIONAL integration - the core runtime doesn't depend on it.
+
+See: https://langfuse.com/
+"""
+
+import logging
+from typing import Optional
+from uuid import UUID
+
+from django_agent_runtime.runtime.interfaces import TraceSink
+
+try:
+    from langfuse import Langfuse
+except ImportError:
+    Langfuse = None
+
+logger = logging.getLogger(__name__)
+
+
+class LangfuseTraceSink(TraceSink):
+    """
+    Langfuse trace sink for LLM observability.
+
+    Sends traces to Langfuse for monitoring, debugging, and analytics.
+    """
+
+    def __init__(
+        self,
+        public_key: Optional[str] = None,
+        secret_key: Optional[str] = None,
+        host: Optional[str] = None,
+    ):
+        if Langfuse is None:
+            raise ImportError("langfuse package is required for LangfuseTraceSink")
+
+        self._client = Langfuse(
+            public_key=public_key,
+            secret_key=secret_key,
+            host=host,
+        )
+        self._traces: dict[UUID, any] = {}
+
+    def start_run(self, run_id: UUID, metadata: dict) -> None:
+        """Start a new trace in Langfuse."""
+        try:
+            trace = self._client.trace(
+                id=str(run_id),
+                name=metadata.get("agent_key", "agent_run"),
+                metadata=metadata,
+            )
+            self._traces[run_id] = trace
+        except Exception as e:
+            logger.warning(f"Failed to start Langfuse trace: {e}")
+
+    def log_event(self, run_id: UUID, event_type: str, payload: dict) -> None:
+        """Log an event to the trace."""
+        trace = self._traces.get(run_id)
+        if not trace:
+            return
+
+        try:
+            # Map event types to Langfuse concepts
+            if event_type == "assistant.message":
+                trace.generation(
+                    name="assistant_message",
+                    output=payload.get("content", ""),
+                    metadata=payload,
+                )
+            elif event_type == "tool.call":
+                trace.span(
+                    name=f"tool:{payload.get('name', 'unknown')}",
+                    input=payload.get("arguments", {}),
+                )
+            elif event_type == "tool.result":
+                # Tool results are logged as part of the span
+                pass
+            else:
+                # Generic event
+                trace.event(
+                    name=event_type,
+                    metadata=payload,
+                )
+        except Exception as e:
+            logger.warning(f"Failed to log Langfuse event: {e}")
+
+    def end_run(self, run_id: UUID, outcome: str, metadata: Optional[dict] = None) -> None:
+        """End the trace."""
+        trace = self._traces.pop(run_id, None)
+        if not trace:
+            return
+
+        try:
+            # Update trace with final status
+            status_map = {
+                "succeeded": "SUCCESS",
+                "failed": "ERROR",
+                "cancelled": "CANCELLED",
+                "timed_out": "ERROR",
+            }
+            trace.update(
+                status=status_map.get(outcome, "UNKNOWN"),
+                metadata=metadata or {},
+            )
+        except Exception as e:
+            logger.warning(f"Failed to end Langfuse trace: {e}")
+
+    def flush(self) -> None:
+        """Flush any buffered traces to Langfuse."""
+        try:
+            self._client.flush()
+        except Exception as e:
+            logger.warning(f"Failed to flush Langfuse traces: {e}")
+