digitalkin 0.3.2.dev7__py3-none-any.whl → 0.3.2.dev10__py3-none-any.whl

monitoring/digitalkin_observability/prometheus.py

@@ -0,0 +1,137 @@
+"""Prometheus metrics exporter for DigitalKin.
+
+This module exports metrics in Prometheus text exposition format.
+No external dependencies required.
+"""
+
+from __future__ import annotations
+
+from digitalkin_observability.metrics import get_metrics
+
+
+class PrometheusExporter:
+    """Exports metrics in Prometheus text format.
+
+    Usage:
+        output = PrometheusExporter.export()
+        # Returns Prometheus-compatible text format
+    """
+
+    @staticmethod
+    def export() -> str:
+        """Generate Prometheus-compatible metrics output."""
+        snapshot = get_metrics().snapshot()
+        lines: list[str] = []
+
+        # Counters
+        lines.extend([
+            "# HELP digitalkin_jobs_started_total Total jobs started",
+            "# TYPE digitalkin_jobs_started_total counter",
+            f"digitalkin_jobs_started_total {snapshot['jobs_started_total']}",
+            "",
+            "# HELP digitalkin_jobs_completed_total Total jobs completed successfully",
+            "# TYPE digitalkin_jobs_completed_total counter",
+            f"digitalkin_jobs_completed_total {snapshot['jobs_completed_total']}",
+            "",
+            "# HELP digitalkin_jobs_failed_total Total jobs failed",
+            "# TYPE digitalkin_jobs_failed_total counter",
+            f"digitalkin_jobs_failed_total {snapshot['jobs_failed_total']}",
+            "",
+            "# HELP digitalkin_jobs_cancelled_total Total jobs cancelled",
+            "# TYPE digitalkin_jobs_cancelled_total counter",
+            f"digitalkin_jobs_cancelled_total {snapshot['jobs_cancelled_total']}",
+            "",
+            "# HELP digitalkin_messages_sent_total Total messages sent",
+            "# TYPE digitalkin_messages_sent_total counter",
+            f"digitalkin_messages_sent_total {snapshot['messages_sent_total']}",
+            "",
+            "# HELP digitalkin_heartbeats_sent_total Total heartbeats sent",
+            "# TYPE digitalkin_heartbeats_sent_total counter",
+            f"digitalkin_heartbeats_sent_total {snapshot['heartbeats_sent_total']}",
+            "",
+            "# HELP digitalkin_errors_total Total errors",
+            "# TYPE digitalkin_errors_total counter",
+            f"digitalkin_errors_total {snapshot['errors_total']}",
+            "",
+        ])
+
+        # Gauges
+        lines.extend([
+            "# HELP digitalkin_active_jobs Current number of active jobs",
+            "# TYPE digitalkin_active_jobs gauge",
+            f"digitalkin_active_jobs {snapshot['active_jobs']}",
+            "",
+            "# HELP digitalkin_active_connections Current number of active connections",
+            "# TYPE digitalkin_active_connections gauge",
+            f"digitalkin_active_connections {snapshot['active_connections']}",
+            "",
+            "# HELP digitalkin_total_queue_depth Total items in all job queues",
+            "# TYPE digitalkin_total_queue_depth gauge",
+            f"digitalkin_total_queue_depth {snapshot['total_queue_depth']}",
+            "",
+        ])
+
+        # Job duration histogram
+        lines.extend(PrometheusExporter._format_histogram(
+            "digitalkin_job_duration_seconds",
+            "Job execution duration in seconds",
+            snapshot["job_duration_seconds"],
+        ))
+
+        # gRPC request duration histogram
+        lines.extend(PrometheusExporter._format_histogram(
+            "digitalkin_grpc_request_duration_seconds",
+            "gRPC request duration in seconds",
+            snapshot["grpc_request_duration_seconds"],
+        ))
+
+        # Per-module breakdown
+        if snapshot["by_module"]:
+            lines.extend([
+                "",
+                "# HELP digitalkin_jobs_by_module Jobs breakdown by module and status",
+                "# TYPE digitalkin_jobs_by_module counter",
+            ])
+            for module_name, counts in snapshot["by_module"].items():
+                for status, value in counts.items():
+                    lines.append(
+                        f'digitalkin_jobs_by_module{{module="{module_name}",status="{status}"}} {value}'
+                    )
+
+        # Per-protocol breakdown
+        if snapshot["by_protocol"]:
+            lines.extend([
+                "",
+                "# HELP digitalkin_messages_by_protocol Messages breakdown by protocol",
+                "# TYPE digitalkin_messages_by_protocol counter",
+            ])
+            for protocol, counts in snapshot["by_protocol"].items():
+                for metric, value in counts.items():
+                    lines.append(
+                        f'digitalkin_messages_by_protocol{{protocol="{protocol}",metric="{metric}"}} {value}'
+                    )
+
+        return "\n".join(lines)
+
+    @staticmethod
+    def _format_histogram(name: str, help_text: str, data: dict) -> list[str]:
+        """Format a histogram for Prometheus output."""
+        lines = [
+            "",
+            f"# HELP {name} {help_text}",
+            f"# TYPE {name} histogram",
+        ]
+
+        # Sort buckets and output cumulative counts
+        cumulative = 0
+        for bucket in sorted(data.get("buckets", {}).keys()):
+            cumulative += data["buckets"][bucket]
+            lines.append(f'{name}_bucket{{le="{bucket}"}} {cumulative}')
+
+        lines.extend([
+            f'{name}_bucket{{le="+Inf"}} {data.get("count", 0)}',
+            f'{name}_sum {data.get("sum", 0)}',
+            f'{name}_count {data.get("count", 0)}',
+        ])
+
+        return lines

