ai-pipeline-core 0.1.1__py3-none-any.whl

ai_pipeline_core/tracing.py

@@ -0,0 +1,205 @@
+"""Tracing utilities that integrate Laminar (``lmnr``) with our code-base.
+
+This module centralises:
+• ``TraceInfo`` - a small helper object for propagating contextual metadata.
+• ``trace`` decorator - augments a callable with Laminar tracing, automatic
+``observe`` instrumentation, and optional support for test runs.
+"""
+
+from __future__ import annotations
+
+import inspect
+import os
+from functools import wraps
+from typing import Any, Callable, ParamSpec, TypeVar, cast, overload
+
+from lmnr import Instruments, Laminar, observe
+from pydantic import BaseModel
+
+from ai_pipeline_core.settings import settings
+
+# ---------------------------------------------------------------------------
+# Typing helpers
+# ---------------------------------------------------------------------------
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+# ---------------------------------------------------------------------------
+# ``TraceInfo`` – metadata container
+# ---------------------------------------------------------------------------
+class TraceInfo(BaseModel):
+    """A container that holds contextual metadata for the current trace."""
+
+    session_id: str | None = None
+    user_id: str | None = None
+    metadata: dict[str, str] = {}
+    tags: list[str] = []
+
+    def get_observe_kwargs(self) -> dict[str, Any]:
+        """Return kwargs suitable for passing to the observe decorator."""
+        kwargs: dict[str, Any] = {}
+
+        # Use environment variable fallback for session_id
+        session_id = self.session_id or os.getenv("LMNR_SESSION_ID")
+        if session_id:
+            kwargs["session_id"] = session_id
+
+        # Use environment variable fallback for user_id
+        user_id = self.user_id or os.getenv("LMNR_USER_ID")
+        if user_id:
+            kwargs["user_id"] = user_id
+
+        if self.metadata:
+            kwargs["metadata"] = self.metadata
+        if self.tags:
+            kwargs["tags"] = self.tags
+        return kwargs
+
+
+# ---------------------------------------------------------------------------
+# ``trace`` decorator
+# ---------------------------------------------------------------------------
+
+
+def _initialise_laminar() -> None:
+    """Ensure Laminar is initialised once per process."""
+    if settings.lmnr_project_api_key:
+        Laminar.initialize(
+            project_api_key=settings.lmnr_project_api_key,
+            disabled_instruments=[Instruments.OPENAI],
+        )
+
+
+# Overload for calls like @trace(name="...", test=True)
+@overload
+def trace(
+    *,
+    name: str | None = None,
+    test: bool = False,
+    debug_only: bool = False,
+    ignore_input: bool = False,
+    ignore_output: bool = False,
+    ignore_inputs: list[str] | None = None,
+    input_formatter: Callable[..., str] | None = None,
+    output_formatter: Callable[..., str] | None = None,
+) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
+
+
+# Overload for the bare @trace call
+@overload
+def trace(func: Callable[P, R]) -> Callable[P, R]: ...
+
+
+# Actual implementation
+def trace(
+    func: Callable[P, R] | None = None,
+    *,
+    name: str | None = None,
+    test: bool = False,
+    debug_only: bool = False,
+    ignore_input: bool = False,
+    ignore_output: bool = False,
+    ignore_inputs: list[str] | None = None,
+    input_formatter: Callable[..., str] | None = None,
+    output_formatter: Callable[..., str] | None = None,
+) -> Callable[[Callable[P, R]], Callable[P, R]] | Callable[P, R]:
+    """Decorator that wires Laminar tracing and observation into a function.
+
+    Args:
+        func: The function to be traced (when used as @trace)
+        name: Custom name for the observation (defaults to function name)
+        test: Mark this trace as a test run
+        debug_only: Only trace when LMNR_DEBUG=true environment variable is set
+        ignore_input: Ignore all inputs in the trace
+        ignore_output: Ignore the output in the trace
+        ignore_inputs: List of specific input parameter names to ignore
+        input_formatter: Custom formatter for inputs (takes any arguments, returns string)
+        output_formatter: Custom formatter for outputs (takes any arguments, returns string)
+
+    Returns:
+        The decorated function with Laminar tracing enabled
+    """
+
+    def decorator(f: Callable[P, R]) -> Callable[P, R]:
+        # --- Pre-computation (done once when the function is decorated) ---
+        _initialise_laminar()
+        sig = inspect.signature(f)
+        is_coroutine = inspect.iscoroutinefunction(f)
+        decorator_test_flag = test
+        observe_name = name or f.__name__
+        _observe = observe
+
+        # Store the new parameters
+        _ignore_input = ignore_input
+        _ignore_output = ignore_output
+        _ignore_inputs = ignore_inputs
+        _input_formatter = input_formatter
+        _output_formatter = output_formatter
+
+        # --- Check debug_only flag and environment variable ---
+        if debug_only and os.getenv("LMNR_DEBUG", "").lower() != "true":
+            # If debug_only is True but LMNR_DEBUG is not set to "true",
+            # return the original function without tracing
+            return f
+
+        # --- Helper function for runtime logic ---
+        def _prepare_and_get_observe_params(runtime_kwargs: dict[str, Any]) -> dict[str, Any]:
+            """
+            Inspects runtime args, manages TraceInfo, and returns params for lmnr.observe.
+            Modifies runtime_kwargs in place to inject TraceInfo if the function expects it.
+            """
+            trace_info = runtime_kwargs.get("trace_info")
+            if not isinstance(trace_info, TraceInfo):
+                trace_info = TraceInfo()
+                if "trace_info" in sig.parameters:
+                    runtime_kwargs["trace_info"] = trace_info
+
+            runtime_test_flag = bool(runtime_kwargs.get("test", False))
+            if (decorator_test_flag or runtime_test_flag) and "test" not in trace_info.tags:
+                trace_info.tags.append("test")
+
+            observe_params = trace_info.get_observe_kwargs()
+            observe_params["name"] = observe_name
+
+            # Add the new Laminar parameters
+            if _ignore_input:
+                observe_params["ignore_input"] = _ignore_input
+            if _ignore_output:
+                observe_params["ignore_output"] = _ignore_output
+            if _ignore_inputs is not None:
+                observe_params["ignore_inputs"] = _ignore_inputs
+            if _input_formatter is not None:
+                observe_params["input_formatter"] = _input_formatter
+            if _output_formatter is not None:
+                observe_params["output_formatter"] = _output_formatter
+
+            return observe_params
+
+        # --- The actual wrappers ---
+        @wraps(f)
+        def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            observe_params = _prepare_and_get_observe_params(kwargs)
+            observed_func = _observe(**observe_params)(f)
+            return observed_func(*args, **kwargs)
+
+        @wraps(f)
+        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            observe_params = _prepare_and_get_observe_params(kwargs)
+            observed_func = _observe(**observe_params)(f)
+            return await observed_func(*args, **kwargs)
+
+        wrapper = async_wrapper if is_coroutine else sync_wrapper
+
+        # Preserve the original function signature
+        try:
+            wrapper.__signature__ = sig  # type: ignore[attr-defined]
+        except (AttributeError, ValueError):
+            pass
+
+        return cast(Callable[P, R], wrapper)
+
+    if func:
+        return decorator(func)  # Called as @trace
+    else:
+        return decorator  # Called as @trace(...)

