flock-core 0.5.0b53__py3-none-any.whl → 0.5.0b55__py3-none-any.whl

flock/logging/trace_and_logged.py

@@ -2,8 +2,10 @@
 
 import functools
 import inspect
+import json
 
 from opentelemetry import trace
+from opentelemetry.trace import Status, StatusCode
 
 from flock.logging.logging import get_logger
 
@@ -12,44 +14,291 @@ 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.
 
-    and logs its inputs,
-    outputs, and exceptions. Supports both synchronous and asynchronous functions.
+    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):
-            with tracer.start_as_current_span(func.__name__) as span:
-                span.set_attribute("args", str(args))
-                span.set_attribute("kwargs", str(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)
-                    span.set_attribute("result", str(result))
-                    logger.debug(f"{func.__name__} executed successfully", result=result)
+
+                    # 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:
-                    logger.exception(f"Error in {func.__name__}", error=str(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):
-        with tracer.start_as_current_span(func.__name__) as span:
-            span.set_attribute("args", str(args))
-            span.set_attribute("kwargs", str(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)
-                span.set_attribute("result", str(result))
-                logger.debug(f"{func.__name__} executed successfully", result=result)
+
+                # 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:
-                logger.exception(f"Error in {func.__name__}", error=str(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/orchestrator.py

@@ -10,11 +10,14 @@ from contextlib import asynccontextmanager
 from typing import TYPE_CHECKING, Any
 from uuid import uuid4
 
+from opentelemetry import trace
+from opentelemetry.trace import Status, StatusCode
 from pydantic import BaseModel
 
 from flock.agent import Agent, AgentBuilder
 from flock.artifacts import Artifact
 from flock.helper.cli_helper import init_console
+from flock.logging.auto_trace import AutoTracedMeta
 from flock.mcp import (
     FlockMCPClientManager,
     FlockMCPConfiguration,
@@ -48,7 +51,12 @@ class BoardHandle:
         return await self._orchestrator.store.list()
 
 
-class Flock:
+class Flock(metaclass=AutoTracedMeta):
+    """Main orchestrator for blackboard-based agent coordination.
+
+    All public methods are automatically traced via OpenTelemetry.
+    """
+
     def _patch_litellm_proxy_imports(self) -> None:
         """Stub litellm proxy_server to avoid optional proxy deps when not used.
 
@@ -91,6 +99,14 @@ class Flock:
         self.max_agent_iterations: int = max_agent_iterations
         self._agent_iteration_count: dict[str, int] = {}
         self.is_dashboard: bool = False
+        # Unified tracing support
+        self._workflow_span = None
+        self._auto_workflow_enabled = os.getenv("FLOCK_AUTO_WORKFLOW_TRACE", "false").lower() in {
+            "true",
+            "1",
+            "yes",
+            "on",
+        }
         if not model:
             self.model = os.getenv("DEFAULT_MODEL")
 
@@ -214,6 +230,119 @@ class Flock:
 
         return self._mcp_manager
 
+    # Unified Tracing ------------------------------------------------------
+
+    @asynccontextmanager
+    async def traced_run(self, name: str = "workflow"):
+        """Context manager for wrapping an entire execution in a single unified trace.
+
+        This creates a parent span that encompasses all operations (publish, run_until_idle, etc.)
+        within the context, ensuring they all belong to the same trace_id for better observability.
+
+        Args:
+            name: Name for the workflow trace (default: "workflow")
+
+        Yields:
+            The workflow span for optional manual attribute setting
+
+        Examples:
+            # Explicit workflow tracing (recommended)
+            async with flock.traced_run("pizza_workflow"):
+                await flock.publish(pizza_idea)
+                await flock.run_until_idle()
+                # All operations now share the same trace_id!
+
+            # Custom attributes
+            async with flock.traced_run("data_pipeline") as span:
+                span.set_attribute("pipeline.version", "2.0")
+                await flock.publish(data)
+                await flock.run_until_idle()
+        """
+        tracer = trace.get_tracer(__name__)
+        with tracer.start_as_current_span(name) as span:
+            # Set workflow-level attributes
+            span.set_attribute("flock.workflow", True)
+            span.set_attribute("workflow.name", name)
+            span.set_attribute("workflow.flock_id", str(id(self)))
+
+            # Store span for nested operations to use
+            prev_workflow_span = self._workflow_span
+            self._workflow_span = span
+
+            try:
+                yield span
+                span.set_status(Status(StatusCode.OK))
+            except Exception as e:
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                raise
+            finally:
+                # Restore previous workflow span
+                self._workflow_span = prev_workflow_span
+
+    @staticmethod
+    def clear_traces(db_path: str = ".flock/traces.duckdb") -> dict[str, Any]:
+        """Clear all traces from the DuckDB database.
+
+        Useful for resetting debug sessions or cleaning up test data.
+
+        Args:
+            db_path: Path to the DuckDB database file (default: ".flock/traces.duckdb")
+
+        Returns:
+            Dictionary with operation results:
+                - deleted_count: Number of spans deleted
+                - success: Whether operation succeeded
+                - error: Error message if failed
+
+        Examples:
+            # Clear all traces
+            result = Flock.clear_traces()
+            print(f"Deleted {result['deleted_count']} spans")
+
+            # Custom database path
+            result = Flock.clear_traces(".flock/custom_traces.duckdb")
+
+            # Check if operation succeeded
+            if result['success']:
+                print("Traces cleared successfully!")
+            else:
+                print(f"Error: {result['error']}")
+        """
+        try:
+            from pathlib import Path
+
+            import duckdb
+
+            db_file = Path(db_path)
+            if not db_file.exists():
+                return {
+                    "success": False,
+                    "deleted_count": 0,
+                    "error": f"Database file not found: {db_path}",
+                }
+
+            # Connect and clear
+            conn = duckdb.connect(str(db_file))
+            try:
+                # Get count before deletion
+                count_result = conn.execute("SELECT COUNT(*) FROM spans").fetchone()
+                deleted_count = count_result[0] if count_result else 0
+
+                # Delete all spans
+                conn.execute("DELETE FROM spans")
+
+                # Vacuum to reclaim space
+                conn.execute("VACUUM")
+
+                return {"success": True, "deleted_count": deleted_count, "error": None}
+
+            finally:
+                conn.close()
+
+        except Exception as e:
+            return {"success": False, "deleted_count": 0, "error": str(e)}
+
     # Runtime --------------------------------------------------------------
 
     async def run_until_idle(self) -> None:

flock/store.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 
 from asyncio import Lock
 from collections import defaultdict
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeVar
 
 from flock.registry import type_registry
 
@@ -17,6 +17,8 @@ if TYPE_CHECKING:
 
     from flock.artifacts import Artifact
 
+T = TypeVar("T")
+
 
 class BlackboardStore:
     async def publish(self, artifact: Artifact) -> None:
@@ -31,6 +33,21 @@ class BlackboardStore:
     async def list_by_type(self, type_name: str) -> builtins.list[Artifact]:
         raise NotImplementedError
 
+    async def get_by_type(self, artifact_type: type[T]) -> builtins.list[T]:
+        """Get artifacts by Pydantic type, returning data already cast.
+
+        Args:
+            artifact_type: The Pydantic model class (e.g., BugAnalysis)
+
+        Returns:
+            List of data objects of the specified type (not Artifact wrappers)
+
+        Example:
+            bug_analyses = await store.get_by_type(BugAnalysis)
+            # Returns list[BugAnalysis] directly, no .data access needed
+        """
+        raise NotImplementedError
+
 
 class InMemoryBlackboardStore(BlackboardStore):
     """Simple in-memory implementation suitable for local dev and tests."""
@@ -58,6 +75,22 @@ class InMemoryBlackboardStore(BlackboardStore):
             canonical = type_registry.resolve_name(type_name)
             return list(self._by_type.get(canonical, []))
 
+    async def get_by_type(self, artifact_type: type[T]) -> builtins.list[T]:
+        """Get artifacts by Pydantic type, returning data already cast.
+
+        Args:
+            artifact_type: The Pydantic model class (e.g., BugAnalysis)
+
+        Returns:
+            List of data objects of the specified type (not Artifact wrappers)
+        """
+        async with self._lock:
+            # Get canonical name from the type
+            canonical = type_registry.resolve_name(artifact_type.__name__)
+            artifacts = self._by_type.get(canonical, [])
+            # Reconstruct Pydantic models from payload dictionaries
+            return [artifact_type(**artifact.payload) for artifact in artifacts]  # type: ignore
+
     async def extend(self, artifacts: Iterable[Artifact]) -> None:  # pragma: no cover - helper
         for artifact in artifacts:
             await self.publish(artifact)