monitoring/tests/test_metrics.py

@@ -0,0 +1,172 @@
+"""Tests for metrics collection.
+
+Run with: python -m pytest tests/test_metrics.py
+"""
+
+import sys
+from pathlib import Path
+
+import pytest
+
+# Add the parent directory to the path so we can import digitalkin_observability
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from digitalkin_observability import MetricsCollector, PrometheusExporter, get_metrics
+
+
+class TestMetricsCollector:
+    """Tests for MetricsCollector singleton."""
+
+    def setup_method(self) -> None:
+        """Reset metrics before each test."""
+        get_metrics().reset()
+
+    def test_singleton_returns_same_instance(self) -> None:
+        """Test that get_metrics returns the same instance."""
+        m1 = get_metrics()
+        m2 = get_metrics()
+        assert m1 is m2
+
+    def test_inc_jobs_started(self) -> None:
+        """Test incrementing jobs started counter."""
+        metrics = get_metrics()
+        metrics.inc_jobs_started("TestModule")
+
+        assert metrics.jobs_started_total == 1
+        assert metrics.active_jobs == 1
+
+    def test_inc_jobs_completed(self) -> None:
+        """Test incrementing jobs completed counter."""
+        metrics = get_metrics()
+        metrics.inc_jobs_started("TestModule")
+        metrics.inc_jobs_completed("TestModule", 1.5)
+
+        assert metrics.jobs_completed_total == 1
+        assert metrics.active_jobs == 0
+        assert metrics.job_duration_seconds.count == 1
+        assert metrics.job_duration_seconds.total_sum == 1.5
+
+    def test_inc_jobs_failed(self) -> None:
+        """Test incrementing jobs failed counter."""
+        metrics = get_metrics()
+        metrics.inc_jobs_started("TestModule")
+        metrics.inc_jobs_failed("TestModule")
+
+        assert metrics.jobs_failed_total == 1
+        assert metrics.active_jobs == 0
+
+    def test_inc_jobs_cancelled(self) -> None:
+        """Test incrementing jobs cancelled counter."""
+        metrics = get_metrics()
+        metrics.inc_jobs_started("TestModule")
+        metrics.inc_jobs_cancelled("TestModule")
+
+        assert metrics.jobs_cancelled_total == 1
+        assert metrics.active_jobs == 0
+
+    def test_inc_messages_sent(self) -> None:
+        """Test incrementing messages sent counter."""
+        metrics = get_metrics()
+        metrics.inc_messages_sent("message")
+        metrics.inc_messages_sent("file")
+        metrics.inc_messages_sent()
+
+        assert metrics.messages_sent_total == 3
+
+    def test_queue_depth_tracking(self) -> None:
+        """Test queue depth tracking."""
+        metrics = get_metrics()
+        metrics.set_queue_depth("job1", 5)
+        metrics.set_queue_depth("job2", 3)
+
+        assert metrics.queue_depth["job1"] == 5
+        assert metrics.queue_depth["job2"] == 3
+
+        metrics.clear_queue_depth("job1")
+        assert "job1" not in metrics.queue_depth
+
+    def test_snapshot(self) -> None:
+        """Test snapshot returns all metrics."""
+        metrics = get_metrics()
+        metrics.inc_jobs_started("TestModule")
+        metrics.inc_jobs_completed("TestModule", 0.5)
+        metrics.inc_messages_sent("message")
+
+        snapshot = metrics.snapshot()
+
+        assert snapshot["jobs_started_total"] == 1
+        assert snapshot["jobs_completed_total"] == 1
+        assert snapshot["messages_sent_total"] == 1
+        assert "job_duration_seconds" in snapshot
+        assert "by_module" in snapshot
+        assert "TestModule" in snapshot["by_module"]
+
+    def test_histogram_observe(self) -> None:
+        """Test histogram observations."""
+        metrics = get_metrics()
+        metrics.observe_grpc_duration(0.05)
+        metrics.observe_grpc_duration(0.15)
+
+        assert metrics.grpc_request_duration_seconds.count == 2
+        assert metrics.grpc_request_duration_seconds.total_sum == pytest.approx(0.2)
+
+    def test_reset_clears_all_metrics(self) -> None:
+        """Test reset clears all metrics."""
+        metrics = get_metrics()
+        metrics.inc_jobs_started("TestModule")
+        metrics.inc_errors()
+
+        metrics.reset()
+
+        assert metrics.jobs_started_total == 0
+        assert metrics.errors_total == 0
+        assert metrics.active_jobs == 0
+
+
+class TestPrometheusExporter:
+    """Tests for Prometheus exporter."""
+
+    def setup_method(self) -> None:
+        """Reset metrics before each test."""
+        get_metrics().reset()
+
+    def test_export_returns_string(self) -> None:
+        """Test that export returns a string."""
+        output = PrometheusExporter.export()
+        assert isinstance(output, str)
+
+    def test_export_contains_job_counters(self) -> None:
+        """Test export contains job counters."""
+        metrics = get_metrics()
+        metrics.inc_jobs_started("TestModule")
+
+        output = PrometheusExporter.export()
+
+        assert "digitalkin_jobs_started_total 1" in output
+        assert "digitalkin_active_jobs 1" in output
+
+    def test_export_contains_histogram(self) -> None:
+        """Test export contains histogram data."""
+        metrics = get_metrics()
+        metrics.observe_grpc_duration(0.05)
+
+        output = PrometheusExporter.export()
+
+        assert "digitalkin_grpc_request_duration_seconds" in output
+        assert "# TYPE digitalkin_grpc_request_duration_seconds histogram" in output
+
+    def test_export_contains_module_breakdown(self) -> None:
+        """Test export contains per-module breakdown."""
+        metrics = get_metrics()
+        metrics.inc_jobs_started("MyModule")
+
+        output = PrometheusExporter.export()
+
+        assert 'digitalkin_jobs_by_module{module="MyModule",status="started"} 1' in output
+
+    def test_export_contains_help_and_type(self) -> None:
+        """Test export contains HELP and TYPE comments."""
+        output = PrometheusExporter.export()
+
+        assert "# HELP digitalkin_jobs_started_total" in output
+        assert "# TYPE digitalkin_jobs_started_total counter" in output