ai_pipeline_core-0.1.1.dist-info/METADATA

@@ -0,0 +1,477 @@
+Metadata-Version: 2.4
+Name: ai-pipeline-core
+Version: 0.1.1
+Summary: Core utilities for AI-powered processing pipelines using prefect
+Project-URL: Homepage, https://github.com/bbarwik/ai-pipeline-core
+Project-URL: Repository, https://github.com/bbarwik/ai-pipeline-core
+Project-URL: Issues, https://github.com/bbarwik/ai-pipeline-core/issues
+Author-email: bbarwik <bbarwik@gmail.com>
+License: MIT
+License-File: LICENSE
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Typing :: Typed
+Requires-Python: >=3.12
+Requires-Dist: httpx>=0.28.1
+Requires-Dist: jinja2>=3.1.6
+Requires-Dist: lmnr>=0.7.4
+Requires-Dist: openai>=1.99.9
+Requires-Dist: prefect>=3.4.13
+Requires-Dist: pydantic-settings>=2.10.1
+Requires-Dist: pydantic>=2.11.7
+Requires-Dist: python-magic>=0.4.27
+Requires-Dist: ruamel-yaml>=0.18.14
+Requires-Dist: tiktoken>=0.11.0
+Provides-Extra: dev
+Requires-Dist: basedpyright>=1.31.2; extra == 'dev'
+Requires-Dist: bump2version>=1.0.1; extra == 'dev'
+Requires-Dist: pre-commit>=4.3.0; extra == 'dev'
+Requires-Dist: pytest-asyncio>=1.1.0; extra == 'dev'
+Requires-Dist: pytest-cov>=5.0.0; extra == 'dev'
+Requires-Dist: pytest-mock>=3.14.0; extra == 'dev'
+Requires-Dist: pytest-xdist>=3.8.0; extra == 'dev'
+Requires-Dist: pytest>=8.4.1; extra == 'dev'
+Requires-Dist: ruff>=0.12.9; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# AI Pipeline Core
+
+A high-performance, type-safe Python library for building AI-powered data processing pipelines with Prefect orchestration and LMNR observability.
+
+[![Python Version](https://img.shields.io/badge/python-3.12%2B-blue)](https://www.python.org/downloads/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![Code Style: Ruff](https://img.shields.io/badge/code%20style-ruff-000000.svg)](https://github.com/astral-sh/ruff)
+[![Type Checked: Basedpyright](https://img.shields.io/badge/type%20checked-basedpyright-blue)](https://github.com/DetachHead/basedpyright)
+[![Test Coverage](https://img.shields.io/badge/coverage-80%25-green)](https://github.com/bbarwik/ai-pipeline-core)
+[![Status: Beta](https://img.shields.io/badge/status-beta-yellow)](https://github.com/bbarwik/ai-pipeline-core)
+[![PyPI version](https://img.shields.io/pypi/v/ai-pipeline-core.svg)](https://pypi.org/project/ai-pipeline-core/)
+[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/ai-pipeline-core.svg)](https://pypi.org/project/ai-pipeline-core/)
+
+> [!NOTE]
+> **Beta Release**
+>
+> This library is in beta. While actively used in production systems, the API may still evolve. We follow semantic versioning for releases.
+
+## Overview
+
+AI Pipeline Core provides a robust foundation for building production-grade AI pipelines with a focus on:
+
+- **100% Async Architecture** - Built for high-throughput, non-blocking operations
+- **Type Safety** - Comprehensive type hints with Pydantic models throughout
+- **Minimal Design** - Every line of code justified, no unnecessary abstractions
+- **Production Ready** - Built-in retry logic, caching, monitoring, and error handling
+- **LLM Optimization** - Smart context/message splitting for efficient token usage
+
+## Key Features
+
+### 🚀 Performance First
+- Fully asynchronous I/O operations
+- Intelligent caching for LLM context
+- Streaming support for large documents
+- Automatic retry with exponential backoff
+
+### 🔒 Type Safety
+- Pydantic models for all data structures
+- Strict type checking with basedpyright
+- Runtime validation for all inputs
+- Immutable configurations by default
+
+### 📊 Observability
+- LMNR (Laminar) tracing integration
+- Structured logging with Prefect
+- Cost tracking for LLM operations
+- Performance metrics out of the box
+
+### 🎯 Developer Experience
+- Self-documenting code for experienced developers
+- Consistent patterns throughout
+- Comprehensive error messages
+- Smart defaults with override capabilities
+
+### 🤖 Advanced LLM Features
+- Search-enabled models (Perplexity Sonar, Gemini Flash Search)
+- Reasoning models support (O1 series)
+- Structured output with Pydantic models
+- Dynamic model selection based on task
+
+## Installation
+
+```bash
+pip install ai-pipeline-core
+```
+
+### Development Installation
+
+For contributors and development:
+
+```bash
+git clone https://github.com/bbarwik/ai-pipeline-core.git
+cd ai-pipeline-core
+pip install -e ".[dev]"
+make install-dev  # Installs pre-commit hooks
+```
+
+### Requirements
+- Python 3.12 or higher
+- Linux/macOS (Windows via WSL2)
+
+## Quick Start
+
+### Basic Document Processing
+```python
+from ai_pipeline_core.documents import Document, FlowDocument
+from ai_pipeline_core.llm import generate_structured, AIMessages, ModelOptions
+from pydantic import BaseModel
+
+class InputDocument(FlowDocument):
+    """Custom document type for your flow"""
+    def get_type(self) -> str:
+        return "input"
+
+class AnalysisResult(BaseModel):
+    """Example Pydantic model for structured output"""
+    summary: str
+    key_points: list[str]
+
+async def process_document(doc: Document):
+    # Generate AI response with structured output
+    response = await generate_structured(
+        model="gemini-2.5-pro",  # Model is required first parameter
+        response_format=AnalysisResult,  # Pydantic model class
+        context=AIMessages([doc]),  # Cached context
+        messages=AIMessages(["Analyze this document"]),  # Dynamic messages
+        options=ModelOptions(max_completion_tokens=5000)  # Optional options
+    )
+    return response.parsed
+```
+
+### Prefect Flow Integration
+```python
+from prefect import flow, task
+from ai_pipeline_core.documents import Document, DocumentList, FlowDocument
+from ai_pipeline_core.flow import FlowConfig
+from ai_pipeline_core.tracing import trace
+
+class OutputDocument(FlowDocument):
+    """Custom output document type"""
+    def get_type(self) -> str:
+        return "output"
+
+class MyFlowConfig(FlowConfig):
+    INPUT_DOCUMENT_TYPES = [InputDocument]
+    OUTPUT_DOCUMENT_TYPE = OutputDocument
+
+@task
+@trace
+async def process_task(doc: Document) -> Document:
+    # Task-level processing with automatic tracing
+    result = await process_document(doc)
+    # Convert result to JSON string for document content
+    import json
+    return OutputDocument(name="result", content=json.dumps(result.model_dump()).encode())
+
+@flow
+async def my_pipeline(documents: DocumentList):
+    config = MyFlowConfig()
+    input_docs = config.get_input_documents(documents)
+
+    results = await process_task.map(input_docs)
+
+    config.validate_output_documents(results)
+    return results
+```
+
+## Core Modules
+
+### Documents System
+The foundation for all data handling. Documents are immutable, type-safe wrappers around content with automatic MIME type detection.
+
+```python
+from ai_pipeline_core.documents import Document, DocumentList
+
+# Documents handle encoding/decoding automatically
+doc = MyDocument(
+    name="report.pdf",
+    content=pdf_bytes,
+    description="Q3 Financial Report"
+)
+
+# Type-safe document collections
+docs = DocumentList([doc1, doc2])
+```
+
+### LLM Module
+Managed AI interactions with built-in retry logic, cost tracking, and structured outputs.
+
+**Supported Models** (via LiteLLM proxy):
+- OpenAI: GPT-4, GPT-5 series
+- Anthropic: Claude 3 series
+- Google: Gemini 2.5 series
+- xAI: Grok models
+- Perplexity: Sonar models (with search capabilities)
+- And many more through LiteLLM compatibility
+
+```python
+from ai_pipeline_core.llm import generate_structured, AIMessages, ModelOptions
+from pydantic import BaseModel
+
+class YourPydanticModel(BaseModel):
+    field1: str
+    field2: int
+
+# Get structured Pydantic model responses
+result = await generate_structured(
+    model="gemini-2.5-pro",  # Model is required first parameter
+    response_format=YourPydanticModel,  # Pydantic model class for structured output
+    context=AIMessages(),  # Optional context (cached)
+    messages=AIMessages(["Your prompt here"]),  # Required messages
+    options=ModelOptions(
+        retries=3,
+        timeout=30,
+        max_completion_tokens=10000
+    )
+)
+# Access the parsed result
+model_instance = result.parsed  # Type: YourPydanticModel
+```
+
+### Prompt Management
+Flexible Jinja2-based prompt system with smart path resolution.
+
+```python
+from ai_pipeline_core import PromptManager
+
+pm = PromptManager(__file__)
+prompt = pm.get("analyze_document.jinja2",
+                 document=doc,
+                 instructions=instructions)
+```
+
+### Tracing & Monitoring
+Automatic observability with LMNR integration.
+
+```python
+from ai_pipeline_core.tracing import trace
+
+@trace(metadata={"workflow": "analysis"})
+async def analyze_data(data: InputData) -> OutputData:
+    # Automatic tracing with performance metrics
+    ...
+```
+
+## Architecture Principles
+
+### 1. Async-First Design
+Every I/O operation is asynchronous. No blocking calls, no synchronous fallbacks.
+
+### 2. Type Safety Throughout
+Complete type annotations with runtime validation. If it compiles, it works.
+
+### 3. Minimal Surface Area
+Less code is better code. Every line must justify its existence.
+
+### 4. Configuration as Code
+All configurations are Pydantic models - validated, typed, and immutable.
+
+## Project Structure
+
+```
+ai_pipeline_core/
+├── documents/          # Document handling system
+│   ├── document.py     # Base document class
+│   ├── flow_document.py # Prefect flow documents
+│   └── task_document.py # Prefect task documents
+├── llm/               # LLM interaction layer
+│   ├── client.py      # Async client implementation
+│   └── model_options.py # Configuration models
+├── flow/              # Prefect flow utilities
+│   └── config.py      # Type-safe flow configuration
+├── logging/           # Structured logging
+├── tracing.py         # Observability decorators
+└── settings.py        # Centralized configuration
+```
+
+## Development
+
+### Running Tests
+```bash
+make test           # Run all tests
+make test-cov      # Run with coverage report
+pytest tests/test_documents.py::TestDocument::test_creation  # Single test
+```
+
+### Code Quality
+```bash
+make lint          # Run linting checks
+make format        # Auto-format code
+make typecheck     # Run type checking
+make pre-commit    # Run all pre-commit hooks
+```
+
+### Development Workflow
+1. Create feature branch
+2. Write tests first (TDD)
+3. Implement minimal solution
+4. Run `make format` and `make typecheck`
+5. Ensure >80% test coverage
+6. Submit PR with clear description
+
+## Best Practices
+
+### DO ✅
+- Use async/await for all I/O operations
+- Define Pydantic models for all data structures
+- Keep functions under 20 lines
+- Use type hints for everything
+- Let Documents handle serialization
+
+### DON'T ❌
+- Import `logging` directly (use pipeline logger)
+- Use raw dictionaries for configuration
+- Write defensive code for unlikely scenarios
+- Add comments explaining what (code should be clear)
+- Use `requests` or other blocking libraries
+
+## Configuration
+
+### Environment Variables
+```bash
+# Required for LLM operations
+OPENAI_API_KEY=sk-...  # Your OpenAI or LiteLLM proxy key
+OPENAI_BASE_URL=http://your-proxy:8000  # LiteLLM proxy endpoint
+
+# Optional - for observability
+LMNR_PROJECT_API_KEY=lmnr_...  # LMNR tracing
+
+# Optional - for orchestration
+PREFECT_API_URL=http://localhost:4200/api
+AI_PIPELINE_LOG_LEVEL=INFO
+```
+
+### Settings Management
+```python
+from ai_pipeline_core.settings import settings
+
+# All settings are validated Pydantic models
+api_key = settings.openai_api_key
+base_url = settings.openai_base_url  # LiteLLM proxy endpoint
+```
+
+## Integration Examples
+
+### With Prefect Cloud
+```python
+from prefect import flow
+from ai_pipeline_core.flow import FlowConfig
+
+@flow(name="document-processor")
+async def process_documents(docs: DocumentList):
+    # Automatic Prefect Cloud integration
+    ...
+```
+
+### With Custom LLM Providers
+```python
+from ai_pipeline_core.settings import settings
+
+# Configure LiteLLM proxy endpoint via environment variables
+# OPENAI_BASE_URL=http://your-litellm-proxy:8000
+# OPENAI_API_KEY=your-proxy-key
+
+# Access in code (settings are immutable)
+base_url = settings.openai_base_url
+```
+
+## Performance Considerations
+
+- **Context Caching**: The LLM module automatically caches context to reduce token usage
+- **Document Streaming**: Large documents are streamed rather than loaded entirely into memory
+- **Batch Processing**: Use Prefect's `.map()` for parallel task execution
+- **Connection Pooling**: HTTP clients use connection pooling by default
+
+## Troubleshooting
+
+### Common Issues
+
+1. **Import Errors**: Ensure Python 3.12+ is installed
+2. **Async Warnings**: All I/O operations must use `await`
+3. **Type Errors**: Run `make typecheck` to identify issues
+4. **MIME Detection**: Install `python-magic` system dependencies
+
+### Debug Mode
+```python
+from ai_pipeline_core.logging import setup_logging, LoggingConfig
+
+# Setup logging with DEBUG level
+setup_logging(LoggingConfig(level="DEBUG"))
+```
+
+## Release Process
+
+See [RELEASE.md](RELEASE.md) for detailed release procedures.
+
+**Important**: All releases require:
+- ✅ Zero errors from `make typecheck`
+- ✅ All unit tests passing with >80% coverage
+- ✅ **Integration tests passing** (with configured API keys)
+
+## Contributing
+
+> [!NOTE]
+> As this is a preview repository used internally, we are not actively accepting external contributions. The codebase may change significantly without notice.
+>
+> **Recommended approach:**
+> 1. Fork the repository
+> 2. Make changes in your fork
+> 3. Share your improvements with the community through your fork
+
+If you've found a critical security issue, please report it via the GitHub Security tab.
+
+For learning purposes, see [CLAUDE.md](CLAUDE.md) for our comprehensive coding standards and architecture guide.
+
+## Documentation
+
+- [CLAUDE.md](CLAUDE.md) - Detailed coding standards and architecture guide
+- [Prefect Integration](docs/prefect.md) - Prefect patterns and best practices
+- [Deployment Guide](docs/prefect_deployment.md) - Production deployment
+- [Prefect Logging](docs/prefect_logging.md) - Logging configuration guide
+
+## License
+
+MIT License - see [LICENSE](LICENSE) file for details.
+
+## Support
+
+> [!CAUTION]
+> This is a preview repository with no guaranteed support. Issues and discussions may not be actively monitored.
+
+- **For Learning**: Review the code, documentation, and examples
+- **For Usage**: Fork the repository and maintain your own version
+- **Security Issues**: Report via GitHub Security tab
+
+## Acknowledgments
+
+Built with:
+- [Prefect](https://www.prefect.io/) - Workflow orchestration
+- [LMNR](https://www.lmnr.ai/) - LLM observability
+- [LiteLLM](https://litellm.ai/) - LLM proxy
+- [Pydantic](https://pydantic-docs.helpmanual.io/) - Data validation
+
+## Stability Notice
+
+**Current Version**: 0.1.1
+**Status**: Internal Preview
+**API Stability**: Unstable - Breaking changes expected
+**Recommended Use**: Learning and reference only
+
+For production use, please fork this repository and maintain your own stable version.
+
+---
+
+**Remember**: The best code is no code. The second best is minimal, typed, async code that does exactly what's needed.