penguiflow 2.2.5__py3-none-any.whl → 2.2.6__py3-none-any.whl

examples/planner_enterprise_agent/telemetry.py

@@ -0,0 +1,245 @@
+"""Enterprise telemetry middleware for comprehensive observability.
+
+This module implements the telemetry patterns from the successful PenguiFlow
+implementation case study, providing full visibility into:
+- Planner lifecycle events
+- Node execution with detailed error payloads
+- Flow events with structured logging
+- MLflow/observability backend integration
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping
+from typing import Any
+
+from examples.planner_enterprise_agent.config import AgentConfig
+from penguiflow.metrics import FlowEvent
+from penguiflow.planner import PlannerEvent
+
+
+class AgentTelemetry:
+    """Comprehensive telemetry for enterprise agent deployments.
+
+    Implements the telemetry middleware pattern that captures:
+    1. Full exception tracebacks from FlowEvents
+    2. Detailed error payloads with context
+    3. LLM call costs and latency
+    4. Planning step metrics
+    5. Structured events for external systems
+
+    Usage:
+        telemetry = AgentTelemetry(config)
+
+        # Add to PenguiFlow
+        flow.add_middleware(log_flow_events(telemetry.logger))
+        flow.add_middleware(telemetry.record_flow_event)
+
+        # Add to ReactPlanner
+        planner = ReactPlanner(..., event_callback=telemetry.record_planner_event)
+    """
+
+    def __init__(self, config: AgentConfig) -> None:
+        self.config = config
+        self.logger = logging.getLogger(f"penguiflow.{config.agent_name}")
+        self.planner_logger = logging.getLogger(
+            f"penguiflow.{config.agent_name}.planner"
+        )
+
+        # Event collection for batch emission
+        self._events: list[dict[str, Any]] = []
+
+        # Metrics tracking
+        self._metrics: dict[str, Any] = {
+            "planner_steps": 0,
+            "planner_llm_calls": 0,
+            "planner_cost_usd": 0.0,
+            "flow_node_errors": 0,
+            "flow_node_successes": 0,
+        }
+
+    async def record_flow_event(self, event: FlowEvent) -> FlowEvent:
+        """Middleware function that intercepts all PenguiFlow events.
+
+        This is the CRITICAL pattern from the case study - it extracts
+        detailed error payloads that would otherwise be trapped in flow state.
+        """
+        event_type = event.event_type
+
+        if event_type == "node_start":
+            self.logger.debug(
+                "node_start",
+                extra={
+                    "node": event.node_name,
+                    "trace_id": event.trace_id,
+                    "node_id": event.node_id,
+                },
+            )
+
+        elif event_type == "node_success":
+            self._metrics["flow_node_successes"] += 1
+            self.logger.info(
+                "node_success",
+                extra={
+                    "node": event.node_name,
+                    "trace_id": event.trace_id,
+                    "latency_ms": event.latency_ms,
+                },
+            )
+
+        elif event_type == "node_error":
+            # THIS IS THE CRITICAL PART - Extract error details!
+            # Without this, you only see "node_error" with no context
+            error_payload = event.error_payload or {}
+
+            self._metrics["flow_node_errors"] += 1
+
+            # Log everything for debugging (the breakthrough from the case study)
+            self.logger.error(
+                "node_error",
+                extra={
+                    "node": event.node_name,
+                    "trace_id": event.trace_id,
+                    "node_id": event.node_id,
+                    "error_class": error_payload.get("error_class"),
+                    "error_message": error_payload.get("error_message"),
+                    "error_traceback": error_payload.get("error_traceback"),
+                    "flow_error_code": error_payload.get("code"),
+                    "flow_error_message": error_payload.get("message"),
+                    # Include full payload for complete visibility
+                    **error_payload,
+                },
+            )
+
+            # Collect for batch emission to observability backend
+            if self.config.enable_telemetry:
+                self._events.append(
+                    {
+                        "event": "flow.node_error",
+                        "payload": {
+                            "node": event.node_name,
+                            "trace_id": event.trace_id,
+                            **error_payload,
+                        },
+                    }
+                )
+
+        # Always return event unmodified - middleware is read-only
+        return event
+
+    def record_planner_event(self, event: PlannerEvent) -> None:
+        """Callback for ReactPlanner events.
+
+        Captures planner-specific telemetry:
+        - Step start/complete with latency
+        - LLM calls with cost tracking
+        - Pause/resume operations
+        - Constraint violations
+        """
+        event_type = event.event_type
+
+        # Extract all event data
+        extra = event.to_payload()
+
+        if event_type == "step_start":
+            self.planner_logger.debug("step_start", extra=extra)
+
+        elif event_type == "step_complete":
+            self._metrics["planner_steps"] += 1
+            self.planner_logger.info("step_complete", extra=extra)
+
+            # Track cost if available
+            cost = extra.get("cost_usd", 0)
+            if cost > 0:
+                self._metrics["planner_cost_usd"] += cost
+
+        elif event_type == "llm_call":
+            self._metrics["planner_llm_calls"] += 1
+            self.planner_logger.debug("llm_call", extra=extra)
+
+        elif event_type == "pause":
+            self.planner_logger.info("pause", extra=extra)
+
+        elif event_type == "resume":
+            self.planner_logger.info("resume", extra=extra)
+
+        elif event_type == "finish":
+            self.planner_logger.info("finish", extra=extra)
+
+        elif event_type.endswith("_error") or "error" in event.extra:
+            self.planner_logger.error(event_type, extra=extra)
+
+        # Collect for observability backend
+        if self.config.enable_telemetry:
+            self._events.append(
+                {
+                    "event": f"planner.{event_type}",
+                    "payload": extra,
+                }
+            )
+
+    def emit_collected_events(self) -> None:
+        """Emit batched events to observability backend.
+
+        Call this after planner execution completes to send all
+        collected telemetry to your monitoring system.
+        """
+        if not self._events:
+            return
+
+        if self.config.telemetry_backend == "mlflow":
+            self._emit_to_mlflow()
+        elif self.config.telemetry_backend == "datadog":
+            self._emit_to_datadog()
+        else:
+            # Default: log as JSON for structured log aggregation
+            self.logger.info(
+                "telemetry_batch",
+                extra={
+                    "events": self._events,
+                    "metrics": self._metrics,
+                },
+            )
+
+        # Clear for next execution
+        self._events.clear()
+
+    def _emit_to_mlflow(self) -> None:
+        """Emit events to MLflow tracking server."""
+        if not self.config.mlflow_tracking_uri:
+            self.logger.warning("mlflow_tracking_uri not configured, skipping emission")
+            return
+
+        # Implementation would use mlflow.log_metrics, mlflow.log_params, etc.
+        # Stub for example purposes
+        self.logger.info(
+            "mlflow_emit",
+            extra={
+                "tracking_uri": self.config.mlflow_tracking_uri,
+                "event_count": len(self._events),
+            },
+        )
+
+    def _emit_to_datadog(self) -> None:
+        """Emit events to DataDog APM."""
+        # Implementation would use datadog client
+        # Stub for example purposes
+        self.logger.info(
+            "datadog_emit",
+            extra={"event_count": len(self._events)},
+        )
+
+    def get_metrics(self) -> Mapping[str, Any]:
+        """Return current metrics snapshot."""
+        return dict(self._metrics)
+
+    def reset_metrics(self) -> None:
+        """Reset metrics counters (useful for testing)."""
+        self._metrics = {
+            "planner_steps": 0,
+            "planner_llm_calls": 0,
+            "planner_cost_usd": 0.0,
+            "flow_node_errors": 0,
+            "flow_node_successes": 0,
+        }

penguiflow/__init__.py

@@ -105,4 +105,4 @@ __all__ = [
     "TrajectoryStep",
 ]
 
-__version__ = "2.2.5"
+__version__ = "2.2.6"

penguiflow/planner/__init__.py

@@ -2,10 +2,13 @@
 
 from __future__ import annotations
 
+from .dspy_client import DSPyLLMClient
 from .react import (
     ParallelCall,
     ParallelJoin,
     PlannerAction,
+    PlannerEvent,
+    PlannerEventCallback,
     PlannerFinish,
     PlannerPause,
     ReactPlanner,
@@ -15,9 +18,12 @@ from .react import (
 )
 
 __all__ = [
+    "DSPyLLMClient",
     "ParallelCall",
     "ParallelJoin",
     "PlannerAction",
+    "PlannerEvent",
+    "PlannerEventCallback",
     "PlannerFinish",
     "PlannerPause",
     "ReactPlanner",

penguiflow/planner/dspy_client.py

@@ -0,0 +1,327 @@
+"""DSPy-based LLM client for ReactPlanner with robust structured outputs.
+
+This module provides a DSPy-powered alternative to direct LiteLLM calls,
+offering better structured output handling across different LLM providers.
+DSPy's signature system with Pydantic models works reliably even with
+providers that don't support native JSON schema mode (like Databricks).
+"""
+
+from __future__ import annotations
+
+import ast
+import asyncio
+import json
+import logging
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel
+
+if TYPE_CHECKING:
+    # PlannerAction imported at runtime in _create_signature to avoid circular import
+    from penguiflow.planner.react import PlannerAction  # noqa: F401
+
+logger = logging.getLogger(__name__)
+
+
+class DSPyLLMClient:
+    """LLM client using DSPy for structured outputs.
+
+    This client implements the JSONLLMClient protocol and uses DSPy's
+    signature system to generate structured outputs. DSPy handles the
+    prompt engineering and parsing internally, providing more reliable
+    structured outputs across different LLM providers.
+
+    Benefits over direct LiteLLM:
+    - Better structured output reliability across providers
+    - Automatic prompt optimization for structure extraction
+    - Works with models that don't support native JSON schema mode
+    - Graceful degradation with retry logic
+
+    Args:
+        llm: Model identifier (e.g., "gpt-4o-mini",
+            "databricks/databricks-gpt-oss-120b")
+        temperature: Sampling temperature (0.0 = deterministic)
+        max_retries: Number of retry attempts for transient failures
+        timeout_s: Timeout per LLM call in seconds
+
+    Example:
+        >>> client = DSPyLLMClient(
+        ...     llm="databricks/databricks-gpt-oss-120b",
+        ...     temperature=0.0,
+        ... )
+        response = await client.complete(
+            messages=[{"role": "user", "content": "..."}],
+            response_format={"type": "json_schema", "json_schema": {...}},
+        )
+    """
+
+    expects_json_schema = True
+
+    def __init__(
+        self,
+        llm: str | dict[str, Any],
+        *,
+        temperature: float = 0.0,
+        max_retries: int = 3,
+        timeout_s: float = 60.0,
+    ) -> None:
+        self._llm = llm
+        self._temperature = temperature
+        self._max_retries = max_retries
+        self._timeout_s = timeout_s
+        self._dspy_module: Any = None
+        self._lm: Any = None
+
+    def _ensure_dspy_initialized(self) -> None:
+        """Lazy-initialize DSPy to avoid import overhead."""
+        if self._lm is not None:
+            return
+
+        try:
+            import dspy
+        except ModuleNotFoundError as exc:  # pragma: no cover
+            raise RuntimeError(
+                "DSPy is not installed. Install penguiflow[planner] or provide "
+                "a custom llm_client."
+            ) from exc
+
+        # Configure DSPy LM
+        model_name = self._llm if isinstance(self._llm, str) else self._llm["model"]
+
+        # DSPy uses LiteLLM under the hood, so all LiteLLM model names work
+        self._lm = dspy.LM(
+            model=model_name,
+            temperature=self._temperature,
+            max_tokens=4096,
+        )
+
+        logger.info(
+            "dspy_lm_initialized",
+            extra={"model": model_name, "temperature": self._temperature},
+        )
+
+    def _create_signature(self, response_format: Mapping[str, Any] | None) -> type[Any]:
+        """Create a DSPy signature for PlannerAction output.
+
+        Args:
+            response_format: OpenAI-style response_format (used to detect schema mode)
+
+        Returns:
+            DSPy Signature class with PlannerAction as output type
+        """
+        import dspy
+
+        # Import at runtime to avoid circular dependency
+        from penguiflow.planner.react import PlannerAction
+
+        if not response_format or "json_schema" not in response_format:
+            # Fallback: simple string output for non-schema requests
+            attrs = {
+                "__doc__": "Generate a response.",
+                "__annotations__": {"messages": str, "response": str},
+                "messages": dspy.InputField(),
+                "response": dspy.OutputField(),
+            }
+            return type("TextOutputSignature", (dspy.Signature,), attrs)
+
+        # Use PlannerAction directly - no schema conversion needed!
+        # DSPy will handle the Pydantic model → JSON → Pydantic validation
+        attrs = {
+            "__doc__": "Generate a structured planner action with proper type safety.",
+            "__annotations__": {"messages": str, "response": PlannerAction},
+            "messages": dspy.InputField(
+                desc="Conversation history and user query requiring a planner action"
+            ),
+            "response": dspy.OutputField(
+                desc=(
+                    "Structured planner action with thought, next_node, "
+                    "args, plan, or join"
+                )
+            ),
+        }
+        return type("PlannerActionSignature", (dspy.Signature,), attrs)
+
+    def _messages_to_text(self, messages: Sequence[Mapping[str, str]]) -> str:
+        """Convert OpenAI-style messages to a single text prompt.
+
+        Args:
+            messages: List of message dicts with role and content
+
+        Returns:
+            Concatenated text suitable for DSPy input
+        """
+        parts = []
+        for msg in messages:
+            role = msg.get("role", "user")
+            content = msg.get("content", "")
+            if role == "system":
+                parts.append(f"System: {content}")
+            elif role == "user":
+                parts.append(f"User: {content}")
+            elif role == "assistant":
+                parts.append(f"Assistant: {content}")
+            else:
+                parts.append(content)
+        return "\n\n".join(parts)
+
+    async def complete(
+        self,
+        *,
+        messages: Sequence[Mapping[str, str]],
+        response_format: Mapping[str, Any] | None = None,
+    ) -> str:
+        """Generate completion with structured output via DSPy.
+
+        Args:
+            messages: OpenAI-style message list
+            response_format: Optional JSON schema for structured output
+
+        Returns:
+            JSON string containing the structured response
+
+        Raises:
+            RuntimeError: If all retry attempts fail
+            TimeoutError: If the call exceeds timeout_s
+        """
+        import dspy
+
+        self._ensure_dspy_initialized()
+
+        # Create signature based on response format
+        signature_class = self._create_signature(response_format)
+
+        # Create DSPy predictor
+        predictor = dspy.Predict(signature_class)
+
+        # Convert messages to text
+        input_text = self._messages_to_text(messages)
+
+        last_error: Exception | None = None
+        for attempt in range(self._max_retries):
+            try:
+                async with asyncio.timeout(self._timeout_s):
+                    # DSPy doesn't have native async support yet, so we run in executor
+                    loop = asyncio.get_running_loop()
+
+                    def _run_dspy() -> Any:
+                        with dspy.context(lm=self._lm):
+                            return predictor(messages=input_text)
+
+                    result = await loop.run_in_executor(None, _run_dspy)
+
+                    # Extract response
+                    if hasattr(result, "response"):
+                        response_obj = result.response
+                        logger.debug(
+                            "dspy_response_extracted",
+                            extra={
+                                "response_type": type(response_obj).__name__,
+                                "response_preview": str(response_obj)[:200],
+                            },
+                        )
+                        if isinstance(response_obj, BaseModel):
+                            # PlannerAction or other Pydantic model - already validated!
+                            json_output = response_obj.model_dump_json()
+                            logger.debug(
+                                "dspy_pydantic_success",
+                                extra={
+                                    "model": type(response_obj).__name__,
+                                    "json_length": len(json_output),
+                                },
+                            )
+                            return json_output
+                        elif isinstance(response_obj, dict):
+                            return json.dumps(response_obj)
+                        else:
+                            # DSPy sometimes returns string - normalise to JSON
+                            response_str = str(response_obj)
+                            logger.debug(
+                                "dspy_string_response",
+                                extra={"response_preview": response_str[:500]},
+                            )
+                            normalised = self._normalise_json(response_str)
+                            if normalised is not None:
+                                logger.debug("dspy_json_normalised_success")
+                                return normalised
+                            logger.warning(
+                                "dspy_invalid_json",
+                                extra={"response": response_str[:500]},
+                            )
+                            raise RuntimeError(
+                                "DSPy returned output that could not be coerced to JSON"
+                            )
+                    else:
+                        raise RuntimeError("DSPy returned no response field")
+
+            except TimeoutError as exc:
+                last_error = exc
+                logger.warning(
+                    "dspy_timeout",
+                    extra={"attempt": attempt + 1, "timeout_s": self._timeout_s},
+                )
+            except Exception as exc:
+                last_error = exc
+                logger.warning(
+                    "dspy_error",
+                    extra={
+                        "attempt": attempt + 1,
+                        "error": str(exc),
+                        "error_type": type(exc).__name__,
+                    },
+                )
+
+            # Exponential backoff
+            if attempt < self._max_retries - 1:
+                await asyncio.sleep(2**attempt)
+
+        # All retries exhausted
+        error_msg = f"DSPy LLM call failed after {self._max_retries} attempts"
+        if last_error:
+            error_msg += f": {last_error}"
+        raise RuntimeError(error_msg)
+
+    def _normalise_json(self, text: str) -> str | None:
+        """Attempt to coerce arbitrary text into canonical JSON string."""
+        candidate = text.strip()
+        if not candidate:
+            return None
+
+        # Remove code fences if present
+        if candidate.startswith("```"):
+            parts = candidate.split("```")
+            candidate = ""
+            for part in parts:
+                stripped = part.strip()
+                if stripped.lower().startswith("json"):
+                    stripped = stripped[4:].strip()
+                if stripped:
+                    candidate = stripped
+                    break
+            if not candidate:
+                candidate = text.strip("` \n")
+
+        # Extract substring bounded by braces if extra commentary exists
+        if candidate.count("{") >= 1 and candidate.count("}") >= 1:
+            start = candidate.find("{")
+            end = candidate.rfind("}")
+            if start != -1 and end != -1 and end > start:
+                candidate = candidate[start : end + 1]
+
+        # First try strict JSON
+        try:
+            payload = json.loads(candidate)
+        except json.JSONDecodeError:
+            # Try python literal eval fallback (handles single quotes, trailing commas)
+            try:
+                payload = ast.literal_eval(candidate)
+            except Exception:
+                return None
+
+        # Ensure payload is JSON-serialisable dict
+        if isinstance(payload, (str, int, float, bool)) or payload is None:
+            return json.dumps(payload)
+        try:
+            return json.dumps(payload)
+        except TypeError:
+            return None