ai-pipeline-core 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

ai_pipeline_core/tracing.py

@@ -11,7 +11,7 @@ from __future__ import annotations
 import inspect
 import os
 from functools import wraps
-from typing import Any, Callable, ParamSpec, TypeVar, cast, overload
+from typing import Any, Callable, Literal, ParamSpec, TypeVar, cast, overload
 
 from lmnr import Instruments, Laminar, observe
 from pydantic import BaseModel
@@ -24,6 +24,8 @@ from ai_pipeline_core.settings import settings
 P = ParamSpec("P")
 R = TypeVar("R")
 
+TraceLevel = Literal["always", "debug", "off"]
+
 
 # ---------------------------------------------------------------------------
 # ``TraceInfo`` – metadata container
@@ -67,22 +69,28 @@ def _initialise_laminar() -> None:
     if settings.lmnr_project_api_key:
         Laminar.initialize(
             project_api_key=settings.lmnr_project_api_key,
-            disabled_instruments=[Instruments.OPENAI],
+            disabled_instruments=[Instruments.OPENAI] if Instruments.OPENAI else [],
         )
 
 
-# Overload for calls like @trace(name="...", test=True)
+# Overload for calls like @trace(name="...", level="debug")
 @overload
 def trace(
     *,
+    level: TraceLevel = "always",
     name: str | None = None,
-    test: bool = False,
-    debug_only: bool = False,
+    session_id: str | None = None,
+    user_id: str | None = None,
+    metadata: dict[str, Any] | None = None,
+    tags: list[str] | None = None,
+    span_type: str | None = None,
     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,
+    ignore_exceptions: bool = False,
+    preserve_global_context: bool = True,
 ) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
 
 
@@ -95,56 +103,76 @@ def trace(func: Callable[P, R]) -> Callable[P, R]: ...
 def trace(
     func: Callable[P, R] | None = None,
     *,
+    level: TraceLevel = "always",
     name: str | None = None,
-    test: bool = False,
-    debug_only: bool = False,
+    session_id: str | None = None,
+    user_id: str | None = None,
+    metadata: dict[str, Any] | None = None,
+    tags: list[str] | None = None,
+    span_type: str | None = None,
     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,
+    ignore_exceptions: bool = False,
     preserve_global_context: bool = True,
 ) -> 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)
+        level: Trace level control:
+            - "always": Always trace (default)
+            - "debug": Only trace when LMNR_DEBUG environment variable is NOT set to "true"
+            - "off": Never 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
+        metadata: Additional metadata for the trace
+        tags: Additional tags for the trace
+        span_type: Type of span for the trace
         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)
+        ignore_exceptions: Whether to ignore exceptions in tracing
+        preserve_global_context: Whether to preserve global context
 
     Returns:
         The decorated function with Laminar tracing enabled
     """
 
+    if level == "off":
+        if func:
+            return func
+        return lambda f: f
+
     def decorator(f: Callable[P, R]) -> Callable[P, R]:
+        # Handle 'debug' level logic - only trace when LMNR_DEBUG is NOT "true"
+        if level == "debug" and os.getenv("LMNR_DEBUG", "").lower() == "true":
+            return f
+
         # --- 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
+        _session_id = session_id
+        _user_id = user_id
+        _metadata = metadata
+        _tags = tags or []
+        _span_type = span_type
         _ignore_input = ignore_input
         _ignore_output = ignore_output
         _ignore_inputs = ignore_inputs
         _input_formatter = input_formatter
         _output_formatter = output_formatter
+        _ignore_exceptions = ignore_exceptions
         _preserve_global_context = preserve_global_context
 
-        # --- 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]:
             """