digitalkin/models/module/module_helpers.py

@@ -1,189 +0,0 @@
-"""Module helpers for inter-module communication."""
-
-from collections.abc import AsyncGenerator, Callable, Coroutine
-from types import SimpleNamespace
-from typing import TYPE_CHECKING, Any
-
-from digitalkin.logger import logger
-
-if TYPE_CHECKING:
-    from digitalkin.models.module.module_context import ModuleContext
-
-
-class ModuleHelpers(SimpleNamespace):
-    """Helpers for module-to-module communication.
-
-    Extends SimpleNamespace to allow dynamic attribute assignment
-    while providing built-in helper methods.
-    """
-
-    def __init__(self, context: "ModuleContext", **kwargs: dict[str, Any]) -> None:
-        """Initialize helpers with context reference.
-
-        Args:
-            context: ModuleContext providing access to services.
-            **kwargs: Additional attributes to set on the namespace.
-        """
-        super().__init__(**kwargs)
-        self._context = context
-
-    async def call_module_by_id(
-        self,
-        module_id: str,
-        input_data: dict,
-        setup_id: str,
-        mission_id: str,
-        callback: Callable[[dict], Coroutine[Any, Any, None]] | None = None,
-    ) -> AsyncGenerator[dict, None]:
-        """Call a module by ID, discovering address/port from registry.
-
-        Args:
-            module_id: Module identifier to look up in registry
-            input_data: Input data as dictionary
-            setup_id: Setup configuration ID
-            mission_id: Mission context ID
-            callback: Optional callback for each response
-
-        Yields:
-            Streaming responses from module as dictionaries
-        """
-        module_info = self._context.registry.discover_by_id(module_id)
-
-        logger.debug(
-            "Calling module by ID",
-            extra={
-                "module_id": module_id,
-                "address": module_info.address,
-                "port": module_info.port,
-            },
-        )
-
-        async for response in self._context.communication.call_module(
-            module_address=module_info.address,
-            module_port=module_info.port,
-            input_data=input_data,
-            setup_id=setup_id,
-            mission_id=mission_id,
-            callback=callback,
-        ):
-            yield response
-
-    async def get_module_schemas_by_id(
-        self,
-        module_id: str,
-        *,
-        llm_format: bool = False,
-    ) -> dict[str, dict]:
-        """Get module schemas by ID, discovering address/port from registry.
-
-        Args:
-            module_id: Module identifier to look up in registry
-            llm_format: If True, return LLM-optimized schema format
-
-        Returns:
-            Dictionary containing schemas: {"input": ..., "output": ..., "setup": ..., "secret": ...}
-        """
-        module_info = self._context.registry.discover_by_id(module_id)
-
-        logger.debug(
-            "Getting module schemas by ID",
-            extra={
-                "module_id": module_id,
-                "address": module_info.address,
-                "port": module_info.port,
-            },
-        )
-
-        return await self._context.communication.get_module_schemas(
-            module_address=module_info.address,
-            module_port=module_info.port,
-            llm_format=llm_format,
-        )
-
-    async def create_openai_style_tool(self, module_id: str) -> dict[str, Any] | None:
-        """Create OpenAI-style function calling schema for a tool.
-
-        Uses tool cache (fast path) with registry fallback. Fetches the tool's
-        input schema and wraps it in OpenAI function calling format.
-
-        Args:
-            module_id: Module ID to look up (checks cache first, then registry)
-
-        Returns:
-            OpenAI-style tool schema if found:
-            {
-                "type": "function",
-                "function": {
-                    "name": "...",
-                    "description": "...",
-                    "parameters": {...}  # Input JSON Schema
-                }
-            }
-            None if tool not found.
-        """
-        module_info = self._context.tool_cache.check_and_get(module_id, self._context.registry)
-        if not module_info:
-            return None
-
-        schemas = await self._context.communication.get_module_schemas(
-            module_address=module_info.address,
-            module_port=module_info.port,
-            llm_format=True,
-        )
-
-        return {
-            "type": "function",
-            "function": {
-                "name": module_info.name or module_info.module_id,
-                "description": module_info.documentation or "",
-                "parameters": schemas["input"],
-            },
-        }
-
-    def create_tool_function(
-        self,
-        module_id: str,
-    ) -> Callable[..., AsyncGenerator[dict, None]] | None:
-        """Create async generator function for a tool.
-
-        Returns an async generator that calls the remote tool module via gRPC
-        and yields each response as it arrives until end_of_stream or gRPC ends.
-
-        Args:
-            module_id: Module ID to look up (checks cache first, then registry)
-
-        Returns:
-            Async generator function if tool found, None otherwise.
-            The function accepts **kwargs matching the tool's input schema
-            and yields dict responses.
-        """
-        module_info = self._context.tool_cache.check_and_get(module_id, self._context.registry)
-        if not module_info:
-            return None
-
-        # Capture references for closure
-        communication = self._context.communication
-        session = self._context.session
-        address = module_info.address
-        port = module_info.port
-
-        async def tool_function(**kwargs: Any) -> AsyncGenerator[dict, None]:  # noqa: ANN401
-            """Call remote tool module and yield responses.
-
-            Yields:
-                dict: Each response from the module until end_of_stream.
-            """
-            wrapped_input = {"root": kwargs}
-            async for response in communication.call_module(
-                module_address=address,
-                module_port=port,
-                input_data=wrapped_input,
-                setup_id=session.setup_id,
-                mission_id=session.mission_id,
-            ):
-                yield response
-
-        tool_function.__name__ = module_info.name or module_info.module_id
-        tool_function.__doc__ = module_info.documentation or ""
-
-        return tool_function