ai-pipeline-core 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

ai_pipeline_core/__init__.py

@@ -118,7 +118,7 @@ from .prompt_manager import PromptManager
 from .settings import Settings
 from .tracing import TraceInfo, TraceLevel, set_trace_cost, trace
 
-__version__ = "0.2.5"
+__version__ = "0.2.7"
 
 __all__ = [
     # Config/Settings

ai_pipeline_core/llm/ai_messages.py

@@ -260,11 +260,14 @@ class AIMessages(list[AIMessageType]):
 
         for message in self:
             if isinstance(message, str):
-                messages.append({"role": "user", "content": message})
+                messages.append({"role": "user", "content": [{"type": "text", "text": message}]})
             elif isinstance(message, Document):
                 messages.append({"role": "user", "content": AIMessages.document_to_prompt(message)})
             elif isinstance(message, ModelResponse):  # type: ignore
-                messages.append({"role": "assistant", "content": message.content})
+                messages.append({
+                    "role": "assistant",
+                    "content": [{"type": "text", "text": message.content}],
+                })
             else:
                 raise ValueError(f"Unsupported message type: {type(message)}")
 

ai_pipeline_core/llm/client.py

@@ -17,7 +17,7 @@ from typing import Any, TypeVar
 
 from lmnr import Laminar
 from openai import AsyncOpenAI
-from openai.lib.streaming.chat import ContentDeltaEvent, ContentDoneEvent
+from openai.lib.streaming.chat import ChunkEvent, ContentDeltaEvent, ContentDoneEvent
 from openai.types.chat import (
     ChatCompletionMessageParam,
 )
@@ -39,7 +39,7 @@ def _process_messages(
     context: AIMessages,
     messages: AIMessages,
     system_prompt: str | None = None,
-    cache_ttl: str | None = "5m",
+    cache_ttl: str | None = "300s",
 ) -> list[ChatCompletionMessageParam]:
     """Process and format messages for LLM API consumption.
 
@@ -51,7 +51,7 @@ def _process_messages(
         context: Messages to be cached (typically expensive/static content).
         messages: Regular messages without caching (dynamic queries).
         system_prompt: Optional system instructions for the model.
-        cache_ttl: Cache TTL for context messages (e.g. "120s", "5m", "1h").
+        cache_ttl: Cache TTL for context messages (e.g. "120s", "300s", "1h").
                    Set to None or empty string to disable caching.
 
     Returns:
@@ -86,12 +86,17 @@ def _process_messages(
         # Use AIMessages.to_prompt() for context
         context_messages = context.to_prompt()
 
-        # Apply caching to last context message if cache_ttl is set
+        # Apply caching to last context message and last content part if cache_ttl is set
         if cache_ttl:
             context_messages[-1]["cache_control"] = {  # type: ignore
                 "type": "ephemeral",
                 "ttl": cache_ttl,
             }
+            assert isinstance(context_messages[-1]["content"], list)  # type: ignore
+            context_messages[-1]["content"][-1]["cache_control"] = {  # type: ignore
+                "type": "ephemeral",
+                "ttl": cache_ttl,
+            }
 
         processed_messages.extend(context_messages)
 
@@ -103,6 +108,42 @@ def _process_messages(
     return processed_messages
 
 
+def _model_name_to_openrouter_model(model: ModelName) -> str:
+    """Convert a model name to an OpenRouter model name.
+
+    Args:
+        model: Model name to convert.
+
+    Returns:
+        OpenRouter model name.
+    """
+    if model == "gpt-4o-search":
+        return "openai/gpt-4o-search-preview"
+    if model == "gemini-2.5-flash-search":
+        return "google/gemini-2.5-flash:online"
+    if model == "grok-4-fast-search":
+        return "x-ai/grok-4-fast:online"
+    if model == "sonar-pro-search":
+        return "perplexity/sonar-reasoning-pro"
+    if model.startswith("gemini"):
+        return f"google/{model}"
+    elif model.startswith("gpt"):
+        return f"openai/{model}"
+    elif model.startswith("grok"):
+        return f"x-ai/{model}"
+    elif model.startswith("claude"):
+        return f"anthropic/{model}"
+    elif model.startswith("qwen3"):
+        return f"qwen/{model}"
+    elif model.startswith("deepseek-"):
+        return f"deepseek/{model}"
+    elif model.startswith("glm-"):
+        return f"z-ai/{model}"
+    elif model.startswith("kimi-"):
+        return f"moonshotai/{model}"
+    return model
+
+
 async def _generate(
     model: str, messages: list[ChatCompletionMessageParam], completion_kwargs: dict[str, Any]
 ) -> ModelResponse:
@@ -128,11 +169,16 @@ async def _generate(
         - Captures response headers for cost tracking
         - Response includes model options for debugging
     """
+    if "openrouter" in settings.openai_base_url.lower():
+        model = _model_name_to_openrouter_model(model)
+
     async with AsyncOpenAI(
         api_key=settings.openai_api_key,
         base_url=settings.openai_base_url,
     ) as client:
-        start_time, first_token_time = time.time(), None
+        start_time = time.time()
+        first_token_time = None
+        usage = None
         async with client.chat.completions.stream(
             model=model,
             messages=messages,
@@ -144,6 +190,9 @@ async def _generate(
                         first_token_time = time.time()
                 elif isinstance(event, ContentDoneEvent):
                     pass
+                elif isinstance(event, ChunkEvent):
+                    if event.chunk.usage:  # used to fix a bug with missing usage data
+                        usage = event.chunk.usage
             if not first_token_time:
                 first_token_time = time.time()
             raw_response = await stream.get_final_completion()
@@ -156,6 +205,7 @@ async def _generate(
             raw_response,
             model_options=completion_kwargs,
             metadata=metadata,
+            usage=usage,
         )
         return response
 
@@ -192,6 +242,10 @@ async def _generate_with_retry(
     if not context and not messages:
         raise ValueError("Either context or messages must be provided")
 
+    if "gemini" in model.lower() and context.approximate_tokens_count < 5000:
+        # Bug fix for minimum explicit context size for Gemini models
+        options.cache_ttl = None
+
     processed_messages = _process_messages(
         context, messages, options.system_prompt, options.cache_ttl
     )
@@ -329,26 +383,11 @@ async def generate(
         ... ])
         >>> response = await llm.generate("gpt-5", messages=messages)
 
-        Configuration via LiteLLM Proxy:
-        >>> # Configure temperature in litellm_config.yaml:
-        >>> # model_list:
-        >>> #   - model_name: gpt-5
-        >>> #     litellm_params:
-        >>> #       model: openai/gpt-4o
-        >>> #       temperature: 0.3
-        >>> #       max_tokens: 1000
-        >>>
-        >>> # Configure retry logic in proxy:
-        >>> # general_settings:
-        >>> #   master_key: sk-1234
-        >>> #   max_retries: 5
-        >>> #   retry_delay: 15
-
     Performance:
         - Context caching saves ~50-90% tokens on repeated calls
         - First call: full token cost
         - Subsequent calls (within cache TTL): only messages tokens
-        - Default cache TTL is 5m (production-optimized)
+        - Default cache TTL is 300s/5 minutes (production-optimized)
         - Default retry logic: 3 attempts with 10s delay (production-optimized)
 
     Caching:

ai_pipeline_core/llm/model_options.py

@@ -45,7 +45,7 @@ class ModelOptions(BaseModel):
 
         timeout: Maximum seconds to wait for response (default: 300).
 
-        cache_ttl: Cache TTL for context messages (default: "5m").
+        cache_ttl: Cache TTL for context messages (default: "300s").
                    String format like "60s", "5m", or None to disable caching.
                    Applied to the last context message for efficient token reuse.
 
@@ -165,7 +165,7 @@ class ModelOptions(BaseModel):
         - search_context_size only works with search models
         - reasoning_effort only works with models that support explicit reasoning
         - response_format is set internally by generate_structured()
-        - cache_ttl accepts formats like "120s", "5m" (default), "1h" or None to disable caching
+        - cache_ttl accepts formats like "120s", "5m", "1h" or None (default: "300s")
         - stop sequences are limited to 4 by most providers
         - user identifier helps track costs per end-user (max 256 chars)
         - extra_body allows passing provider-specific parameters
@@ -179,7 +179,7 @@ class ModelOptions(BaseModel):
     retries: int = 3
     retry_delay_seconds: int = 20
     timeout: int = 600
-    cache_ttl: str | None = "5m"
+    cache_ttl: str | None = "300s"
     service_tier: Literal["auto", "default", "flex", "scale", "priority"] | None = None
     max_completion_tokens: int | None = None
     stop: str | list[str] | None = None

ai_pipeline_core/llm/model_response.py

@@ -11,6 +11,7 @@ from copy import deepcopy
 from typing import Any, Generic, TypeVar
 
 from openai.types.chat import ChatCompletion
+from openai.types.completion_usage import CompletionUsage
 from pydantic import BaseModel
 
 T = TypeVar(
@@ -61,6 +62,7 @@ class ModelResponse(ChatCompletion):
         chat_completion: ChatCompletion,
         model_options: dict[str, Any],
         metadata: dict[str, Any],
+        usage: CompletionUsage | None = None,
     ) -> None:
         """Initialize ModelResponse from ChatCompletion.
 
@@ -73,6 +75,7 @@ class ModelResponse(ChatCompletion):
                           Stored for metadata extraction and tracing.
             metadata: Custom metadata for tracking (time_taken, first_token_time, etc.).
                      Includes timing information and custom tags.
+            usage: Optional usage information from streaming response.
 
         Example:
             >>> # Usually created internally by generate()
@@ -83,10 +86,19 @@ class ModelResponse(ChatCompletion):
             ... )
         """
         data = chat_completion.model_dump()
+
+        # fixes issue where the role is "assistantassistant" instead of "assistant"
+        for i in range(len(data["choices"])):
+            if role := data["choices"][i]["message"].get("role"):
+                if role.startswith("assistant") and role != "assistant":
+                    data["choices"][i]["message"]["role"] = "assistant"
+
         super().__init__(**data)
 
         self._model_options = model_options
         self._metadata = metadata
+        if usage:
+            self.usage = usage
 
     @property
     def content(self) -> str:
@@ -254,7 +266,7 @@ class ModelResponse(ChatCompletion):
 
         other_fields = self.__dict__
         for key, value in other_fields.items():
-            if key in ["_model_options", "_metadata", "choices", "usage"]:
+            if key in ["_model_options", "_metadata", "choices"]:
                 continue
             try:
                 metadata[f"response.raw.{key}"] = json.dumps(value, indent=2, default=str)

ai_pipeline_core/llm/model_types.py

@@ -20,7 +20,7 @@ ModelName: TypeAlias = (
         "grok-4",
         # Small models
         "gemini-2.5-flash",
-        "gpt-5-nano",
+        "gpt-5-mini",
         "grok-4-fast",
         # Search models
         "gemini-2.5-flash-search",

ai_pipeline_core/settings.py

@@ -126,6 +126,10 @@ class Settings(BaseSettings):
     # Prefect Configuration
     prefect_api_url: str = ""
     prefect_api_key: str = ""
+    prefect_api_auth_string: str = ""
+    prefect_work_pool_name: str = "default"
+    prefect_work_queue_name: str = "default"
+    prefect_gcs_bucket: str = ""
 
     # Observability
     lmnr_project_api_key: str = ""
@@ -135,6 +139,4 @@ class Settings(BaseSettings):
     gcs_service_account_file: str = ""  # Path to GCS service account JSON file
 
 
-# Legacy: Module-level instance for backwards compatibility
-# Applications should create their own settings instance
 settings = Settings()

ai_pipeline_core/simple_runner/cli.py

@@ -1,7 +1,5 @@
 """Command-line interface for simple pipeline execution."""
 
-from __future__ import annotations
-
 import asyncio
 import os
 import sys

ai_pipeline_core/tracing.py

@@ -6,8 +6,6 @@ This module centralizes:
   ``observe`` instrumentation, and optional support for test runs.
 """
 
-from __future__ import annotations
-
 import inspect
 import json
 import os

ai_pipeline_core/utils/__init__.py

@@ -0,0 +1,8 @@
+"""Experimental utilities for deployment and remote execution.
+
+These features are experimental and subject to change.
+"""
+
+from .remote_deployment import remote_deployment, run_remote_deployment
+
+__all__ = ["remote_deployment", "run_remote_deployment"]

ai_pipeline_core/utils/deploy.py

@@ -0,0 +1,358 @@
+#!/usr/bin/env python3
+"""Universal Prefect deployment script using Python API.
+
+This script:
+1. Builds a Python package from pyproject.toml
+2. Uploads it to Google Cloud Storage
+3. Creates/updates a Prefect deployment using the RunnerDeployment pattern
+
+Requirements:
+- Settings configured with PREFECT_API_URL and optionally PREFECT_API_KEY
+- Settings configured with PREFECT_GCS_BUCKET
+- pyproject.toml with project name and version
+- Local package installed for flow metadata extraction
+
+Usage:
+    python -m ai_pipeline_core.utils.deploy
+"""
+
+import argparse
+import asyncio
+import subprocess
+import sys
+import tomllib
+import traceback
+from pathlib import Path
+from typing import Any, Optional
+
+from prefect.cli.deploy._storage import _PullStepStorage  # type: ignore
+from prefect.client.orchestration import get_client
+from prefect.deployments.runner import RunnerDeployment
+from prefect.flows import load_flow_from_entrypoint
+
+from ai_pipeline_core.settings import settings
+from ai_pipeline_core.storage import Storage
+
+# ============================================================================
+# Deployer Class
+# ============================================================================
+
+
+class Deployer:
+    """Deploy Prefect flows using the RunnerDeployment pattern.
+
+    This is the official Prefect approach that handles flow registration,
+    deployment creation/updates, and all edge cases automatically.
+    """
+
+    def __init__(self):
+        """Initialize deployer."""
+        self.config = self._load_config()
+        self._validate_prefect_settings()
+
+    def _load_config(self) -> dict[str, Any]:
+        """Load and normalize project configuration from pyproject.toml.
+
+        Returns:
+            Configuration dictionary with project metadata and deployment settings.
+        """
+        if not settings.prefect_gcs_bucket:
+            self._die(
+                "PREFECT_GCS_BUCKET not configured in settings.\n"
+                "Configure via environment variable or .env file:\n"
+                "  PREFECT_GCS_BUCKET=your-bucket-name"
+            )
+
+        pyproject_path = Path("pyproject.toml")
+        if not pyproject_path.exists():
+            self._die("pyproject.toml not found. Run from project root.")
+
+        with open(pyproject_path, "rb") as f:
+            data = tomllib.load(f)
+
+        project = data.get("project", {})
+        name = project.get("name")
+        version = project.get("version")
+
+        if not name:
+            self._die("Project name not found in pyproject.toml")
+        if not version:
+            self._die("Project version not found in pyproject.toml")
+
+        # Normalize naming conventions
+        # Hyphens in package names become underscores in Python imports
+        package_name = name.replace("-", "_")
+        flow_folder = name.replace("_", "-")
+
+        return {
+            "name": name,
+            "package": package_name,
+            "version": version,
+            "bucket": settings.prefect_gcs_bucket,
+            "folder": f"flows/{flow_folder}",
+            "tarball": f"{package_name}-{version}.tar.gz",
+            "work_pool": settings.prefect_work_pool_name,
+            "work_queue": settings.prefect_work_queue_name,
+        }
+
+    def _validate_prefect_settings(self):
+        """Validate that required Prefect settings are configured."""
+        self.api_url = settings.prefect_api_url
+        if not self.api_url:
+            self._die(
+                "PREFECT_API_URL not configured in settings.\n"
+                "Configure via environment variable or .env file:\n"
+                "  PREFECT_API_URL=https://api.prefect.cloud/api/accounts/.../workspaces/..."
+            )
+
+    def _run(self, cmd: str, check: bool = True) -> Optional[str]:
+        """Execute shell command and return output.
+
+        Args:
+            cmd: Shell command to execute
+            check: Whether to raise on non-zero exit code
+
+        Returns:
+            Command stdout if successful, None if failed and check=False
+        """
+        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
+
+        if check and result.returncode != 0:
+            self._die(f"Command failed: {cmd}\n{result.stderr}")
+
+        return result.stdout.strip() if result.returncode == 0 else None
+
+    def _info(self, msg: str):
+        """Print info message."""
+        print(f"→ {msg}")
+
+    def _success(self, msg: str):
+        """Print success message."""
+        print(f"✓ {msg}")
+
+    def _die(self, msg: str):
+        """Print error and exit."""
+        print(f"✗ {msg}", file=sys.stderr)
+        sys.exit(1)
+
+    def _build_package(self) -> Path:
+        """Build Python package using `python -m build`.
+
+        Returns:
+            Path to the built tarball
+        """
+        self._info(f"Building {self.config['name']} v{self.config['version']}")
+
+        # Build sdist (source distribution)
+        build_cmd = "python -m build --sdist"
+
+        self._run(build_cmd)
+
+        # Verify tarball was created
+        tarball_path = Path("dist") / self.config["tarball"]
+        if not tarball_path.exists():
+            self._die(
+                f"Build artifact not found: {tarball_path}\n"
+                f"Expected tarball name: {self.config['tarball']}\n"
+                f"Check that pyproject.toml version matches."
+            )
+
+        self._success(f"Built {tarball_path.name} ({tarball_path.stat().st_size // 1024} KB)")
+        return tarball_path
+
+    async def _upload_package(self, tarball: Path):
+        """Upload package tarball to Google Cloud Storage using Storage abstraction.
+
+        Args:
+            tarball: Path to the tarball to upload
+        """
+        # Extract flow_folder from the config folder path
+        # e.g., "flows/ai-document-writer" -> "ai-document-writer"
+        flow_folder = self.config["folder"].split("/", 1)[1] if "/" in self.config["folder"] else ""
+
+        # Initialize storage with gs://bucket-name/flows and set subfolder to flow_folder
+        base_uri = f"gs://{self.config['bucket']}/flows"
+        storage = await Storage.from_uri(base_uri)
+        storage = storage.with_base(flow_folder)
+
+        dest_uri = storage.url_for(tarball.name)
+        self._info(f"Uploading to {dest_uri}")
+
+        # Read and upload the tarball
+        tarball_bytes = tarball.read_bytes()
+        await storage.write_bytes(tarball.name, tarball_bytes)
+
+        self._success(f"Package uploaded to {self.config['folder']}/{tarball.name}")
+
+    async def _deploy_via_api(self):
+        """Create or update Prefect deployment using RunnerDeployment pattern.
+
+        This is the official Prefect approach that:
+        1. Automatically creates/updates the flow registration
+        2. Handles deployment create vs update logic
+        3. Properly formats all parameters for the API
+        """
+        # Define entrypoint (assumes flow function has same name as package)
+        entrypoint = f"{self.config['package']}:{self.config['package']}"
+
+        # Load flow to get metadata
+        # This requires the package to be installed locally (typical dev workflow)
+        self._info(f"Loading flow from entrypoint: {entrypoint}")
+        try:
+            flow = load_flow_from_entrypoint(entrypoint)
+            self._success(f"Loaded flow: {flow.name}")
+        except ImportError as e:
+            self._die(
+                f"Failed to import flow: {e}\n\n"
+                f"The package must be installed locally to extract flow metadata.\n"
+                f"Install it with: pip install -e .\n\n"
+                f"Expected entrypoint: {entrypoint}\n"
+                f"This means: Python package '{self.config['package']}' "
+                f"with flow function '{self.config['package']}'"
+            )
+        except AttributeError as e:
+            self._die(
+                f"Flow function not found: {e}\n\n"
+                f"Expected flow function named '{self.config['package']}' "
+                f"in package '{self.config['package']}'.\n"
+                f"Check that your flow is decorated with @flow and named correctly."
+            )
+
+        # Define pull steps for workers
+        # These steps tell workers how to get and install the flow code
+        pull_steps = [
+            {
+                "prefect_gcp.deployments.steps.pull_from_gcs": {
+                    "id": "pull_code",
+                    "requires": "prefect-gcp>=0.6",
+                    "bucket": self.config["bucket"],
+                    "folder": self.config["folder"],
+                }
+            },
+            {
+                "prefect.deployments.steps.run_shell_script": {
+                    "id": "install_project",
+                    "stream_output": True,
+                    "directory": "{{ pull_code.directory }}",
+                    # Use uv for fast installation (worker has it installed)
+                    "script": f"uv pip install --system ./{self.config['tarball']}",
+                }
+            },
+        ]
+
+        # Create RunnerDeployment
+        # This is the official Prefect pattern that handles all the complexity
+        self._info(f"Creating deployment for flow '{flow.name}'")
+
+        deployment = RunnerDeployment(
+            name=self.config["package"],
+            flow_name=flow.name,
+            entrypoint=entrypoint,
+            work_pool_name=self.config["work_pool"],
+            work_queue_name=self.config["work_queue"],
+            tags=[self.config["name"]],
+            version=self.config["version"],
+            description=flow.description
+            or f"Deployment for {self.config['package']} v{self.config['version']}",
+            storage=_PullStepStorage(pull_steps),
+            parameters={},
+            job_variables={},
+            paused=False,
+        )
+
+        # Verify work pool exists before deploying
+        async with get_client() as client:
+            try:
+                work_pool = await client.read_work_pool(self.config["work_pool"])
+                self._success(
+                    f"Work pool '{self.config['work_pool']}' verified (type: {work_pool.type})"
+                )
+            except Exception as e:
+                self._die(
+                    f"Work pool '{self.config['work_pool']}' not accessible: {e}\n"
+                    "Create it in the Prefect UI or with: prefect work-pool create"
+                )
+
+        # Apply deployment
+        # This automatically handles create vs update based on whether deployment exists
+        self._info("Applying deployment (create or update)...")
+        try:
+            deployment_id = await deployment.apply()  # type: ignore
+            self._success(f"Deployment ID: {deployment_id}")
+
+            # Print helpful URLs
+            if self.api_url:
+                ui_url = self.api_url.replace("/api/", "/")
+                print(f"\n🌐 View deployment: {ui_url}/deployments/deployment/{deployment_id}")
+                print(f"🚀 Run now: prefect deployment run '{flow.name}/{self.config['package']}'")
+        except Exception as e:
+            self._die(f"Failed to apply deployment: {e}")
+
+    async def run(self):
+        """Execute the complete deployment pipeline."""
+        print("=" * 70)
+        print(f"Prefect Deployment: {self.config['name']} v{self.config['version']}")
+        print(f"Target: gs://{self.config['bucket']}/{self.config['folder']}")
+        print("=" * 70)
+        print()
+
+        # Phase 1: Build
+        tarball = self._build_package()
+
+        # Phase 2: Upload
+        await self._upload_package(tarball)
+
+        # Phase 3: Deploy
+        await self._deploy_via_api()
+
+        print()
+        print("=" * 70)
+        self._success("Deployment complete!")
+        print("=" * 70)
+
+
+# ============================================================================
+# CLI Entry Point
+# ============================================================================
+
+
+def main():
+    """Command-line interface for deployment script."""
+    parser = argparse.ArgumentParser(
+        description="Deploy Prefect flows to GCP using the official RunnerDeployment pattern",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Example:
+  python -m ai_pipeline_core.utils.deploy
+
+Prerequisites:
+  - Settings configured with PREFECT_API_URL (and optionally PREFECT_API_KEY)
+  - Settings configured with PREFECT_GCS_BUCKET
+  - pyproject.toml with project name and version
+  - Package installed locally: pip install -e .
+  - GCP authentication configured (via service account or default credentials)
+  - Work pool created in Prefect UI or CLI
+
+Settings can be configured via:
+  - Environment variables (e.g., export PREFECT_API_URL=...)
+  - .env file in the current directory
+        """,
+    )
+
+    parser.parse_args()
+
+    try:
+        deployer = Deployer()
+        asyncio.run(deployer.run())
+    except KeyboardInterrupt:
+        print("\n✗ Deployment cancelled by user", file=sys.stderr)
+        sys.exit(1)
+    except Exception as e:
+        print(f"\n✗ Unexpected error: {e}", file=sys.stderr)
+
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

ai_pipeline_core/utils/remote_deployment.py

@@ -0,0 +1,269 @@
+"""Experimental remote deployment utilities.
+
+EXPERIMENTAL: This module provides utilities for calling remotely deployed Prefect flows.
+Subject to change in future versions.
+"""
+
+import inspect
+from functools import wraps
+from typing import Any, Callable, ParamSpec, Type, TypeVar
+
+from prefect import get_client
+from prefect.client.orchestration import PrefectClient
+from prefect.client.schemas import FlowRun
+from prefect.context import AsyncClientContext
+from prefect.deployments.flow_runs import run_deployment
+from prefect.exceptions import ObjectNotFound
+
+from ai_pipeline_core import DocumentList, FlowDocument
+from ai_pipeline_core.settings import settings
+from ai_pipeline_core.tracing import TraceLevel, set_trace_cost, trace
+
+# --------------------------------------------------------------------------- #
+# Utility functions (copied from pipeline.py for consistency)
+# --------------------------------------------------------------------------- #
+
+
+def _callable_name(obj: Any, fallback: str) -> str:
+    """Safely extract callable's name for error messages.
+
+    Args:
+        obj: Any object that might have a __name__ attribute.
+        fallback: Default name if extraction fails.
+
+    Returns:
+        The callable's __name__ if available, fallback otherwise.
+
+    Note:
+        Internal helper that never raises exceptions.
+    """
+    try:
+        n = getattr(obj, "__name__", None)
+        return n if isinstance(n, str) else fallback
+    except Exception:
+        return fallback
+
+
+def _is_already_traced(func: Callable[..., Any]) -> bool:
+    """Check if a function has already been wrapped by the trace decorator.
+
+    This checks both for the explicit __is_traced__ marker and walks
+    the __wrapped__ chain to detect nested trace decorations.
+
+    Args:
+        func: Function to check for existing trace decoration.
+
+    Returns:
+        True if the function is already traced, False otherwise.
+    """
+    # Check for explicit marker
+    if hasattr(func, "__is_traced__") and func.__is_traced__:  # type: ignore[attr-defined]
+        return True
+
+    # Walk the __wrapped__ chain to detect nested traces
+    current = func
+    depth = 0
+    max_depth = 10  # Prevent infinite loops
+
+    while hasattr(current, "__wrapped__") and depth < max_depth:
+        wrapped = current.__wrapped__  # type: ignore[attr-defined]
+        # Check if the wrapped function has the trace marker
+        if hasattr(wrapped, "__is_traced__") and wrapped.__is_traced__:  # type: ignore[attr-defined]
+            return True
+        current = wrapped
+        depth += 1
+
+    return False
+
+
+# --------------------------------------------------------------------------- #
+# Remote deployment execution
+# --------------------------------------------------------------------------- #
+
+
+async def run_remote_deployment(deployment_name: str, parameters: dict[str, Any]) -> Any:
+    """Run a remote Prefect deployment.
+
+    Args:
+        deployment_name: Name of the deployment to run.
+        parameters: Parameters to pass to the deployment.
+
+    Returns:
+        Result from the deployment execution.
+
+    Raises:
+        ValueError: If deployment is not found in local or remote Prefect API.
+    """
+
+    async def _run(client: PrefectClient, as_subflow: bool) -> Any:
+        fr: FlowRun = await run_deployment(
+            client=client, name=deployment_name, parameters=parameters, as_subflow=as_subflow
+        )  # type: ignore
+        return await fr.state.result()  # type: ignore
+
+    async with get_client() as client:
+        try:
+            await client.read_deployment_by_name(name=deployment_name)
+            return await _run(client, True)
+        except ObjectNotFound:
+            pass
+
+    if not settings.prefect_api_url:
+        raise ValueError(f"{deployment_name} deployment not found, PREFECT_API_URL is not set")
+
+    async with PrefectClient(
+        api=settings.prefect_api_url,
+        api_key=settings.prefect_api_key,
+        auth_string=settings.prefect_api_auth_string,
+    ) as client:
+        try:
+            await client.read_deployment_by_name(name=deployment_name)
+            with AsyncClientContext.model_construct(
+                client=client, _httpx_settings=None, _context_stack=0
+            ):
+                return await _run(client, False)
+        except ObjectNotFound:
+            pass
+
+    raise ValueError(f"{deployment_name} deployment not found")
+
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+def remote_deployment(
+    output_document_type: Type[FlowDocument],
+    *,
+    # tracing
+    name: str | None = None,
+    trace_level: TraceLevel = "always",
+    trace_ignore_input: bool = False,
+    trace_ignore_output: bool = False,
+    trace_ignore_inputs: list[str] | None = None,
+    trace_input_formatter: Callable[..., str] | None = None,
+    trace_output_formatter: Callable[..., str] | None = None,
+    trace_cost: float | None = None,
+    trace_trim_documents: bool = True,
+) -> Callable[[Callable[P, T]], Callable[P, T]]:
+    """Decorator for calling remote Prefect deployments with automatic tracing.
+
+    EXPERIMENTAL: Decorator for calling remote Prefect deployments with automatic
+    parameter serialization, result deserialization, and LMNR tracing.
+
+    IMPORTANT: Never combine with @trace decorator - this includes tracing automatically.
+    The framework will raise TypeError if you try to use both decorators together.
+
+    Best Practice - Use Defaults:
+        For most use cases, only specify output_document_type. The defaults provide
+        automatic tracing with optimal settings.
+
+    Args:
+        output_document_type: The FlowDocument type to deserialize results into.
+        name: Custom trace name (defaults to function name).
+        trace_level: When to trace ("always", "debug", "off").
+                    - "always": Always trace (default)
+                    - "debug": Only trace when LMNR_DEBUG="true"
+                    - "off": Disable tracing
+        trace_ignore_input: Don't trace input arguments.
+        trace_ignore_output: Don't trace return value.
+        trace_ignore_inputs: List of parameter names to exclude from tracing.
+        trace_input_formatter: Custom formatter for input tracing.
+        trace_output_formatter: Custom formatter for output tracing.
+        trace_cost: Optional cost value to track in metadata. When provided and > 0,
+             sets gen_ai.usage.output_cost, gen_ai.usage.cost, and cost metadata.
+        trace_trim_documents: Trim document content in traces to first 100 chars (default True).
+                             Reduces trace size with large documents.
+
+    Returns:
+        Decorator function that wraps the target function.
+
+    Example:
+        >>> # RECOMMENDED - Minimal usage
+        >>> @remote_deployment(output_document_type=OutputDoc)
+        >>> async def process_remotely(
+        ...     project_name: str,
+        ...     documents: DocumentList,
+        ...     flow_options: FlowOptions
+        >>> ) -> DocumentList:
+        ...     pass  # This stub is replaced by remote call
+        >>>
+        >>> # With custom tracing
+        >>> @remote_deployment(
+        ...     output_document_type=OutputDoc,
+        ...     trace_cost=0.05,  # Track cost of remote execution
+        ...     trace_level="debug"  # Only trace in debug mode
+        >>> )
+        >>> async def debug_remote_flow(...) -> DocumentList:
+        ...     pass
+
+    Note:
+        - Remote calls are automatically traced with LMNR
+        - The decorated function's body is never executed - it serves as a signature template
+        - Deployment name is auto-derived from function name
+        - DocumentList parameters are automatically serialized/deserialized
+
+    Raises:
+        TypeError: If function is already decorated with @trace.
+        ValueError: If deployment is not found.
+    """
+
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
+        fname = _callable_name(func, "remote_deployment")
+
+        # Check if function is already traced
+        if _is_already_traced(func):
+            raise TypeError(
+                f"@remote_deployment target '{fname}' is already decorated "
+                f"with @trace. Remove the @trace decorator - @remote_deployment includes "
+                f"tracing automatically."
+            )
+
+        @wraps(func)
+        async def _wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+            sig = inspect.signature(func)
+            bound = sig.bind(*args, **kwargs)
+            bound.apply_defaults()
+
+            # Serialize parameters, converting DocumentList to list[dict]
+            parameters = {}
+            for pname, value in bound.arguments.items():
+                if isinstance(value, DocumentList):
+                    parameters[pname] = [doc for doc in value]
+                else:
+                    parameters[pname] = value
+
+            # Auto-derive deployment name
+            deployment_name = f"{func.__name__.replace('_', '-')}/{func.__name__}"
+
+            result = await run_remote_deployment(
+                deployment_name=deployment_name, parameters=parameters
+            )
+
+            # Set trace cost if provided
+            if trace_cost is not None and trace_cost > 0:
+                set_trace_cost(trace_cost)
+
+            assert isinstance(result, list), "Result must be a list"
+
+            # Auto-handle return type conversion from list[dict] to DocumentList
+            return_type = sig.return_annotation
+
+            assert return_type is DocumentList, "Return type must be a DocumentList"
+            return DocumentList([output_document_type(**item) for item in result])  # type: ignore
+
+        # Apply trace decorator
+        traced_wrapper = trace(
+            level=trace_level,
+            name=name or fname,
+            ignore_input=trace_ignore_input,
+            ignore_output=trace_ignore_output,
+            ignore_inputs=trace_ignore_inputs,
+            input_formatter=trace_input_formatter,
+            output_formatter=trace_output_formatter,
+            trim_documents=trace_trim_documents,
+        )(_wrapper)
+
+        return traced_wrapper  # type: ignore
+
+    return decorator

{ai_pipeline_core-0.2.5.dist-info → ai_pipeline_core-0.2.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-pipeline-core
-Version: 0.2.5
+Version: 0.2.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

{ai_pipeline_core-0.2.5.dist-info → ai_pipeline_core-0.2.7.dist-info}/RECORD

@@ -1,11 +1,11 @@
-ai_pipeline_core/__init__.py,sha256=XjQnfQELFAbS9MOd41N59AQHkZuepKmHGeY_m8TzKSk,5720
+ai_pipeline_core/__init__.py,sha256=SCyd40nB8yF10KylHMBhEdpF0slV35OTsIePEKi9GP8,5720
 ai_pipeline_core/exceptions.py,sha256=vx-XLTw2fJSPs-vwtXVYtqoQUcOc0JeI7UmHqRqQYWU,1569
 ai_pipeline_core/pipeline.py,sha256=fWTVmrnOEIFge6o2NUYW2ndGef5UurpL8_fK5tkXbzI,28700
 ai_pipeline_core/prefect.py,sha256=91ZgLJHsDsRUW77CpNmkKxYs3RCJuucPM3pjKmNBeDg,2199
 ai_pipeline_core/prompt_manager.py,sha256=FAtb1yK7bGuAeuIJ523LOX9bd7TrcHG-TqZ7Lz4RJC0,12087
 ai_pipeline_core/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ai_pipeline_core/settings.py,sha256=-a9jVGg77xifj2SagCR9shXfzXUd-2MlrlquEu4htG8,5035
-ai_pipeline_core/tracing.py,sha256=mmK64s1lw18EE_7PQgfZb0sJhAuhkVDxXw_wBpR7UGE,31530
+ai_pipeline_core/settings.py,sha256=IMrFaX0i-WIlaOA5O53ipNSta6KQVSFHc1aJXmS3nSo,5078
+ai_pipeline_core/tracing.py,sha256=HT8heSwsVot6D6u8dPi-BHVlaemkPsPs5aXtG-iIzNk,31494
 ai_pipeline_core/documents/__init__.py,sha256=WHStvGZiSyybOcMTYxSV24U6MA3Am_0_Az5p-DuMFrk,738
 ai_pipeline_core/documents/document.py,sha256=hdTh36KGEcrDollTnQmTI66DJIqYfe4X42Y0q7Cm4fY,68153
 ai_pipeline_core/documents/document_list.py,sha256=Y_NCjfM_CjkIwHRD2iyGgYBuIykN8lT2IIH_uWOiGis,16254
@@ -18,21 +18,24 @@ ai_pipeline_core/flow/__init__.py,sha256=2BfWYMOPYW5teGzwo-qzpn_bom1lxxry0bPsjVg
 ai_pipeline_core/flow/config.py,sha256=3PCDph2n8dj-txqAvd9Wflbi_6lmfXFR9rUhM-szGSQ,18887
 ai_pipeline_core/flow/options.py,sha256=2rKR2GifhXcyw8avI_oiEDMLC2jm5Qzpw8z56pbxUMo,2285
 ai_pipeline_core/llm/__init__.py,sha256=3B_vtEzxrzidP1qOUNQ4RxlUmxZ2MBKQcUhQiTybM9g,661
-ai_pipeline_core/llm/ai_messages.py,sha256=tseyncD-T1IjaXKzUkzEgS3CGvz-WEDsA6v8wt9Emx0,14295
-ai_pipeline_core/llm/client.py,sha256=6gedRSeReT0SIbFZfqe73zjeSfgIXhBDRqeXUEHq4x8,22177
-ai_pipeline_core/llm/model_options.py,sha256=5XTBDdGZMxAo8JvhfEBIrtTOJC6piaMsgJ_J2cGiByo,11775
-ai_pipeline_core/llm/model_response.py,sha256=KhdlgR1vL5LmM_HHWKBAgI0LaqLdyaZqahEvtItpBEM,12706
-ai_pipeline_core/llm/model_types.py,sha256=OCRdTbQ1ZZ95nT_2PgOm16n2et25QOQyBtB1zsqm_3U,2791
+ai_pipeline_core/llm/ai_messages.py,sha256=Onin3UPdbJQNl3WfY3-_jE5KRmF-ciXsa5K6UPOiy5s,14410
+ai_pipeline_core/llm/client.py,sha256=385nKrr5fbDha8lNe9AbGLJ9Eszzo3_ArC22WWd8T_s,23646
+ai_pipeline_core/llm/model_options.py,sha256=uRNIHfVeh2sgt1mZBiOUx6hPQ6GKjB8b7TytZJ6afKg,11768
+ai_pipeline_core/llm/model_response.py,sha256=6kEr9ss3UGlykvtWAvh1l55rGw2-wyVup3QJhm0Oggc,13264
+ai_pipeline_core/llm/model_types.py,sha256=2J4Qsb1x21I4eo_VPeaMMOW8shOGPqzJuoGjTLcBFPM,2791
 ai_pipeline_core/logging/__init__.py,sha256=Nz6-ghAoENsgNmLD2ma9TW9M0U2_QfxuQ5DDW6Vt6M0,651
 ai_pipeline_core/logging/logging.yml,sha256=YTW48keO_K5bkkb-KXGM7ZuaYKiquLsjsURei8Ql0V4,1353
 ai_pipeline_core/logging/logging_config.py,sha256=pV2x6GgMPXrzPH27sicCSXfw56beio4C2JKCJ3NsXrg,6207
 ai_pipeline_core/logging/logging_mixin.py,sha256=OTye2pbUbG5oYZkI06TNkGCEa4y0ldePz5IAfdmNUPU,8090
 ai_pipeline_core/simple_runner/__init__.py,sha256=9krT-CcDAZ0jB2MjWqFYhaK5qtUDMpB5qWzjRLa4Zhk,322
-ai_pipeline_core/simple_runner/cli.py,sha256=yVyuxLY2RZvdNwmwT5LCe-km2nQJzWTPI0vSWn4_yms,9344
+ai_pipeline_core/simple_runner/cli.py,sha256=p9Z1jtRMH10T5Bl3QfHPxyW6LL4qYvvXeOXbPGeeXeE,9308
 ai_pipeline_core/simple_runner/simple_runner.py,sha256=f6cIodYkul-Apu1d63T6kR5DZpiaCWpphUcEPp5XjFo,9102
 ai_pipeline_core/storage/__init__.py,sha256=tcIkjJ3zPBLCyetwiJDewBvS2sbRJrDlBh3gEsQm08E,184
 ai_pipeline_core/storage/storage.py,sha256=ClMr419Y-eU2RuOjZYd51dC0stWQk28Vb56PvQaoUwc,20007
-ai_pipeline_core-0.2.5.dist-info/METADATA,sha256=y6-nOwAztmZCHnCY4Z5tmrDmR7pQGMyxCj8OONVMz7E,15159
-ai_pipeline_core-0.2.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-ai_pipeline_core-0.2.5.dist-info/licenses/LICENSE,sha256=kKj8mfbdWwkyG3U6n7ztB3bAZlEwShTkAsvaY657i3I,1074
-ai_pipeline_core-0.2.5.dist-info/RECORD,,
+ai_pipeline_core/utils/__init__.py,sha256=TJSmEm1Quf-gKwXrxM96u2IGzVolUyeNNfLMPoLstXI,254
+ai_pipeline_core/utils/deploy.py,sha256=rAtRuwkmGkc-fqvDMXpt08OzLrD7KTDMAmLDC9wYg7Y,13147
+ai_pipeline_core/utils/remote_deployment.py,sha256=cPTgnS5InK08qiWnuPz3e8YKjoT3sPBloSaDfNTzghs,10137
+ai_pipeline_core-0.2.7.dist-info/METADATA,sha256=Ec4yFJTtt8qFtQD4Hg7KXT2VO80_i8JmLbwfXNia2GE,15159
+ai_pipeline_core-0.2.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ai_pipeline_core-0.2.7.dist-info/licenses/LICENSE,sha256=kKj8mfbdWwkyG3U6n7ztB3bAZlEwShTkAsvaY657i3I,1074
+ai_pipeline_core-0.2.7.dist-info/RECORD,,