@@ -157,13 +185,23 @@ def trace(
                 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
 
+            # Override with decorator-level session_id and user_id if provided
+            if _session_id:
+                observe_params["session_id"] = _session_id
+            if _user_id:
+                observe_params["user_id"] = _user_id
+
+            # Merge decorator-level metadata and tags
+            if _metadata:
+                observe_params["metadata"] = {**observe_params.get("metadata", {}), **_metadata}
+            if _tags:
+                observe_params["tags"] = observe_params.get("tags", []) + _tags
+            if _span_type:
+                observe_params["span_type"] = _span_type
+
             # Add the new Laminar parameters
             if _ignore_input:
                 observe_params["ignore_input"] = _ignore_input
@@ -175,6 +213,8 @@ def trace(
                 observe_params["input_formatter"] = _input_formatter
             if _output_formatter is not None:
                 observe_params["output_formatter"] = _output_formatter
+            if _ignore_exceptions:
+                observe_params["ignore_exceptions"] = _ignore_exceptions
             if _preserve_global_context:
                 observe_params["preserve_global_context"] = _preserve_global_context
 
@@ -207,3 +247,6 @@ def trace(
         return decorator(func)  # Called as @trace
     else:
         return decorator  # Called as @trace(...)
+
+
+__all__ = ["trace", "TraceLevel", "TraceInfo"]

{ai_pipeline_core-0.1.5.dist-info → ai_pipeline_core-0.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-pipeline-core
-Version: 0.1.5
+Version: 0.1.7
 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
@@ -20,7 +20,7 @@ Classifier: Typing :: Typed
 Requires-Python: >=3.12
 Requires-Dist: httpx>=0.28.1
 Requires-Dist: jinja2>=3.1.6
-Requires-Dist: lmnr>=0.7.5
+Requires-Dist: lmnr>=0.7.6
 Requires-Dist: openai>=1.99.9
 Requires-Dist: prefect>=3.4.13
 Requires-Dist: pydantic-settings>=2.10.1
@@ -151,40 +151,76 @@ async def process_document(doc: Document):
     return response.parsed
 ```
 
-### Prefect Flow Integration
+### Enhanced Pipeline Decorators (New in v0.1.7)
 ```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
+from ai_pipeline_core import pipeline_flow, pipeline_task
+from ai_pipeline_core.flow import FlowOptions
+from ai_pipeline_core.documents import DocumentList, FlowDocument
 
-class OutputDocument(FlowDocument):
-    """Custom output document type"""
-    def get_type(self) -> str:
-        return "output"
+class CustomFlowOptions(FlowOptions):
+    """Extend base options with your custom fields"""
+    batch_size: int = 100
+    temperature: float = 0.7
 
-class MyFlowConfig(FlowConfig):
-    INPUT_DOCUMENT_TYPES = [InputDocument]
-    OUTPUT_DOCUMENT_TYPE = OutputDocument
-
-@task
-@trace
+@pipeline_task(trace_level="always", retries=3)
 async def process_task(doc: Document) -> Document:
-    # Task-level processing with automatic tracing
+    # Task with automatic tracing and retries
     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())
+    return OutputDocument(name="result", content=result.encode())
+
+@pipeline_flow(trace_level="always")
+async def my_pipeline(
+    project_name: str,
+    documents: DocumentList,
+    flow_options: CustomFlowOptions  # Type-safe custom options
+) -> DocumentList:
+    # Pipeline flow with enforced signature and tracing
+    results = []
+    for doc in documents:
+        result = await process_task(doc)
+        results.append(result)
+    return DocumentList(results)
+```
 
-@flow
-async def my_pipeline(documents: DocumentList):
-    config = MyFlowConfig()
-    input_docs = config.get_input_documents(documents)
+### Simple Runner Utility (New in v0.1.7)
+```python
+from ai_pipeline_core.simple_runner import run_cli, run_pipeline
+from ai_pipeline_core.flow import FlowOptions
+
+# CLI-based pipeline execution
+if __name__ == "__main__":
+    run_cli(
+        flows=[my_pipeline],
+        flow_configs=[MyFlowConfig],
+        options_cls=CustomFlowOptions
+    )
 
-    results = await process_task.map(input_docs)
+# Or programmatic execution
+async def main():
+    result = await run_pipeline(
+        project_name="my-project",
+        output_dir=Path("./output"),
+        flow=my_pipeline,
+        flow_config=MyFlowConfig,
+        flow_options=CustomFlowOptions(batch_size=50)
+    )
+```
 
-    config.validate_output_documents(results)
-    return results
+### Clean Prefect Decorators (New in v0.1.7)
+```python
+# Import clean Prefect decorators without tracing
+from ai_pipeline_core.prefect import flow, task
+
+# Or use pipeline decorators with tracing
+from ai_pipeline_core import pipeline_flow, pipeline_task
+
+@task  # Clean Prefect task
+def compute(x: int) -> int:
+    return x * 2
+
+@pipeline_task(trace_level="always")  # With tracing
+def compute_traced(x: int) -> int:
+    return x * 2
 ```
 
 ## Core Modules
@@ -291,8 +327,14 @@ ai_pipeline_core/
 │   ├── client.py      # Async client implementation
 │   └── model_options.py # Configuration models
 ├── flow/              # Prefect flow utilities
-│   └── config.py      # Type-safe flow configuration
+│   ├── config.py      # Type-safe flow configuration
+│   └── options.py     # FlowOptions base class (v0.1.7)
+├── simple_runner/     # Pipeline execution utilities (v0.1.7)
+│   ├── cli.py         # CLI interface
+│   └── simple_runner.py # Core runner logic
 ├── logging/           # Structured logging
+├── pipeline.py        # Enhanced decorators (v0.1.7)
+├── prefect.py         # Clean Prefect exports (v0.1.7)
 ├── tracing.py         # Observability decorators
 └── settings.py        # Centralized configuration
 ```
@@ -469,9 +511,29 @@ Built with:
 - [LiteLLM](https://litellm.ai/) - LLM proxy
 - [Pydantic](https://pydantic-docs.helpmanual.io/) - Data validation
 
+## What's New in v0.1.7
+
+### Major Additions
+- **Enhanced Pipeline Decorators**: New `pipeline_flow` and `pipeline_task` decorators combining Prefect functionality with automatic LMNR tracing
+- **FlowOptions Base Class**: Extensible configuration system for flows with type-safe inheritance
+- **Simple Runner Module**: CLI and programmatic utilities for easy pipeline execution
+- **Clean Prefect Exports**: Separate imports for Prefect decorators with and without tracing
+- **Expanded Exports**: All major components now accessible from top-level package import
+
+### API Improvements
+- Better type inference for document flows with custom options
+- Support for custom FlowOptions inheritance in pipeline flows
+- Improved error messages for invalid flow signatures
+- Enhanced document utility functions (`canonical_name_key`, `sanitize_url`)
+
+### Developer Experience
+- Simplified imports - most components available from `ai_pipeline_core` directly
+- Better separation of concerns between clean Prefect and traced pipeline decorators
+- More intuitive flow configuration with `FlowOptions` inheritance
+
 ## Stability Notice
 
-**Current Version**: 0.1.5
+**Current Version**: 0.1.7
 **Status**: Internal Preview
 **API Stability**: Unstable - Breaking changes expected
 **Recommended Use**: Learning and reference only

{ai_pipeline_core-0.1.5.dist-info → ai_pipeline_core-0.1.7.dist-info}/RECORD

@@ -1,21 +1,24 @@
-ai_pipeline_core/__init__.py,sha256=xjZh6D4fkepTm3LK42qaTpNHgN2meK4uthrW9pjwLfE,779
+ai_pipeline_core/__init__.py,sha256=INcTtHr2TFY8bR0eCg7RwvIRYY6px8knCgjyIvSSKP4,1602
 ai_pipeline_core/exceptions.py,sha256=_vW0Hbw2LGb5tcVvH0YzTKMff7QOPfCRr3w-w_zPyCE,968
+ai_pipeline_core/pipeline.py,sha256=GOrPC53j756Xhpg_CShnkAKxSdkC16XHEoPeIhkjLIA,16569
+ai_pipeline_core/prefect.py,sha256=VHYkkRcUmSpdwyWosOOxuExVCncIQgT6MypqGdjcYnM,241
 ai_pipeline_core/prompt_manager.py,sha256=XmNUdMIC0WrE9fF0LIcfozAKOGrlYwj8AfXvCndIH-o,4693
 ai_pipeline_core/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ai_pipeline_core/settings.py,sha256=Zl2BPa6IHzh-B5V7cg5mtySr1dhWZQYYKxXz3BwrHlQ,615
-ai_pipeline_core/tracing.py,sha256=_bijptKWXh7V_xENFQGF11-B70rOGwV6g0qdBoF-VCw,7890
-ai_pipeline_core/documents/__init__.py,sha256=rEnKj-sSlZ9WnFlZAmSGVi1P8vnsHmU9O9_YwtP40ms,242
-ai_pipeline_core/documents/document.py,sha256=cKa9v0GpDYh48Aqj0u_tCwh4_uCcn5YsKzC537f0v0I,10089
+ai_pipeline_core/tracing.py,sha256=T-3fTyA37TejXxotkVzTNqL2a5nOfZ0bcHg9TClLvmg,9471
+ai_pipeline_core/documents/__init__.py,sha256=TLW8eOEmthfDHOTssXjyBlqhgrZe9ZIyxlkd0LBJ3_s,340
+ai_pipeline_core/documents/document.py,sha256=e3IBr0TThucBAaOHvdqv0X--iCcBrqh2jzFTyaOp7O0,12418
 ai_pipeline_core/documents/document_list.py,sha256=HOG_uZDazA9CJB7Lr_tNcDFzb5Ff9RUt0ELWQK_eYNM,4940
 ai_pipeline_core/documents/flow_document.py,sha256=qsV-2JYOMhkvAj7lW54ZNH_4QUclld9h06CoU59tWww,815
-ai_pipeline_core/documents/mime_type.py,sha256=tMWGH9PVmHe6a_IzdaJUqIHf4qnwQOwOCBhsgW2AyTE,2244
+ai_pipeline_core/documents/mime_type.py,sha256=sBhNRoBJQ35JoHWhJzBGpp00WFDfMdEX0JZKKkR7QH0,3371
 ai_pipeline_core/documents/task_document.py,sha256=WjHqtl1d60XFBBqewNRdz1OqBErGI0jRx15oQYCTHo8,907
 ai_pipeline_core/documents/utils.py,sha256=BdE4taSl1vrBhxnFbOP5nDA7lXIcvY__AMRTHoaNb5M,2764
-ai_pipeline_core/flow/__init__.py,sha256=_Sji2yY1ICkvVX6QiiGWKzqIXtg9UAiuvhjHSK_gdO8,57
+ai_pipeline_core/flow/__init__.py,sha256=54DRfZnjXQVrimgtKEVEm5u5ErImx31cjK2PpBvHjU4,116
 ai_pipeline_core/flow/config.py,sha256=crbe_OvNE6qulIKv1D8yKoe8xrEsIlvICyxjhqHHBxQ,2266
+ai_pipeline_core/flow/options.py,sha256=WygJEwjqOa14l23a_Hp36hJX-WgxHMq-YzSieC31Z4Y,701
 ai_pipeline_core/llm/__init__.py,sha256=3XVK-bSJdOe0s6KmmO7PDbsXHfjlcZEG1MVBmaz3EeU,442
 ai_pipeline_core/llm/ai_messages.py,sha256=DwJJe05BtYdnMZeHbBbyEbDCqrW63SRvprxptoJUCn4,4586
-ai_pipeline_core/llm/client.py,sha256=IOcyjwyAKQWlqnwC5p2Hl4FeRCzOJAHC5Yqr_oCBQ8s,7703
+ai_pipeline_core/llm/client.py,sha256=VMs1nQKCfoxbcvE2mypn5QF19u90Ua87-5IiZxWOj98,7784
 ai_pipeline_core/llm/model_options.py,sha256=TvAAlDFZN-TP9-J-RZBuU_dpSocskf6paaQMw1XY9UE,1321
 ai_pipeline_core/llm/model_response.py,sha256=fIWueaemgo0cMruvToMZyKsRPzKwL6IlvUJN7DLG710,5558
 ai_pipeline_core/llm/model_types.py,sha256=rIwY6voT8-xdfsKPDC0Gkdl2iTp9Q2LuvWGSRU9Mp3k,342
@@ -23,7 +26,10 @@ ai_pipeline_core/logging/__init__.py,sha256=DOO6ckgnMVXl29Sy7q6jhO-iW96h54pCHQDz
 ai_pipeline_core/logging/logging.yml,sha256=YTW48keO_K5bkkb-KXGM7ZuaYKiquLsjsURei8Ql0V4,1353
 ai_pipeline_core/logging/logging_config.py,sha256=6MBz9nnVNvqiLDoyy9-R3sWkn6927Re5hdz4hwTptpI,4903
 ai_pipeline_core/logging/logging_mixin.py,sha256=RDaR2ju2-vKTJRzXGa0DquGPT8_UxahWjvKJnaD0IV8,7810
-ai_pipeline_core-0.1.5.dist-info/METADATA,sha256=U1OIPjGwAGsuyJ3QnhUtJQWMzj-OqkXDuyH6cW8Dq70,15869
-ai_pipeline_core-0.1.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-ai_pipeline_core-0.1.5.dist-info/licenses/LICENSE,sha256=kKj8mfbdWwkyG3U6n7ztB3bAZlEwShTkAsvaY657i3I,1074
-ai_pipeline_core-0.1.5.dist-info/RECORD,,
+ai_pipeline_core/simple_runner/__init__.py,sha256=OPbTCZvqpnYdwi1Knnkj-MpmD0Nvtg5O7UwIdAKz_AY,384
+ai_pipeline_core/simple_runner/cli.py,sha256=TjiSh7lr1VnTbO1jA2DuVzC2AA6V_5sA5Z8XSuldQmc,3054
+ai_pipeline_core/simple_runner/simple_runner.py,sha256=70BHT1iz-G368H2t4tsWAVni0jw2VkWVdnKICuVtLPw,5009
+ai_pipeline_core-0.1.7.dist-info/METADATA,sha256=2Pi815TCTBlKnTp2duTaUJiKaextafqZ5yfPZdD_--o,18361
+ai_pipeline_core-0.1.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ai_pipeline_core-0.1.7.dist-info/licenses/LICENSE,sha256=kKj8mfbdWwkyG3U6n7ztB3bAZlEwShTkAsvaY657i3I,1074
+ai_pipeline_core-0.1.7.dist-info/RECORD,,