ai-pipeline-core 0.2.6__py3-none-any.whl → 0.4.1__py3-none-any.whl

ai_pipeline_core/llm/model_options.py

@@ -41,11 +41,11 @@ class ModelOptions(BaseModel):
 
         retries: Number of retry attempts on failure (default: 3).
 
-        retry_delay_seconds: Seconds to wait between retries (default: 10).
+        retry_delay_seconds: Seconds to wait between retries (default: 20).
 
-        timeout: Maximum seconds to wait for response (default: 300).
+        timeout: Maximum seconds to wait for response (default: 600).
 
-        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.
 
@@ -99,77 +99,11 @@ class ModelOptions(BaseModel):
                    Merged with usage_tracking if both are set.
                    Useful for beta features or provider-specific capabilities.
 
-    Example:
-        >>> # Basic configuration
-        >>> options = ModelOptions(
-        ...     temperature=0.7,
-        ...     max_completion_tokens=1000
-        ... )
-        >>>
-        >>> # With system prompt
-        >>> options = ModelOptions(
-        ...     system_prompt="You are a helpful coding assistant",
-        ...     temperature=0.3  # Lower for code generation
-        ... )
-        >>>
-        >>> # With custom cache TTL
-        >>> options = ModelOptions(
-        ...     cache_ttl="300s",  # Cache context for 5 minutes
-        ...     max_completion_tokens=1000
-        ... )
-        >>>
-        >>> # Disable caching
-        >>> options = ModelOptions(
-        ...     cache_ttl=None,  # No context caching
-        ...     temperature=0.5
-        ... )
-        >>>
-        >>> # For search-enabled models
-        >>> options = ModelOptions(
-        ...     search_context_size="high",  # Get more search results
-        ...     max_completion_tokens=2000
-        ... )
-        >>>
-        >>> # For reasoning models
-        >>> options = ModelOptions(
-        ...     reasoning_effort="high",  # Deep reasoning
-        ...     timeout=600  # More time for complex reasoning
-        ... )
-        >>>
-        >>> # With stop sequences
-        >>> options = ModelOptions(
-        ...     stop=["STOP", "END", "\n\n"],  # Stop on these sequences
-        ...     temperature=0.7
-        ... )
-        >>>
-        >>> # With custom extra_body parameters
-        >>> options = ModelOptions(
-        ...     extra_body={"custom_param": "value", "beta_feature": True},
-        ...     usage_tracking=True  # Still tracks usage alongside custom params
-        ... )
-        >>>
-        >>> # With user tracking for cost monitoring
-        >>> options = ModelOptions(
-        ...     user="user_12345",  # Track costs per user
-        ...     temperature=0.7
-        ... )
-        >>>
-        >>> # With metadata for tracking and observability
-        >>> options = ModelOptions(
-        ...     metadata={"experiment": "v1", "version": "2.0", "feature": "search"},
-        ...     temperature=0.7
-        ... )
-
-    Note:
-        - Not all options apply to all models
-        - 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
-        - 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
-        - usage_tracking is enabled by default for cost monitoring
+    Not all options apply to all models. search_context_size only works with search models,
+    reasoning_effort only works with models that support explicit reasoning, and
+    response_format is set internally by generate_structured(). cache_ttl accepts formats
+    like "120s", "5m", "1h" or None (default: "300s"). Stop sequences are limited to 4 by
+    most providers.
     """
 
     temperature: float | None = None
@@ -179,18 +113,19 @@ 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
     response_format: type[BaseModel] | None = None
     verbosity: Literal["low", "medium", "high"] | None = None
+    stream: bool = False
     usage_tracking: bool = True
     user: str | None = None
     metadata: dict[str, str] | None = None
     extra_body: dict[str, Any] | None = None
 
-    def to_openai_completion_kwargs(self) -> dict[str, Any]:
+    def to_openai_completion_kwargs(self) -> dict[str, Any]:  # noqa: C901
         """Convert options to OpenAI API completion parameters.
 
         Transforms ModelOptions fields into the format expected by
