ai-pipeline-core 0.2.6__py3-none-any.whl → 0.2.8__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.6"
+__version__ = "0.2.8"
 
 __all__ = [
     # Config/Settings

ai_pipeline_core/flow/options.py

@@ -61,11 +61,11 @@ class FlowOptions(BaseSettings):
     """
 
     core_model: ModelName = Field(
-        default="gpt-5",
+        default="gemini-2.5-pro",
         description="Primary model for complex analysis and generation tasks.",
     )
     small_model: ModelName = Field(
-        default="gpt-5-mini",
+        default="grok-4-fast",
         description="Fast, cost-effective model for simple tasks and orchestration.",
     )
 

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

@@ -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,38 @@ def _process_messages(
     return processed_messages
 
 
+def _remove_cache_control(
+    messages: list[ChatCompletionMessageParam],
+) -> list[ChatCompletionMessageParam]:
+    """Remove cache control directives from messages.
+
+    Internal utility that strips cache_control fields from both message-level
+    and content-level entries. Used in retry logic when cache-related errors
+    occur during LLM API calls.
+
+    Args:
+        messages: List of messages that may contain cache_control directives.
+
+    Returns:
+        The same message list (modified in-place) with all cache_control
+        fields removed from both messages and their content items.
+
+    Note:
+        This function modifies the input list in-place but also returns it
+        for convenience. Handles both list-based content (multipart) and
+        string content (simple messages).
+    """
+    for message in messages:
+        if content := message.get("content"):
+            if isinstance(content, list):
+                for item in content:
+                    if "cache_control" in item:
+                        del item["cache_control"]
+        if "cache_control" in message:
+            del message["cache_control"]
+    return messages  # type: ignore
+
+
 def _model_name_to_openrouter_model(model: ModelName) -> str:
     """Convert a model name to an OpenRouter model name.
 
@@ -237,6 +274,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 < 10000:
+        # 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
     )
@@ -263,6 +304,8 @@ async def _generate_with_retry(
             if not isinstance(e, asyncio.TimeoutError):
                 # disable cache if it's not a timeout because it may cause an error
                 completion_kwargs["extra_body"]["cache"] = {"no-cache": True}
+                # sometimes there are issues with cache so cache is removed in case of failure
+                processed_messages = _remove_cache_control(processed_messages)
 
             logger.warning(
                 f"LLM generation failed (attempt {attempt + 1}/{options.retries}): {e}",
@@ -374,26 +417,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

@@ -266,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/utils/deploy.py

@@ -7,8 +7,8 @@ This script:
 3. Creates/updates a Prefect deployment using the RunnerDeployment pattern
 
 Requirements:
-- .env file with PREFECT_API_URL and optionally PREFECT_API_KEY
-- .env file with PREFECT_GCS_BUCKET
+- 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
 
@@ -18,7 +18,6 @@ Usage:
 
 import argparse
 import asyncio
-import os
 import subprocess
 import sys
 import tomllib
@@ -34,14 +33,6 @@ from prefect.flows import load_flow_from_entrypoint
 from ai_pipeline_core.settings import settings
 from ai_pipeline_core.storage import Storage
 
-# ============================================================================
-# Configuration
-# ============================================================================
-
-WORK_POOL_NAME = settings.prefect_work_pool_name
-DEFAULT_WORK_QUEUE = settings.prefect_work_queue_name
-PREDEFINED_BUCKET = settings.prefect_gcs_bucket
-
 # ============================================================================
 # Deployer Class
 # ============================================================================
@@ -57,7 +48,7 @@ class Deployer:
     def __init__(self):
         """Initialize deployer."""
         self.config = self._load_config()
-        self._setup_prefect_env()
+        self._validate_prefect_settings()
 
     def _load_config(self) -> dict[str, Any]:
         """Load and normalize project configuration from pyproject.toml.
@@ -65,10 +56,10 @@ class Deployer:
         Returns:
             Configuration dictionary with project metadata and deployment settings.
         """
-        if not PREDEFINED_BUCKET:
+        if not settings.prefect_gcs_bucket:
             self._die(
-                "PREFECT_GCS_BUCKET not found in .env file.\n"
-                "Create a .env file with:\n"
+                "PREFECT_GCS_BUCKET not configured in settings.\n"
+                "Configure via environment variable or .env file:\n"
                 "  PREFECT_GCS_BUCKET=your-bucket-name"
             )
 
@@ -97,33 +88,23 @@ class Deployer:
             "name": name,
             "package": package_name,
             "version": version,
-            "bucket": PREDEFINED_BUCKET,
+            "bucket": settings.prefect_gcs_bucket,
             "folder": f"flows/{flow_folder}",
             "tarball": f"{package_name}-{version}.tar.gz",
-            "work_pool": WORK_POOL_NAME,
-            "work_queue": DEFAULT_WORK_QUEUE,
+            "work_pool": settings.prefect_work_pool_name,
+            "work_queue": settings.prefect_work_queue_name,
         }
 
-    def _setup_prefect_env(self):
-        """Configure Prefect environment variables from .env file."""
-        self.api_url = os.getenv("PREFECT_API_URL")
+    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 found in .env file.\n"
-                "Create a .env file with:\n"
+                "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/..."
             )
 
