ai-pipeline-core 0.2.6__py3-none-any.whl → 0.4.1__py3-none-any.whl

ai_pipeline_core/logging/logging_config.py

@@ -6,7 +6,7 @@ Provides logging configuration management that integrates with Prefect's logging
 import logging.config
 import os
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any
 
 import yaml
 from prefect.logging import get_logger
@@ -16,7 +16,7 @@ DEFAULT_LOG_LEVELS = {
     "ai_pipeline_core": "INFO",
     "ai_pipeline_core.documents": "INFO",
     "ai_pipeline_core.llm": "INFO",
-    "ai_pipeline_core.flow": "INFO",
+    "ai_pipeline_core.pipeline": "INFO",
     "ai_pipeline_core.testing": "DEBUG",
 }
 
@@ -32,22 +32,19 @@ class LoggingConfig:
         3. PREFECT_LOGGING_SETTINGS_PATH environment variable
         4. Default configuration
 
-    Example:
-        >>> config = LoggingConfig()
-        >>> config.apply()
     """
 
-    def __init__(self, config_path: Optional[Path] = None):
+    def __init__(self, config_path: Path | None = None):
         """Initialize logging configuration.
 
         Args:
             config_path: Optional path to YAML configuration file.
         """
         self.config_path = config_path or self._get_default_config_path()
-        self._config: Optional[Dict[str, Any]] = None
+        self._config: dict[str, Any] | None = None
 
     @staticmethod
-    def _get_default_config_path() -> Optional[Path]:
+    def _get_default_config_path() -> Path | None:
         """Get default config path from environment variables.
 
         Returns:
@@ -63,7 +60,7 @@ class LoggingConfig:
 
         return None
 
-    def load_config(self) -> Dict[str, Any]:
+    def load_config(self) -> dict[str, Any]:
         """Load logging configuration from file or defaults.
 
         Returns:
@@ -71,7 +68,7 @@ class LoggingConfig:
         """
         if self._config is None:
             if self.config_path and self.config_path.exists():
-                with open(self.config_path, "r") as f:
+                with open(self.config_path, encoding="utf-8") as f:
                     self._config = yaml.safe_load(f)
             else:
                 self._config = self._get_default_config()
@@ -80,7 +77,7 @@ class LoggingConfig:
         return self._config
 
     @staticmethod
-    def _get_default_config() -> Dict[str, Any]:
+    def _get_default_config() -> dict[str, Any]:
         """Get default logging configuration.
 
         Returns:
@@ -95,10 +92,7 @@ class LoggingConfig:
                     "datefmt": "%H:%M:%S",
                 },
                 "detailed": {
-                    "format": (
-                        "%(asctime)s | %(levelname)-7s | %(name)s | "
-                        "%(funcName)s:%(lineno)d - %(message)s"
-                    ),
+                    "format": ("%(asctime)s | %(levelname)-7s | %(name)s | %(funcName)s:%(lineno)d - %(message)s"),
                     "datefmt": "%Y-%m-%d %H:%M:%S",
                 },
             },
@@ -134,10 +128,10 @@ class LoggingConfig:
 
 
 # Global configuration instance
-_logging_config: Optional[LoggingConfig] = None
+_logging_config: LoggingConfig | None = None
 
 
-def setup_logging(config_path: Optional[Path] = None, level: Optional[str] = None):
+def setup_logging(config_path: Path | None = None, level: str | None = None):
     """Setup logging for the AI Pipeline Core library.
 
     Initializes logging configuration for the pipeline system.