@@ -221,16 +156,9 @@ class ModelOptions(BaseModel):
             {"web_search_options": {"search_context_size": "low|medium|high"}}
             Non-search models silently ignore this parameter.
 
-        Example:
-            >>> options = ModelOptions(temperature=0.5, timeout=60)
-            >>> kwargs = options.to_openai_completion_kwargs()
-            >>> kwargs
-            {'timeout': 60, 'extra_body': {}, 'temperature': 0.5}
-
-        Note:
-            - system_prompt is handled separately in _process_messages()
-            - retries and retry_delay_seconds are used by retry logic
-            - extra_body always includes usage tracking for cost monitoring
+        system_prompt is handled separately in _process_messages().
+        retries and retry_delay_seconds are used by retry logic.
+        extra_body always includes usage tracking for cost monitoring.
         """
         kwargs: dict[str, Any] = {
             "timeout": self.timeout,

ai_pipeline_core/llm/model_response.py

@@ -1,13 +1,12 @@
 """Model response structures for LLM interactions.
 
-@public
-
 Provides enhanced response classes that use OpenAI-compatible base types via LiteLLM
 with additional metadata, cost tracking, and structured output support.
 """
 
 import json
 from copy import deepcopy
+from dataclasses import dataclass
 from typing import Any, Generic, TypeVar
 
 from openai.types.chat import ChatCompletion
@@ -21,14 +20,20 @@ T = TypeVar(
 """Type parameter for structured response Pydantic models."""
 
 
+@dataclass(frozen=True)
+class Citation:
+    """A URL citation returned by search-enabled models (e.g. sonar-pro-search, gemini-3-flash-search)."""
+
+    title: str
+    url: str
+
+
 class ModelResponse(ChatCompletion):
     """Response wrapper for LLM text generation.
 
-    @public
-
     Primary usage is adding to AIMessages for multi-turn conversations:
 
-        >>> response = await llm.generate("gpt-5", messages=messages)
+        >>> response = await llm.generate("gpt-5.1", messages=messages)
         >>> messages.append(response)  # Add assistant response to conversation
         >>> print(response.content)  # Access generated text
 
@@ -39,22 +44,9 @@ class ModelResponse(ChatCompletion):
     Almost all use cases are covered by these two patterns. Advanced features
     like token usage and cost tracking are available but rarely needed.
 
-    Example:
-        >>> from ai_pipeline_core import llm, AIMessages
-        >>>
-        >>> messages = AIMessages(["Explain quantum computing"])
-        >>> response = await llm.generate("gpt-5", messages=messages)
-        >>>
-        >>> # Primary usage: add to conversation
-        >>> messages.append(response)
-        >>>
-        >>> # Access generated text
-        >>> print(response.content)
-
-    Note:
-        Inherits from OpenAI's ChatCompletion for compatibility.
-        Other properties (usage, model, id) should only be accessed
-        when absolutely necessary.
+    Inherits from OpenAI's ChatCompletion for compatibility.
+    Other properties (usage, model, id) should only be accessed
+    when absolutely necessary.
     """
 
     def __init__(
@@ -77,21 +69,21 @@ class ModelResponse(ChatCompletion):
                      Includes timing information and custom tags.
             usage: Optional usage information from streaming response.
 
-        Example:
-            >>> # Usually created internally by generate()
-            >>> response = ModelResponse(
-            ...     chat_completion=completion,
-            ...     model_options={"temperature": 0.7, "model": "gpt-4"},
-            ...     metadata={"time_taken": 1.5, "first_token_time": 0.3}
-            ... )
         """
         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"
+            # Strip annotations with unsupported types (e.g. Grok returns type="file" for PDFs,
+            # but OpenAI's ChatCompletion only accepts type="url_citation")
+            if annotations := data["choices"][i]["message"].get("annotations"):
+                data["choices"][i]["message"]["annotations"] = [a for a in annotations if a.get("type") == "url_citation"]
 
         super().__init__(**data)
 
