flock-core 0.4.542__py3-none-any.whl → 0.5.0__py3-none-any.whl

flock/logging/telemetry_exporter/duckdb_exporter.py

@@ -0,0 +1,216 @@
+"""DuckDB exporter for OpenTelemetry spans - optimized for analytical queries."""
+
+import json
+from pathlib import Path
+
+import duckdb
+from opentelemetry.sdk.trace.export import SpanExportResult
+from opentelemetry.trace import Status, StatusCode
+
+from flock.logging.telemetry_exporter.base_exporter import TelemetryExporter
+
+
+class DuckDBSpanExporter(TelemetryExporter):
+    """Export spans to DuckDB for fast analytical queries.
+
+    DuckDB is a columnar analytical database optimized for OLAP workloads,
+    making it 10-100x faster than SQLite for trace analytics like:
+    - Aggregations (avg/p95/p99 duration)
+    - Time-range queries
+    - Service/operation filtering
+    - Complex analytical queries
+
+    The database is a single file with zero configuration required.
+    """
+
+    def __init__(self, dir: str, db_name: str = "traces.duckdb", ttl_days: int | None = None):
+        """Initialize the DuckDB exporter.
+
+        Args:
+            dir: Directory where the database file will be created
+            db_name: Name of the DuckDB file (default: traces.duckdb)
+            ttl_days: Delete traces older than this many days (default: None = keep forever)
+        """
+        super().__init__()
+        self.telemetry_path = Path(dir)
+        self.telemetry_path.mkdir(parents=True, exist_ok=True)
+        self.db_path = self.telemetry_path / db_name
+        self.ttl_days = ttl_days
+
+        # Initialize database and create schema
+        self._init_database()
+
+    def _init_database(self):
+        """Create the spans table if it doesn't exist."""
+        with duckdb.connect(str(self.db_path)) as conn:
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS spans (
+                    trace_id VARCHAR NOT NULL,
+                    span_id VARCHAR PRIMARY KEY,
+                    parent_id VARCHAR,
+                    name VARCHAR NOT NULL,
+                    service VARCHAR,
+                    operation VARCHAR,
+                    kind VARCHAR,
+                    start_time BIGINT NOT NULL,
+                    end_time BIGINT NOT NULL,
+                    duration_ms DOUBLE NOT NULL,
+                    status_code VARCHAR NOT NULL,
+                    status_description VARCHAR,
+                    attributes JSON,
+                    events JSON,
+                    links JSON,
+                    resource JSON,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            """)
+
+            # Create indexes for common query patterns
+            conn.execute("CREATE INDEX IF NOT EXISTS idx_trace_id ON spans(trace_id)")
+            conn.execute("CREATE INDEX IF NOT EXISTS idx_service ON spans(service)")
+            conn.execute("CREATE INDEX IF NOT EXISTS idx_start_time ON spans(start_time)")
+            conn.execute("CREATE INDEX IF NOT EXISTS idx_name ON spans(name)")
+            conn.execute("CREATE INDEX IF NOT EXISTS idx_created_at ON spans(created_at)")
+
+        # Cleanup old traces if TTL is configured
+        if self.ttl_days is not None:
+            self._cleanup_old_traces()
+
+    def _cleanup_old_traces(self):
+        """Delete traces older than TTL_DAYS.
+
+        This runs on exporter initialization to keep the database size manageable.
+        """
+        if self.ttl_days is None:
+            return
+
+        try:
+            with duckdb.connect(str(self.db_path)) as conn:
+                # Delete spans older than TTL
+                # Note: DuckDB doesn't support ? placeholders inside INTERVAL expressions
+                # Safe: ttl_days is guaranteed to be an int, no injection risk
+                query = f"""
+                    DELETE FROM spans
+                    WHERE created_at < CURRENT_TIMESTAMP - INTERVAL '{self.ttl_days} DAYS'
+                """  # nosec B608
+                result = conn.execute(query)
+
+                deleted_count = result.fetchall()[0][0] if result else 0
+
+                if deleted_count > 0:
+                    print(
+                        f"[DuckDB TTL] Deleted {deleted_count} spans older than {self.ttl_days} days"
+                    )
+
+        except Exception as e:
+            print(f"[DuckDB TTL] Error cleaning up old traces: {e}")
+
+    def _span_to_record(self, span):
+        """Convert a ReadableSpan to a database record."""
+        context = span.get_span_context()
+        status = span.status or Status(StatusCode.UNSET)
+
+        # Extract service and operation from span name
+        # Format: "ServiceName.operation_name"
+        parts = span.name.split(".", 1)
+        service = parts[0] if len(parts) > 0 else "unknown"
+        operation = parts[1] if len(parts) > 1 else span.name
+
+        # Calculate duration in milliseconds
+        duration_ms = (span.end_time - span.start_time) / 1_000_000
+
+        # Serialize complex fields to JSON
+        attributes_json = json.dumps(dict(span.attributes or {}))
+        events_json = json.dumps(
+            [
+                {
+                    "name": event.name,
+                    "timestamp": event.timestamp,
+                    "attributes": dict(event.attributes or {}),
+                }
+                for event in span.events
+            ]
+        )
+        links_json = json.dumps(
+            [
+                {
+                    "context": {
+                        "trace_id": format(link.context.trace_id, "032x"),
+                        "span_id": format(link.context.span_id, "016x"),
+                    },
+                    "attributes": dict(link.attributes or {}),
+                }
+                for link in span.links
+            ]
+        )
+        resource_json = json.dumps(dict(span.resource.attributes.items()))
+
+        # Get parent span ID if exists
+        parent_id = None
+        if span.parent and span.parent.span_id != 0:
+            parent_id = format(span.parent.span_id, "016x")
+
+        return {
+            "trace_id": format(context.trace_id, "032x"),
+            "span_id": format(context.span_id, "016x"),
+            "parent_id": parent_id,
+            "name": span.name,
+            "service": service,
+            "operation": operation,
+            "kind": span.kind.name if span.kind else None,
+            "start_time": span.start_time,
+            "end_time": span.end_time,
+            "duration_ms": duration_ms,
+            "status_code": status.status_code.name,
+            "status_description": status.description,
+            "attributes": attributes_json,
+            "events": events_json,
+            "links": links_json,
+            "resource": resource_json,
+        }
+
+    def export(self, spans):
+        """Export spans to DuckDB."""
+        try:
+            with duckdb.connect(str(self.db_path)) as conn:
+                for span in spans:
+                    record = self._span_to_record(span)
+
+                    # Insert span record
+                    conn.execute(
+                        """
+                        INSERT OR REPLACE INTO spans (
+                            trace_id, span_id, parent_id, name, service, operation,
+                            kind, start_time, end_time, duration_ms,
+                            status_code, status_description,
+                            attributes, events, links, resource
+                        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                    """,
+                        (
+                            record["trace_id"],
+                            record["span_id"],
+                            record["parent_id"],
+                            record["name"],
+                            record["service"],
+                            record["operation"],
+                            record["kind"],
+                            record["start_time"],
+                            record["end_time"],
+                            record["duration_ms"],
+                            record["status_code"],
+                            record["status_description"],
+                            record["attributes"],
+                            record["events"],
+                            record["links"],
+                            record["resource"],
+                        ),
+                    )
+
+            return SpanExportResult.SUCCESS
+        except Exception as e:
+            print(f"Error exporting spans to DuckDB: {e}")
+            return SpanExportResult.FAILURE
+
+    def shutdown(self) -> None:
+        """Cleanup resources."""
+        # DuckDB connections are managed per-transaction, no cleanup needed

flock/{core/logging → logging}/telemetry_exporter/file_exporter.py

@@ -2,17 +2,17 @@
 
 import json
 
+# with workflow.unsafe.imports_passed_through():
+from pathlib import Path
+
 from opentelemetry.sdk.trace.export import SpanExportResult
 from opentelemetry.trace import Status, StatusCode
-from temporalio import workflow
 
-from flock.core.logging.telemetry_exporter.base_exporter import (
+# from temporalio import workflow
+from flock.logging.telemetry_exporter.base_exporter import (
     TelemetryExporter,
 )
 
-with workflow.unsafe.imports_passed_through():
-    from pathlib import Path
-
 
 class FileSpanExporter(TelemetryExporter):
     """A simple exporter that writes span data as JSON lines into a file."""
@@ -29,7 +29,7 @@ class FileSpanExporter(TelemetryExporter):
         context = span.get_span_context()
         status = span.status or Status(StatusCode.UNSET)
 
-        return {
+        result = {
             "name": span.name,
             "context": {
                 "trace_id": format(context.trace_id, "032x"),
@@ -63,12 +63,15 @@ class FileSpanExporter(TelemetryExporter):
                 }
                 for link in span.links
             ],
-            "resource": {
-                attr_key: attr_value
-                for attr_key, attr_value in span.resource.attributes.items()
-            },
+            "resource": dict(span.resource.attributes.items()),
         }
 
+        # Add parent_id if this span has a parent
+        if span.parent and span.parent.span_id != 0:
+            result["parent_id"] = format(span.parent.span_id, "016x")
+
+        return result
+
     def export(self, spans):
         """Write spans to a log file."""
         try:

flock/{core/logging → logging}/telemetry_exporter/sqlite_exporter.py

@@ -7,7 +7,7 @@ from typing import Any
 
 from opentelemetry.sdk.trace.export import SpanExportResult
 
-from flock.core.logging.telemetry_exporter.base_exporter import (
+from flock.logging.telemetry_exporter.base_exporter import (
     TelemetryExporter,
 )
 
@@ -77,7 +77,7 @@ class SqliteTelemetryExporter(TelemetryExporter):
                 trace_id = format(span.context.trace_id, "032x")
                 cursor.execute(
                     """
-                    INSERT OR REPLACE INTO spans 
+                    INSERT OR REPLACE INTO spans
                     (id, name, trace_id, span_id, start_time, end_time, attributes, status)
                     VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                     """,
@@ -100,4 +100,3 @@ class SqliteTelemetryExporter(TelemetryExporter):
 
     def shutdown(self) -> None:
         """Cleanup resources."""
-        pass

flock/logging/trace_and_logged.py

@@ -0,0 +1,304 @@
+"""A decorator that wraps a function in an OpenTelemetry span and logs its inputs, outputs, and exceptions."""
+
+import functools
+import inspect
+import json
+
+from opentelemetry import trace
+from opentelemetry.trace import Status, StatusCode
+
+from flock.logging.logging import get_logger
+
+
+logger = get_logger("tools")
+tracer = trace.get_tracer(__name__)
+
+
+# Global trace filter configuration
+class TraceFilterConfig:
+    """Configuration for filtering which operations get traced."""
+
+    def __init__(self):
+        self.services: set[str] | None = None  # Whitelist: only trace these services
+        self.ignore_operations: set[str] = set()  # Blacklist: never trace these operations
+
+    def should_trace(self, service: str, operation: str) -> bool:
+        """Check if an operation should be traced based on filters.
+
+        Args:
+            service: Service name (e.g., "Flock", "Agent")
+            operation: Full operation name (e.g., "Flock.publish")
+
+        Returns:
+            True if should trace, False otherwise
+        """
+        # Check blacklist first (highest priority)
+        if operation in self.ignore_operations:
+            return False
+
+        # Check whitelist if configured
+        if self.services is not None:
+            service_lower = service.lower()
+            if service_lower not in self.services:
+                return False
+
+        return True
+
+
+# Global instance
+_trace_filter_config = TraceFilterConfig()
+
+
+def _serialize_value(value, max_depth=10, current_depth=0):
+    """Serialize a value to JSON-compatible format for span attributes.
+
+    Args:
+        value: The value to serialize
+        max_depth: Maximum recursion depth to prevent infinite loops
+        current_depth: Current recursion depth
+
+    Returns:
+        A JSON-serializable representation of the value
+    """
+    if current_depth >= max_depth:
+        return f"<max_depth_reached: {type(value).__name__}>"
+
+    try:
+        # Handle None
+        if value is None:
+            return None
+
+        # Handle primitives - these are already JSON-serializable
+        if isinstance(value, (str, int, float, bool)):
+            return value
+
+        # Handle lists/tuples
+        if isinstance(value, (list, tuple)):
+            return [_serialize_value(item, max_depth, current_depth + 1) for item in value]
+
+        # Handle dicts
+        if isinstance(value, dict):
+            return {
+                str(k): _serialize_value(v, max_depth, current_depth + 1) for k, v in value.items()
+            }
+
+        # Handle sets
+        if isinstance(value, set):
+            return [_serialize_value(item, max_depth, current_depth + 1) for item in value]
+
+        # For custom objects with __dict__, serialize their attributes
+        if hasattr(value, "__dict__"):
+            class_name = value.__class__.__name__
+            try:
+                obj_dict = {}
+                for k, v in value.__dict__.items():
+                    # Skip private attributes and methods
+                    if k.startswith("_"):
+                        continue
+                    # Skip methods and callables
+                    if callable(v):
+                        continue
+                    try:
+                        obj_dict[k] = _serialize_value(v, max_depth, current_depth + 1)
+                    except Exception:
+                        obj_dict[k] = f"<error serializing {k}>"
+
+                return {
+                    "__class__": class_name,
+                    "__module__": value.__class__.__module__,
+                    **obj_dict,
+                }
+            except Exception as e:
+                return {"__class__": class_name, "__error__": str(e)}
+
+        # For objects with a useful string representation
+        result = str(value)
+        # If the string is too long (> 5000 chars), truncate it
+        if len(result) > 5000:
+            return result[:5000] + "... (string truncated at 5000 chars)"
+        return result
+
+    except Exception as e:
+        # If all else fails, return type information with error
+        return {"__type__": type(value).__name__, "__error__": str(e)}
+
+
+def _extract_span_attributes(func, args, kwargs):
+    """Extract useful attributes from function arguments for OTEL spans.
+
+    Returns a dict of attributes and a display name for the span.
+    """
+    attributes = {}
+    span_name = func.__name__
+
+    # Try to get class name if this is a method
+    if args and hasattr(args[0], "__class__"):
+        obj = args[0]
+        class_name = obj.__class__.__name__
+        span_name = f"{class_name}.{func.__name__}"
+        attributes["class"] = class_name
+
+        # Extract agent-specific attributes
+        if hasattr(obj, "name"):
+            attributes["agent.name"] = str(obj.name)
+        if hasattr(obj, "description"):
+            attributes["agent.description"] = str(obj.description)[:200]  # Truncate
+
+    # Extract context attributes (correlation_id, task_id)
+    for arg_name, arg_value in kwargs.items():
+        if arg_name == "ctx" and hasattr(arg_value, "correlation_id"):
+            if arg_value.correlation_id:
+                attributes["correlation_id"] = str(arg_value.correlation_id)
+            if hasattr(arg_value, "task_id"):
+                attributes["task_id"] = str(arg_value.task_id)
+
+    # Check positional args for Context
+    for arg in args[1:]:  # Skip self
+        if hasattr(arg, "correlation_id"):
+            if arg.correlation_id:
+                attributes["correlation_id"] = str(arg.correlation_id)
+            if hasattr(arg, "task_id"):
+                attributes["task_id"] = str(arg.task_id)
+            break
+
+    # Add function metadata
+    attributes["function"] = func.__name__
+    attributes["module"] = func.__module__
+
+    # Capture input arguments (skip 'self' for methods)
+    try:
+        sig = inspect.signature(func)
+        bound_args = sig.bind(*args, **kwargs)
+        bound_args.apply_defaults()
+
+        # Serialize arguments
+        for param_name, param_value in bound_args.arguments.items():
+            # Skip 'self' and 'cls'
+            if param_name in ("self", "cls"):
+                continue
+            # Serialize the argument value to JSON-compatible format
+            serialized = _serialize_value(param_value)
+            # Convert to JSON string for OTEL attribute storage
+            try:
+                attributes[f"input.{param_name}"] = json.dumps(serialized, default=str)
+            except Exception:
+                # If JSON serialization fails, use string representation
+                attributes[f"input.{param_name}"] = str(serialized)
+    except Exception as e:
+        # If we can't capture inputs, just note that
+        attributes["input.error"] = str(e)
+
+    return attributes, span_name
+
+
+def traced_and_logged(func):
+    """A decorator that wraps a function in an OpenTelemetry span.
+
+    Creates proper parent-child span relationships and extracts relevant
+    attributes for observability in Grafana/Jaeger.
+
+    Automatically extracts:
+    - Agent name and description
+    - Correlation ID and task ID from Context
+    - Class and method names
+    - Exception information
+
+    Supports both synchronous and asynchronous functions.
+    """
+    if inspect.iscoroutinefunction(func):
+
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            attributes, span_name = _extract_span_attributes(func, args, kwargs)
+
+            # Check if we should trace this operation
+            service_name = span_name.split(".")[0] if "." in span_name else span_name
+            if not _trace_filter_config.should_trace(service_name, span_name):
+                # Skip tracing, just call the function
+                return await func(*args, **kwargs)
+
+            with tracer.start_as_current_span(span_name) as span:
+                # Set all extracted attributes
+                for key, value in attributes.items():
+                    span.set_attribute(key, value)
+
+                try:
+                    result = await func(*args, **kwargs)
+
+                    # Capture output value as JSON
+                    try:
+                        serialized_result = _serialize_value(result)
+                        span.set_attribute(
+                            "output.value", json.dumps(serialized_result, default=str)
+                        )
+                    except Exception as e:
+                        span.set_attribute("output.value", str(result))
+                        span.set_attribute("output.serialization_error", str(e))
+
+                    # Set result type and metadata
+                    if result is not None:
+                        span.set_attribute("output.type", type(result).__name__)
+                        if hasattr(result, "__len__"):
+                            try:
+                                span.set_attribute("output.length", len(result))
+                            except TypeError:
+                                pass
+
+                    span.set_status(Status(StatusCode.OK))
+                    logger.debug(f"{span_name} executed successfully")
+                    return result
+
+                except Exception as e:
+                    span.set_status(Status(StatusCode.ERROR, str(e)))
+                    span.record_exception(e)
+                    logger.exception(f"Error in {span_name}", error=str(e))
+                    raise
+
+        return async_wrapper
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        attributes, span_name = _extract_span_attributes(func, args, kwargs)
+
+        # Check if we should trace this operation
+        service_name = span_name.split(".")[0] if "." in span_name else span_name
+        if not _trace_filter_config.should_trace(service_name, span_name):
+            # Skip tracing, just call the function
+            return func(*args, **kwargs)
+
+        with tracer.start_as_current_span(span_name) as span:
+            # Set all extracted attributes
+            for key, value in attributes.items():
+                span.set_attribute(key, value)
+
+            try:
+                result = func(*args, **kwargs)
+
+                # Capture output value as JSON
+                try:
+                    serialized_result = _serialize_value(result)
+                    span.set_attribute("output.value", json.dumps(serialized_result, default=str))
+                except Exception as e:
+                    span.set_attribute("output.value", str(result))
+                    span.set_attribute("output.serialization_error", str(e))
+
+                # Set result type and metadata
+                if result is not None:
+                    span.set_attribute("output.type", type(result).__name__)
+                    if hasattr(result, "__len__"):
+                        try:
+                            span.set_attribute("output.length", len(result))
+                        except TypeError:
+                            pass
+
+                span.set_status(Status(StatusCode.OK))
+                logger.debug(f"{span_name} executed successfully")
+                return result
+
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                logger.exception(f"Error in {span_name}", error=str(e))
+                raise
+
+    return wrapper

flock/mcp/__init__.py

@@ -0,0 +1,91 @@
+"""MCP (Model Context Protocol) Integration for Flock-Flow.
+
+This package provides integration with MCP servers, enabling agents to
+dynamically discover and use external tools following the Model Context Protocol.
+
+Architecture Decisions:
+- AD001: Two-Level Architecture (orchestrator + agent)
+- AD003: Tool Namespacing ({server}__{tool})
+- AD004: Per-(agent_id, run_id) Connection Isolation
+- AD005: Lazy Connection Establishment
+- AD007: Graceful Degradation on MCP Failures
+
+Key Components:
+- FlockMCPConfiguration: Server configuration
+- FlockMCPClient: Individual server connection
+- FlockMCPClientManager: Connection pooling and lifecycle
+- FlockMCPTool: MCP tool wrapper compatible with DSPy
+
+Example Usage:
+    ```python
+    from flock import Flock
+    from flock.mcp import StdioServerParameters
+
+    # Create orchestrator
+    orchestrator = Flock()
+
+    # Register MCP server
+    orchestrator.add_mcp(
+        name="filesystem",
+        connection_params=StdioServerParameters(
+            command="uvx",
+            args=["mcp-server-filesystem", "/tmp"]
+        )
+    )
+
+    # Build agent with MCP access
+    agent = (
+        orchestrator.agent("file_agent")
+        .with_mcps(["filesystem"])
+        .build()
+    )
+    ```
+"""
+
+from flock.mcp.client import FlockMCPClient
+from flock.mcp.config import (
+    FlockMCPCachingConfiguration,
+    FlockMCPCallbackConfiguration,
+    FlockMCPConfiguration,
+    FlockMCPConnectionConfiguration,
+    FlockMCPFeatureConfiguration,
+)
+from flock.mcp.manager import FlockMCPClientManager
+from flock.mcp.tool import FlockMCPTool
+from flock.mcp.types import (
+    FlockListRootsMCPCallback,
+    FlockLoggingMCPCallback,
+    FlockMessageHandlerMCPCallback,
+    FlockSamplingMCPCallback,
+    MCPRoot,
+    ServerParameters,
+    SseServerParameters,
+    StdioServerParameters,
+    StreamableHttpServerParameters,
+    WebsocketServerParameters,
+)
+
+
+__all__ = [
+    "FlockListRootsMCPCallback",
+    "FlockLoggingMCPCallback",
+    "FlockMCPCachingConfiguration",
+    "FlockMCPCallbackConfiguration",
+    # Client and Manager
+    "FlockMCPClient",
+    "FlockMCPClientManager",
+    # Configuration
+    "FlockMCPConfiguration",
+    "FlockMCPConnectionConfiguration",
+    "FlockMCPFeatureConfiguration",
+    "FlockMCPTool",
+    "FlockMessageHandlerMCPCallback",
+    "FlockSamplingMCPCallback",
+    "MCPRoot",
+    # Types
+    "ServerParameters",
+    "SseServerParameters",
+    "StdioServerParameters",
+    "StreamableHttpServerParameters",
+    "WebsocketServerParameters",
+]