-        os.environ["PREFECT_API_URL"] = self.api_url
-
-        # Optional: API key for authentication
-        if api_key := os.getenv("PREFECT_API_KEY"):
-            os.environ["PREFECT_API_KEY"] = api_key
-
-        # Optional: Alternative auth method
-        if api_auth := os.getenv("PREFECT_API_AUTH_STRING"):
-            os.environ["PREFECT_API_AUTH_STRING"] = api_auth
-
     def _run(self, cmd: str, check: bool = True) -> Optional[str]:
         """Execute shell command and return output.
 
@@ -345,12 +326,16 @@ Example:
   python -m ai_pipeline_core.utils.deploy
 
 Prerequisites:
-  - .env file with PREFECT_API_URL (and optionally PREFECT_API_KEY)
-  - .env file with PREFECT_GCS_BUCKET
+  - 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
         """,
     )
 

{ai_pipeline_core-0.2.6.dist-info → ai_pipeline_core-0.2.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-pipeline-core
-Version: 0.2.6
+Version: 0.2.8
 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.6.dist-info → ai_pipeline_core-0.2.8.dist-info}/RECORD

@@ -1,4 +1,4 @@
-ai_pipeline_core/__init__.py,sha256=BWdVF64no62Cg9b_6GLiRub7ytBQQPk1RC38V0FGoAA,5720
+ai_pipeline_core/__init__.py,sha256=NMpeyF-rJdNnNWfi9eoOgWNnk7kOVHqvSZq3cnyDuCI,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
@@ -16,13 +16,13 @@ ai_pipeline_core/documents/temporary_document.py,sha256=Sam344Mm5AlZTm3_l01YdDWe
 ai_pipeline_core/documents/utils.py,sha256=ZyJNjFN7ihWno0K7dJZed7twYmmPLA0z40UzFw1A3A8,5465
 ai_pipeline_core/flow/__init__.py,sha256=2BfWYMOPYW5teGzwo-qzpn_bom1lxxry0bPsjVgcsCk,188
 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/flow/options.py,sha256=zn3N5DgYtlxLq0AvXfana3UOhym7A3XCheQSBIIarZE,2295
 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=Ojb0Ew2w88KlTfNP7VyRsIjjHh5Z8A1WhfHsNQwfeak,23673
-ai_pipeline_core/llm/model_options.py,sha256=5XTBDdGZMxAo8JvhfEBIrtTOJC6piaMsgJ_J2cGiByo,11775
-ai_pipeline_core/llm/model_response.py,sha256=VsuEp8tMbxWOkpZBFH1WIJisbZSvg7Z4JPqPejz7YM8,13273
-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=Da1NgGzfIkFRw_aDASK36MMdKe3DXaj5_3wMg0gR-Hk,24999
+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
@@ -33,9 +33,9 @@ ai_pipeline_core/simple_runner/simple_runner.py,sha256=f6cIodYkul-Apu1d63T6kR5DZ
 ai_pipeline_core/storage/__init__.py,sha256=tcIkjJ3zPBLCyetwiJDewBvS2sbRJrDlBh3gEsQm08E,184
 ai_pipeline_core/storage/storage.py,sha256=ClMr419Y-eU2RuOjZYd51dC0stWQk28Vb56PvQaoUwc,20007
 ai_pipeline_core/utils/__init__.py,sha256=TJSmEm1Quf-gKwXrxM96u2IGzVolUyeNNfLMPoLstXI,254
-ai_pipeline_core/utils/deploy.py,sha256=Y15-xoIWfs-HkEQeTobjq4UdTiXSTKzTghL0Jo8iMgU,13565
+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.6.dist-info/METADATA,sha256=rC6CVxTMsoL6iCIIw3b_Lr5zWdBBmzqCgboiS4vFXpY,15159
-ai_pipeline_core-0.2.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-ai_pipeline_core-0.2.6.dist-info/licenses/LICENSE,sha256=kKj8mfbdWwkyG3U6n7ztB3bAZlEwShTkAsvaY657i3I,1074
-ai_pipeline_core-0.2.6.dist-info/RECORD,,
+ai_pipeline_core-0.2.8.dist-info/METADATA,sha256=Sl_8s24ar0lwa7iio4d0QYDhzsmAIGHHqOKuXnFxP7s,15159
+ai_pipeline_core-0.2.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ai_pipeline_core-0.2.8.dist-info/licenses/LICENSE,sha256=kKj8mfbdWwkyG3U6n7ztB3bAZlEwShTkAsvaY657i3I,1074
+ai_pipeline_core-0.2.8.dist-info/RECORD,,