@@ -104,22 +96,12 @@ class ModelResponse(ChatCompletion):
     def content(self) -> str:
         """Get the generated text content.
 
-        @public
-
         Primary property for accessing the LLM's response text.
         This is the main property you'll use with ModelResponse.
 
         Returns:
             Generated text from the model, or empty string if none.
 
-        Example:
-            >>> response = await generate("gpt-5", messages="Hello")
-            >>> text = response.content  # The generated response
-            >>>
-            >>> # Common pattern: add to messages then use content
-            >>> messages.append(response)
-            >>> if "error" in response.content.lower():
-            ...     # Handle error case
         """
         content = self.choices[0].message.content or ""
         return content.split("</think>")[-1].strip()
@@ -128,8 +110,6 @@ class ModelResponse(ChatCompletion):
     def reasoning_content(self) -> str:
         """Get the reasoning content.
 
-        @public
-
         Returns:
             The reasoning content from the model, or empty string if none.
         """
@@ -140,7 +120,19 @@ class ModelResponse(ChatCompletion):
             return ""
         return message.content.split("</think>")[0].strip()
 
-    def get_laminar_metadata(self) -> dict[str, str | int | float]:
+    @property
+    def citations(self) -> list[Citation]:
+        """Get URL citations from search-enabled models.
+
+        Returns:
+            List of Citation objects with title and url. Empty list for non-search models.
+        """
+        annotations = self.choices[0].message.annotations
+        if not annotations:
+            return []
+        return [Citation(title=a.url_citation.title, url=a.url_citation.url) for a in annotations if a.url_citation]
+
+    def get_laminar_metadata(self) -> dict[str, str | int | float]:  # noqa: C901
         """Extract metadata for LMNR (Laminar) observability including cost tracking.
 
         Collects comprehensive metadata about the generation for tracing,
@@ -175,56 +167,26 @@ class ModelResponse(ChatCompletion):
             1. x-litellm-response-cost header (primary)
             2. usage.cost attribute (fallback)
 
-            Cost is stored in three fields for compatibility:
-            - gen_ai.usage.output_cost (standard)
-            - gen_ai.usage.cost (alternative)
-            - gen_ai.cost (simple)
-
-        Example:
-            >>> response = await llm.generate(
-            ...     "gpt-5",
-            ...     context=large_doc,
-            ...     messages="Summarize this"
-            ... )
-            >>>
-            >>> # Get comprehensive metadata
-            >>> metadata = response.get_laminar_metadata()
-            >>>
-            >>> # Track generation cost
-            >>> cost = metadata.get('gen_ai.usage.output_cost', 0)
-            >>> if cost > 0:
-            ...     print(f"Generation cost: ${cost:.4f}")
-            >>>
-            >>> # Monitor token usage
-            >>> print(f"Input: {metadata.get('gen_ai.usage.prompt_tokens', 0)} tokens")
-            >>> print(f"Output: {metadata.get('gen_ai.usage.completion_tokens', 0)} tokens")
-            >>> print(f"Total: {metadata.get('gen_ai.usage.total_tokens', 0)} tokens")
-            >>>
-            >>> # Check cache effectiveness
-            >>> cached = metadata.get('gen_ai.usage.cached_tokens', 0)
-            >>> if cached > 0:
-            ...     total = metadata.get('gen_ai.usage.total_tokens', 1)
-            ...     savings = (cached / total) * 100
-            ...     print(f"Cache hit: {cached} tokens ({savings:.1f}% savings)")
-            >>>
-            >>> # Calculate cost per token
-            >>> if cost > 0 and metadata.get('gen_ai.usage.total_tokens'):
-            ...     cost_per_1k = (cost / metadata['gen_ai.usage.total_tokens']) * 1000
-            ...     print(f"Cost per 1K tokens: ${cost_per_1k:.4f}")
-
-        Note:
-            - Cost availability depends on LiteLLM proxy configuration
-            - Not all providers return cost information
-            - Cached tokens reduce actual cost but may not be reflected
-            - Used internally by tracing but accessible for cost analysis
+            Cost is stored in three fields for observability tool consumption:
+            - gen_ai.usage.output_cost (OpenTelemetry GenAI semantic convention)
+            - gen_ai.usage.cost (aggregated cost)
+            - gen_ai.cost (short-form)
+
+        Cost availability depends on LiteLLM proxy configuration. Not all providers
+        return cost information. Cached tokens reduce actual cost but may not be reflected.
+        Used internally by tracing but accessible for cost analysis.
         """
         metadata: dict[str, str | int | float] = deepcopy(self._metadata)
 
         # Add base metadata