@@ -149,18 +143,8 @@ def setup_logging(config_path: Optional[Path] = None, level: Optional[str] = Non
         config_path: Optional path to YAML logging configuration file.
         level: Optional log level override (INFO, DEBUG, WARNING, etc.).
 
-    Example:
-        >>> # In your main.py or application entry point:
-        >>> def main():
-        ...     setup_logging()  # Call once at startup
-        ...     # Your application code here
-        ...
-        >>> # Or with custom level:
-        >>> if __name__ == "__main__":
-        ...     setup_logging(level="DEBUG")
-        ...     run_application()
     """
-    global _logging_config
+    global _logging_config  # noqa: PLW0603
 
     _logging_config = LoggingConfig(config_path)
     _logging_config.apply()
@@ -179,22 +163,28 @@ def setup_logging(config_path: Optional[Path] = None, level: Optional[str] = Non
 def get_pipeline_logger(name: str):
     """Get a logger for pipeline components.
 
-    @public
-
-    Returns a Prefect-integrated logger with proper configuration.
+    Returns a Prefect-integrated logger with the OTel span-event bridge
+    attached.  Any log record at INFO+ emitted while an OTel span is
+    recording will be captured as a span event in the trace.
 
     Args:
         name: Logger name, typically __name__.
 
     Returns:
-        Prefect logger instance.
+        Prefect logger instance with bridge handler.
 
-    Example:
-        >>> logger = get_pipeline_logger(__name__)
-        >>> logger.info("Module initialized")
     """
-    # Ensure logging is setup
     if _logging_config is None:
         setup_logging()
 
-    return get_logger(name)
+    logger = get_logger(name)
+
+    # Attach the singleton bridge handler so log records become OTel span events.
+    # The handler is a no-op when no span is recording, so early attachment is safe.
+    from ai_pipeline_core.observability._logging_bridge import get_bridge_handler  # noqa: PLC0415
+
+    handler = get_bridge_handler()
+    if handler not in logger.handlers:
+        logger.addHandler(handler)
+
+    return logger

ai_pipeline_core/logging/logging_mixin.py

@@ -2,9 +2,10 @@
 
 import contextlib
 import time
+from collections.abc import Generator
 from contextlib import contextmanager
 from functools import cached_property
-from typing import Any, Dict, Generator, Optional
+from typing import Any
 
 from prefect import get_run_logger
 from prefect.context import FlowRunContext, TaskRunContext
@@ -23,7 +24,7 @@ class LoggerMixin:
     - Internal routing when outside flow/task context
     """
 
-    _logger_name: Optional[str] = None
+    _logger_name: str | None = None
 
     @cached_property
     def logger(self):
@@ -32,7 +33,8 @@ class LoggerMixin:
             return logger
         return get_logger(self._logger_name or self.__class__.__module__)
 
-    def _get_run_logger(self):
+    @staticmethod
+    def _get_run_logger():
         """Attempt to get Prefect run logger.
 
         Returns:
@@ -56,15 +58,15 @@ class LoggerMixin:
         """Log warning message with optional context."""
         self.logger.warning(message, extra=kwargs)
 
-    def log_error(self, message: str, exc_info: bool = False, **kwargs: Any) -> None:
+    def log_error(self, message: str, *, exc_info: bool = False, **kwargs: Any) -> None:
         """Log error message with optional exception info."""
         self.logger.error(message, exc_info=exc_info, extra=kwargs)
 
-    def log_critical(self, message: str, exc_info: bool = False, **kwargs: Any) -> None:
+    def log_critical(self, message: str, *, exc_info: bool = False, **kwargs: Any) -> None:
         """Log critical message with optional exception info."""
         self.logger.critical(message, exc_info=exc_info, extra=kwargs)
 
-    def log_with_context(self, level: str, message: str, context: Dict[str, Any]) -> None:
+    def log_with_context(self, level: str, message: str, context: dict[str, Any]) -> None:
         """Log message with structured context.
 
         Args:
@@ -72,12 +74,6 @@ class LoggerMixin:
             message: Log message
             context: Additional context as dictionary
 
-        Example:
-            self.log_with_context("info", "Processing document", {
-                "document_id": doc.id,
-                "document_size": doc.size,
-                "document_type": doc.type
-            })
         """
         log_method = getattr(self.logger, level.lower(), self.logger.info)
 
@@ -98,11 +94,6 @@ class StructuredLoggerMixin(LoggerMixin):
             event: Event name
             **kwargs: Event attributes
 
-        Example:
-            self.log_event("document_processed",
-                          document_id=doc.id,
-                          duration_ms=processing_time,
-                          status="success")
         """
         self.logger.info(event, extra={"event": event, "structured": True, **kwargs})
 
@@ -115,9 +106,6 @@ class StructuredLoggerMixin(LoggerMixin):
             unit: Unit of measurement
             **tags: Additional tags
 
-        Example:
-            self.log_metric("processing_time", 1.23, "seconds",
-                          document_type="pdf", model="gpt-4")
         """
         self.logger.info(
             f"Metric: {metric_name}",
@@ -138,9 +126,6 @@ class StructuredLoggerMixin(LoggerMixin):
             duration_ms: Duration in milliseconds
             **attributes: Additional attributes
 
-        Example:
-            self.log_span("llm_generation", 1234.5,
-                         model="gpt-4", tokens=500)
         """
         self.logger.info(
             f"Span: {operation}",
@@ -160,9 +145,6 @@ class StructuredLoggerMixin(LoggerMixin):
             operation: Operation name
             **context: Additional context
 
-        Example:
-            with self.log_operation("document_processing", doc_id=doc.id):
-                process_document(doc)
         """
         start_time = time.perf_counter()
 
@@ -171,14 +153,12 @@ class StructuredLoggerMixin(LoggerMixin):
         try:
             yield
             duration_ms = (time.perf_counter() - start_time) * 1000
-            self.log_info(
-                f"Completed {operation}", duration_ms=duration_ms, status="success", **context
-            )
+            self.log_info(f"Completed {operation}", duration_ms=duration_ms, status="success", **context)
         except Exception as e:
             # Intentionally broad: Context manager must catch all exceptions to log them
             duration_ms = (time.perf_counter() - start_time) * 1000
             self.log_error(
-                f"Failed {operation}: {str(e)}",
+                f"Failed {operation}: {e!s}",
                 exc_info=True,
                 duration_ms=duration_ms,
                 status="failure",
@@ -190,31 +170,25 @@ class StructuredLoggerMixin(LoggerMixin):
 class PrefectLoggerMixin(StructuredLoggerMixin):
     """Enhanced mixin specifically for Prefect flows and tasks."""
 
-    def log_flow_start(self, flow_name: str, parameters: Dict[str, Any]) -> None:
+    def log_flow_start(self, flow_name: str, parameters: dict[str, Any]) -> None:
         """Log flow start with parameters."""
         self.log_event("flow_started", flow_name=flow_name, parameters=parameters)
 
     def log_flow_end(self, flow_name: str, status: str, duration_ms: float) -> None:
         """Log flow completion."""
-        self.log_event(
-            "flow_completed", flow_name=flow_name, status=status, duration_ms=duration_ms
-        )
+        self.log_event("flow_completed", flow_name=flow_name, status=status, duration_ms=duration_ms)
 
-    def log_task_start(self, task_name: str, inputs: Dict[str, Any]) -> None:
+    def log_task_start(self, task_name: str, inputs: dict[str, Any]) -> None:
         """Log task start with inputs."""
         self.log_event("task_started", task_name=task_name, inputs=inputs)
 
     def log_task_end(self, task_name: str, status: str, duration_ms: float) -> None:
         """Log task completion."""
-        self.log_event(
-            "task_completed", task_name=task_name, status=status, duration_ms=duration_ms
-        )
+        self.log_event("task_completed", task_name=task_name, status=status, duration_ms=duration_ms)
 
     def log_retry(self, operation: str, attempt: int, max_attempts: int, error: str) -> None:
         """Log retry attempt."""
-        self.log_warning(
-            f"Retrying {operation}", attempt=attempt, max_attempts=max_attempts, error=error
-        )
+        self.log_warning(f"Retrying {operation}", attempt=attempt, max_attempts=max_attempts, error=error)
 
     def log_checkpoint(self, checkpoint_name: str, **data: Any) -> None:
         """Log a checkpoint in processing."""

ai_pipeline_core/observability/__init__.py

@@ -0,0 +1,32 @@
+"""Observability system for AI pipelines.
+
+Contains debug tracing, ClickHouse-based tracking, and initialization utilities.
+"""
+
+from ai_pipeline_core.observability._debug import (
+    ArtifactStore,
+    ContentRef,
+    ContentWriter,
+    LocalDebugSpanProcessor,
+    LocalTraceWriter,
+    SpanInfo,
+    TraceDebugConfig,
+    TraceState,
+    WriteJob,
+    generate_summary,
+)
+from ai_pipeline_core.observability._debug._content import reconstruct_span_content
+
+__all__ = [
+    "ArtifactStore",
+    "ContentRef",
+    "ContentWriter",
+    "LocalDebugSpanProcessor",
+    "LocalTraceWriter",
+    "SpanInfo",
+    "TraceDebugConfig",
+    "TraceState",
+    "WriteJob",
+    "generate_summary",
+    "reconstruct_span_content",
+]

ai_pipeline_core/observability/_debug/__init__.py

@@ -0,0 +1,30 @@
+"""Local trace debugging system for AI pipelines.
+
+This module provides filesystem-based trace debugging that saves all spans
+with their inputs/outputs for LLM-assisted debugging. Includes static
+summary generation and LLM-powered auto-summary capabilities.
+
+Enabled automatically in CLI mode (``run_cli``), writing to ``<working_dir>/.trace``.
+Disable with ``--no-trace``.
+"""
+
+from ._config import TraceDebugConfig
+from ._content import ArtifactStore, ContentRef, ContentWriter, reconstruct_span_content
+from ._processor import LocalDebugSpanProcessor
+from ._summary import generate_summary
+from ._types import SpanInfo, TraceState, WriteJob
+from ._writer import LocalTraceWriter
+
+__all__ = [
+    "ArtifactStore",
+    "ContentRef",
+    "ContentWriter",
+    "LocalDebugSpanProcessor",
+    "LocalTraceWriter",
+    "SpanInfo",
+    "TraceDebugConfig",
+    "TraceState",
+    "WriteJob",
+    "generate_summary",
+    "reconstruct_span_content",
+]

ai_pipeline_core/observability/_debug/_auto_summary.py

@@ -0,0 +1,94 @@
+"""LLM-powered auto-summary generation for trace debugging.
+
+Separated from _summary.py to avoid circular imports: this module depends on
+ai_pipeline_core.llm, which cannot be imported during the initial package load
+chain that includes _debug/__init__.py.
+"""
+
+from pydantic import BaseModel, ConfigDict
+
+from ai_pipeline_core.llm import generate_structured
+from ai_pipeline_core.llm.ai_messages import AIMessages
+from ai_pipeline_core.llm.model_options import ModelOptions
+
+from ._types import TraceState
+
+
+class AutoTraceSummary(BaseModel):
+    """LLM-generated trace analysis."""
+
+    model_config = ConfigDict(frozen=True)
+
+    overview: str
+    outcome: str
+    error_analysis: str
+    bottlenecks: tuple[str, ...] = ()
+    cost_assessment: str
+    recommendations: tuple[str, ...] = ()
+
+
+async def generate_auto_summary(
+    trace: TraceState,  # noqa: ARG001
+    static_summary: str,
+    model: str,
+) -> str | None:
+    """Generate LLM-powered auto-summary of the trace.
+
+    Args:
+        trace: Completed trace state with all span data.
+        static_summary: Pre-generated static summary text used as LLM input context.
+        model: LLM model name for summary generation.
+
+    Returns:
+        Formatted markdown auto-summary string, or None if generation fails.
+    """
+    messages = AIMessages()
+    messages.append(static_summary)
+
+    options = ModelOptions(
+        system_prompt=(
+            "You are analyzing an AI pipeline execution trace. "
+            "Provide concise, actionable analysis based on the execution data. "
+            "Focus on cost efficiency, performance bottlenecks, and errors."
+        ),
+    )
+
+    result = await generate_structured(
+        model=model,
+        response_format=AutoTraceSummary,
+        messages=messages,
+        options=options,
+        purpose="trace_auto_summary",
+    )
+
+    if not result or not result.parsed:
+        return None
+
+    summary = result.parsed
+    lines = [
+        "# Auto-Summary (LLM-Generated)",
+        "",
+        f"**Overview:** {summary.overview}",
+        "",
+        f"**Outcome:** {summary.outcome}",
+        "",
+    ]
+
+    if summary.error_analysis:
+        lines.append(f"**Error Analysis:** {summary.error_analysis}")
+        lines.append("")
+
+    if summary.bottlenecks:
+        lines.append("**Bottlenecks:**")
+        lines.extend(f"- {b}" for b in summary.bottlenecks)
+        lines.append("")
+
+    lines.append(f"**Cost Assessment:** {summary.cost_assessment}")
+    lines.append("")
+
+    if summary.recommendations:
+        lines.append("**Recommendations:**")
+        lines.extend(f"- {r}" for r in summary.recommendations)
+        lines.append("")
+
+    return "\n".join(lines)

ai_pipeline_core/observability/_debug/_config.py

@@ -0,0 +1,95 @@
+"""Configuration for local trace debugging."""
+
+from pathlib import Path
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class TraceDebugConfig(BaseModel):
+    """Configuration for local trace debugging.
+
+    Controls how traces are written to the local filesystem for debugging.
+    Enabled automatically in CLI mode, writing to ``<working_dir>/.trace``.
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    path: Path = Field(description="Directory for debug traces")
+    enabled: bool = Field(default=True, description="Whether debug tracing is enabled")
+
+    # Content size limits (Issue #2)
+    max_file_bytes: int = Field(
+        default=50_000,
+        description="Max bytes for input.yaml or output.yaml. Elements externalized to stay under.",
+    )
+    max_element_bytes: int = Field(
+        default=10_000,
+        description="Max bytes for single element. Above this, partial + artifact ref.",
+    )
+    element_excerpt_bytes: int = Field(
+        default=2_000,
+        description="Bytes of content to keep inline when element exceeds max_element_bytes.",
+    )
+    max_content_bytes: int = Field(
+        default=10_000_000,
+        description="Max bytes for any single artifact. Above this, truncate.",
+    )
+
+    # Image handling (Issue #7 - no changes per user)
+    extract_base64_images: bool = Field(
+        default=True,
+        description="Extract base64 images to artifact files",
+    )
+
+    # Span optimization (Issue #4)
+    merge_wrapper_spans: bool = Field(
+        default=True,
+        description="Merge Prefect wrapper spans with inner traced function spans",
+    )
+
+    # Indexes (Issue #1)
+    include_llm_index: bool = Field(
+        default=True,
+        description="Generate _llm_calls.yaml with LLM-specific details",
+    )
+    include_error_index: bool = Field(
+        default=True,
+        description="Generate _errors.yaml with failed span details",
+    )
+
+    # Cleanup
+    max_traces: int | None = Field(
+        default=None,
+        description="Max number of traces to keep. None for unlimited.",
+    )
+
+    # Security - default redaction patterns for common secrets
+    redact_patterns: tuple[str, ...] = Field(
+        default=(
+            r"sk-[a-zA-Z0-9]{20,}",  # OpenAI API keys
+            r"sk-proj-[a-zA-Z0-9\-_]{20,}",  # OpenAI project keys
+            r"AKIA[0-9A-Z]{16}",  # AWS access keys
+            r"ghp_[a-zA-Z0-9]{36}",  # GitHub personal tokens
+            r"gho_[a-zA-Z0-9]{36}",  # GitHub OAuth tokens
+            r"xoxb-[a-zA-Z0-9\-]+",  # Slack bot tokens
+            r"xoxp-[a-zA-Z0-9\-]+",  # Slack user tokens
+            r"(?i)password\s*[:=]\s*['\"]?[^\s'\"]+",  # Passwords
+            r"(?i)secret\s*[:=]\s*['\"]?[^\s'\"]+",  # Secrets
+            r"(?i)api[_\-]?key\s*[:=]\s*['\"]?[^\s'\"]+",  # API keys
+            r"(?i)bearer\s+[a-zA-Z0-9\-_\.]+",  # Bearer tokens
+        ),
+        description="Regex patterns for secrets to redact",
+    )
+
+    # Summary
+    generate_summary: bool = Field(default=True, description="Generate _summary.md")
+
+    # Auto-summary (LLM-powered)
+    auto_summary_enabled: bool = Field(
+        default=False,
+        description="Generate LLM-powered auto-summary after trace completion",
+    )
+    auto_summary_model: str = Field(
+        default="gemini-3-flash",
+        description="Model to use for auto-summary generation",
+    )