ai-pipeline-core 0.2.7__tar.gz → 0.2.9__tar.gz

{ai_pipeline_core-0.2.7 → ai_pipeline_core-0.2.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-pipeline-core
-Version: 0.2.7
+Version: 0.2.9
 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.7 → ai_pipeline_core-0.2.9}/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.7"
+__version__ = "0.2.9"
 
 __all__ = [
     # Config/Settings

{ai_pipeline_core-0.2.7 → ai_pipeline_core-0.2.9}/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-0.2.7 → ai_pipeline_core-0.2.9}/ai_pipeline_core/llm/client.py

@@ -45,31 +45,30 @@ def _process_messages(
 
     Internal function that combines context and messages into a single
     list of API-compatible messages. Applies caching directives to
-    context messages for efficiency.
+    system prompt and context messages for efficiency.
 
     Args:
         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", "300s", "1h").
+        cache_ttl: Cache TTL for system and context messages (e.g. "120s", "300s", "1h").
                    Set to None or empty string to disable caching.
 
     Returns:
         List of formatted messages ready for API calls, with:
-        - System prompt at the beginning (if provided)
-        - Context messages with cache_control on the last one (if cache_ttl)
+        - System prompt at the beginning with cache_control (if provided and cache_ttl set)
+        - Context messages with cache_control on all messages (if cache_ttl set)
         - Regular messages without caching
 
     System Prompt Location:
         The system prompt parameter is always injected as the FIRST message
-        with role="system". It is NOT cached with context, allowing dynamic
-        system prompts without breaking cache efficiency.
+        with role="system". It is cached along with context when cache_ttl is set.
 
     Cache behavior:
-        The last context message gets ephemeral caching with specified TTL
+        All system and context messages get ephemeral caching with specified TTL
         to reduce token usage on repeated calls with same context.
         If cache_ttl is None or empty string (falsy), no caching is applied.
-        Only the last context message receives cache_control to maximize efficiency.
+        All system and context messages receive cache_control to maximize cache efficiency.
 
     Note:
         This is an internal function used by _generate_with_retry().
@@ -79,26 +78,28 @@ def _process_messages(
 
     # Add system prompt if provided
     if system_prompt:
-        processed_messages.append({"role": "system", "content": system_prompt})
+        processed_messages.append({
+            "role": "system",
+            "content": [{"type": "text", "text": system_prompt}],
+        })
 
     # Process context messages with caching if provided
     if context:
         # Use AIMessages.to_prompt() for context
         context_messages = context.to_prompt()
+        processed_messages.extend(context_messages)
 
-        # 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
+    if cache_ttl:
+        for message in processed_messages:
+            message["cache_control"] = {  # type: ignore
                 "type": "ephemeral",
                 "ttl": cache_ttl,
             }
-
-        processed_messages.extend(context_messages)
+            if isinstance(message["content"], list):  # type: ignore
+                message["content"][-1]["cache_control"] = {  # type: ignore
+                    "type": "ephemeral",
+                    "ttl": cache_ttl,
+                }
 
     # Process regular messages without caching
     if messages:
@@ -108,6 +109,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.
 
@@ -124,7 +157,7 @@ def _model_name_to_openrouter_model(model: ModelName) -> str:
     if model == "grok-4-fast-search":
         return "x-ai/grok-4-fast:online"
     if model == "sonar-pro-search":
-        return "perplexity/sonar-reasoning-pro"
+        return "perplexity/sonar-pro-search"
     if model.startswith("gemini"):
         return f"google/{model}"
     elif model.startswith("gpt"):
@@ -242,7 +275,7 @@ 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:
+    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
 
@@ -272,6 +305,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}",

{ai_pipeline_core-0.2.7 → ai_pipeline_core-0.2.9}/ai_pipeline_core/llm/model_response.py

@@ -88,10 +88,13 @@ class ModelResponse(ChatCompletion):
         data = chat_completion.model_dump()
 
         # fixes issue where the role is "assistantassistant" instead of "assistant"
+        valid_finish_reasons = {"stop", "length", "tool_calls", "content_filter", "function_call"}
         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"
+            data["choices"][i]["message"]["role"] = "assistant"
+            # Only update finish_reason if it's not already a valid value
+            current_finish_reason = data["choices"][i].get("finish_reason")
+            if current_finish_reason not in valid_finish_reasons:
+                data["choices"][i]["finish_reason"] = "stop"
 
         super().__init__(**data)
 

{ai_pipeline_core-0.2.7 → ai_pipeline_core-0.2.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ai-pipeline-core"
-version = "0.2.7"
+version = "0.2.9"
 description = "Core utilities for AI-powered processing pipelines using prefect"
 readme = "README.md"
 license = {text = "MIT"}
@@ -178,7 +178,7 @@ reportIncompatibleVariableOverride = "error"
 reportMissingParameterType = "warning"
 
 [tool.bumpversion]
-current_version = "0.2.7"
+current_version = "0.2.9"
 commit = true
 tag = true
 tag_name = "v{new_version}"