+        # NOTE: gen_ai.response.model is intentionally omitted — Laminar's UI uses it
+        # to override the span display name in the tree view, hiding the actual span name
+        # (set via `purpose` parameter). Tracked upstream: Laminar's getSpanDisplayName()
+        # in frontend/components/traces/trace-view/utils.ts prefers model over span name
+        # for LLM spans. Restore once Laminar shows both or prefers span name.
         metadata.update({
             "gen_ai.response.id": self.id,
-            "gen_ai.response.model": self.model,
-            "get_ai.system": "litellm",
+            "gen_ai.system": "litellm",
         })
 
         # Add usage metadata if available
@@ -242,21 +204,19 @@ class ModelResponse(ChatCompletion):
                 cost = float(self.usage.cost)  # type: ignore[attr-defined]
 
             # Add reasoning tokens if available
-            if completion_details := self.usage.completion_tokens_details:
-                if reasoning_tokens := completion_details.reasoning_tokens:
-                    metadata["gen_ai.usage.reasoning_tokens"] = reasoning_tokens
+            if (completion_details := self.usage.completion_tokens_details) and (reasoning_tokens := completion_details.reasoning_tokens):
+                metadata["gen_ai.usage.reasoning_tokens"] = reasoning_tokens
 
             # Add cached tokens if available
-            if prompt_details := self.usage.prompt_tokens_details:
-                if cached_tokens := prompt_details.cached_tokens:
-                    metadata["gen_ai.usage.cached_tokens"] = cached_tokens
+            if (prompt_details := self.usage.prompt_tokens_details) and (cached_tokens := prompt_details.cached_tokens):
+                metadata["gen_ai.usage.cached_tokens"] = cached_tokens
 
         # Add cost metadata if available
         if cost and cost > 0:
             metadata.update({
                 "gen_ai.usage.output_cost": cost,
                 "gen_ai.usage.cost": cost,
-                "get_ai.cost": cost,
+                "gen_ai.cost": cost,
             })
 
         for key, value in self._model_options.items():
@@ -266,7 +226,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)
@@ -275,7 +235,7 @@ class ModelResponse(ChatCompletion):
 
         message = self.choices[0].message
         for key, value in message.__dict__.items():
-            if key in ["content"]:
+            if key in {"content"}:
                 continue
             metadata[f"response.raw.message.{key}"] = json.dumps(value, indent=2, default=str)
 
@@ -294,16 +254,13 @@ class ModelResponse(ChatCompletion):
         if not self.content:
             raise ValueError("Empty response content")
 
-        if response_format := self._model_options.get("response_format"):
-            if isinstance(response_format, BaseModel):
-                response_format.model_validate_json(self.content)
+        if (response_format := self._model_options.get("response_format")) and isinstance(response_format, BaseModel):
+            response_format.model_validate_json(self.content)
 
 
-class StructuredModelResponse(ModelResponse, Generic[T]):
+class StructuredModelResponse(ModelResponse, Generic[T]):  # noqa: UP046
     """Response wrapper for structured/typed LLM output.
 
-    @public
-
     Primary usage is accessing the .parsed property for the structured data.
     """
 

ai_pipeline_core/llm/model_types.py

@@ -10,44 +10,41 @@ Model categories:
 - Search models: Models with web search capabilities
 """
 
-from typing import Literal, TypeAlias
+from typing import Literal
 
-ModelName: TypeAlias = (
+type ModelName = (
     Literal[
         # Core models
-        "gemini-2.5-pro",
-        "gpt-5",
-        "grok-4",
+        "gemini-3-pro",
+        "gpt-5.1",
         # Small models
-        "gemini-2.5-flash",
-        "gpt-5-nano",
-        "grok-4-fast",
+        "gemini-3-flash",
+        "gpt-5-mini",
+        "grok-4.1-fast",
         # Search models
-        "gemini-2.5-flash-search",
+        "gemini-3-flash-search",
+        "gpt-5-mini-search",
+        "grok-4.1-fast-search",
         "sonar-pro-search",
-        "gpt-4o-search",
-        "grok-4-fast-search",
     ]
     | str
 )
 """Type-safe model name identifiers with support for custom models.
 
-@public
-
 Provides IDE autocompletion for common model names while allowing any
 string for custom models. The type is a union of predefined literals
 and str, giving you the best of both worlds: suggestions for known
 models and flexibility for custom ones.
 
-Note: These are example common model names as of Q3 2025. Actual availability
+These are example common model names as of Q1 2026. Actual availability
 depends on your LiteLLM proxy configuration and provider access.
 
 Model categories:
-    Core models (gemini-2.5-pro, gpt-5, grok-4):
+    Core models (gemini-3-pro, gpt-5.1):
         High-capability models for complex tasks requiring deep reasoning,
         nuanced understanding, or creative generation.
 
-    Small models (gemini-2.5-flash, gpt-5-mini, grok-4-fast):
+    Small models (gemini-3-flash, gpt-5-mini, grok-4.1-fast):
         Efficient models optimized for speed and cost, suitable for
         simpler tasks or high-volume processing.
 
@@ -61,22 +58,7 @@ Using custom models:
     - Custom models work seamlessly as strings
     - No need for Union types or additional type aliases
 
-Example:
-    >>> from ai_pipeline_core import llm, ModelName
-    >>>
-    >>> # Predefined model with IDE autocomplete
-    >>> model: ModelName = "gpt-5"  # IDE suggests common models
-    >>> response = await llm.generate(model, messages="Hello")
-    >>>
-    >>> # Custom model works directly
-    >>> model: ModelName = "custom-model-v2"  # Any string is valid
-    >>> response = await llm.generate(model, messages="Hello")
-    >>>
-    >>> # Both types work seamlessly
-    >>> models: list[ModelName] = ["gpt-5", "custom-llm", "gemini-2.5-pro"]
-
-Note:
-    The ModelName type includes both predefined literals and str,
-    allowing full flexibility while maintaining IDE support for
-    common models.
+The ModelName type includes both predefined literals and str,
+allowing full flexibility while maintaining IDE support for
+common models.
 """

ai_pipeline_core/logging/__init__.py

@@ -2,11 +2,6 @@
 
 Provides a Prefect-integrated logging facade for unified logging across pipelines.
 Prefer get_pipeline_logger instead of logging.getLogger to ensure proper integration.
-
-Example:
-    >>> from ai_pipeline_core import get_pipeline_logger
-    >>> logger = get_pipeline_logger(__name__)
-    >>> logger.info("Processing started")
 """
 
 from .logging_config import LoggingConfig, get_pipeline_logger, setup_logging
@@ -14,8 +9,8 @@ from .logging_mixin import LoggerMixin, StructuredLoggerMixin
 
 __all__ = [
     "LoggerMixin",
-    "StructuredLoggerMixin",
     "LoggingConfig",
-    "setup_logging",
+    "StructuredLoggerMixin",
     "get_pipeline_logger",
+    "setup_logging",
 ]

ai_pipeline_core/logging/logging.yml

@@ -48,7 +48,7 @@ loggers:
   ai_pipeline_core.llm:
     level: INFO
 
-  ai_pipeline_core.flow:
+  ai_pipeline_core.pipeline:
     level: INFO
 
   ai_pipeline_core.testing: