ai-pipeline-core 0.1.1__py3-none-any.whl

ai_pipeline_core/flow/config.py

@@ -0,0 +1,66 @@
+"""Flow configuration base class."""
+
+from abc import ABC
+from typing import ClassVar
+
+from ai_pipeline_core.documents import DocumentList, FlowDocument
+
+
+class FlowConfig(ABC):
+    """
+    Configuration for a flow. It makes flow easier to implement and test.
+    """
+
+    INPUT_DOCUMENT_TYPES: ClassVar[list[type[FlowDocument]]]
+    OUTPUT_DOCUMENT_TYPE: ClassVar[type[FlowDocument]]
+
+    @classmethod
+    def get_input_document_types(cls) -> list[type[FlowDocument]]:
+        """
+        Get the input document types for the flow.
+        """
+        return cls.INPUT_DOCUMENT_TYPES
+
+    @classmethod
+    def get_output_document_type(cls) -> type[FlowDocument]:
+        """
+        Get the output document type for the flow.
+        """
+        return cls.OUTPUT_DOCUMENT_TYPE
+
+    @classmethod
+    def has_input_documents(cls, documents: DocumentList) -> bool:
+        """
+        Check if the flow has all required input documents.
+        """
+        for doc_cls in cls.INPUT_DOCUMENT_TYPES:
+            if not any(isinstance(doc, doc_cls) for doc in documents):
+                return False
+        return True
+
+    @classmethod
+    def get_input_documents(cls, documents: DocumentList) -> DocumentList:
+        """
+        Get the input documents for the flow.
+        """
+        input_documents = DocumentList()
+        for doc_cls in cls.INPUT_DOCUMENT_TYPES:
+            filtered_documents = [doc for doc in documents if isinstance(doc, doc_cls)]
+            if not filtered_documents:
+                raise ValueError(f"No input document found for class {doc_cls.__name__}")
+            input_documents.extend(filtered_documents)
+        return input_documents
+
+    @classmethod
+    def validate_output_documents(cls, documents: DocumentList) -> None:
+        """
+        Validate the output documents of the flow.
+        """
+        assert isinstance(documents, DocumentList), "Documents must be a DocumentList"
+        output_document_class = cls.get_output_document_type()
+
+        invalid = [type(d).__name__ for d in documents if not isinstance(d, output_document_class)]
+        assert not invalid, (
+            "Documents must be of the correct type. "
+            f"Expected: {output_document_class.__name__}, Got invalid: {invalid}"
+        )

ai_pipeline_core/llm/__init__.py

@@ -0,0 +1,19 @@
+from .ai_messages import AIMessages, AIMessageType
+from .client import (
+    generate,
+    generate_structured,
+)
+from .model_options import ModelOptions
+from .model_response import ModelResponse, StructuredModelResponse
+from .model_types import ModelName
+
+__all__ = [
+    "AIMessages",
+    "AIMessageType",
+    "ModelName",
+    "ModelOptions",
+    "ModelResponse",
+    "StructuredModelResponse",
+    "generate",
+    "generate_structured",
+]

ai_pipeline_core/llm/ai_messages.py

@@ -0,0 +1,129 @@
+import base64
+import json
+
+from openai.types.chat import (
+    ChatCompletionContentPartParam,
+    ChatCompletionMessageParam,
+)
+from prefect.logging import get_logger
+
+from ai_pipeline_core.documents import Document
+
+from .model_response import ModelResponse
+
+AIMessageType = str | Document | ModelResponse
+
+
+class AIMessages(list[AIMessageType]):
+    def get_last_message(self) -> AIMessageType:
+        return self[-1]
+
+    def get_last_message_as_str(self) -> str:
+        last_message = self.get_last_message()
+        if isinstance(last_message, str):
+            return last_message
+        raise ValueError(f"Wrong message type: {type(last_message)}")
+
+    def to_prompt(self) -> list[ChatCompletionMessageParam]:
+        """Convert AIMessages to OpenAI-compatible format.
+
+        Returns:
+            List of ChatCompletionMessageParam for OpenAI API
+        """
+        messages: list[ChatCompletionMessageParam] = []
+
+        for message in self:
+            if isinstance(message, str):
+                messages.append({"role": "user", "content": 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})
+            else:
+                raise ValueError(f"Unsupported message type: {type(message)}")
+
+        return messages
+
+    def to_tracing_log(self) -> list[str]:
+        """Convert AIMessages to a list of strings for tracing."""
+        messages: list[str] = []
+        for message in self:
+            if isinstance(message, Document):
+                serialized_document = message.serialize_model()
+                del serialized_document["content"]
+                messages.append(json.dumps(serialized_document, indent=2))
+            elif isinstance(message, ModelResponse):
+                messages.append(message.content)
+            else:
+                assert isinstance(message, str)
+                messages.append(message)
+        return messages
+
+    @staticmethod
+    def document_to_prompt(document: Document) -> list[ChatCompletionContentPartParam]:
+        """
+        Convert a document to prompt format for LLM consumption.
+
+        Args:
+            document: The document to convert
+
+        Returns:
+            List of chat completion content parts for the prompt
+        """
+        prompt: list[ChatCompletionContentPartParam] = []
+
+        # Build the text header
+        description = (
+            f"<description>{document.description}</description>\n" if document.description else ""
+        )
+        header_text = (
+            f"<document>\n<id>{document.id}</id>\n<name>{document.name}</name>\n{description}"
+        )
+
+        # Handle text documents
+        if document.is_text:
+            content_text = (
+                f"{header_text}<content>\n{document.as_text()}\n</content>\n</document>\n"
+            )
+            prompt.append({"type": "text", "text": content_text})
+            return prompt
+
+        # Handle non-text documents
+        if not document.is_image and not document.is_pdf:
+            get_logger(__name__).error(
+                f"Document is not a text, image or PDF: {document.name} - {document.mime_type}"
+            )
+            return []
+
+        # Add header for binary content
+        prompt.append(
+            {
+                "type": "text",
+                "text": f"{header_text}<content>\n",
+            }
+        )
+
+        # Encode binary content
+        base64_content = base64.b64encode(document.content).decode("utf-8")
+        data_uri = f"data:{document.mime_type};base64,{base64_content}"
+
+        # Add appropriate content type
+        if document.is_pdf:
+            prompt.append(
+                {
+                    "type": "file",
+                    "file": {"file_data": data_uri},
+                }
+            )
+        else:  # is_image
+            prompt.append(
+                {
+                    "type": "image_url",
+                    "image_url": {"url": data_uri, "detail": "high"},
+                }
+            )
+
+        # Close the document tag
+        prompt.append({"type": "text", "text": "</content>\n</document>\n"})
+
+        return prompt

ai_pipeline_core/llm/client.py

@@ -0,0 +1,218 @@
+import asyncio
+from typing import Any, TypeVar
+
+from lmnr import Laminar
+from openai import AsyncOpenAI
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+)
+from prefect.logging import get_logger
+from pydantic import BaseModel
+
+from ai_pipeline_core.exceptions import LLMError
+from ai_pipeline_core.settings import settings
+from ai_pipeline_core.tracing import trace
+
+from .ai_messages import AIMessages
+from .model_options import ModelOptions
+from .model_response import ModelResponse, StructuredModelResponse
+from .model_types import ModelName
+
+logger = get_logger()
+
+
+def _process_messages(
+    context: AIMessages,
+    messages: AIMessages,
+    system_prompt: str | None = None,
+) -> list[ChatCompletionMessageParam]:
+    """Convert context and messages to OpenAI-compatible format.
+
+    Args:
+        context: Messages to be cached (optional)
+        messages: Regular messages that won't be cached
+        system_prompt: Optional system prompt
+
+    Returns:
+        List of formatted messages for OpenAI API
+    """
+
+    processed_messages: list[ChatCompletionMessageParam] = []
+
+    # Add system prompt if provided
+    if system_prompt:
+        processed_messages.append({"role": "system", "content": system_prompt})
+
+    # Process context messages with caching if provided
+    if context:
+        # Use AIMessages.to_prompt() for context
+        context_messages = context.to_prompt()
+
+        # Apply caching to context messages
+        for msg in context_messages:
+            if msg.get("role") == "user":
+                # Add cache control to user messages in context
+                msg["cache_control"] = {  # type: ignore
+                    "type": "ephemeral",
+                    "ttl": "120s",  # Cache for 2m
+                }
+            processed_messages.append(msg)
+
+    # Process regular messages without caching
+    if messages:
+        regular_messages = messages.to_prompt()
+        processed_messages.extend(regular_messages)
+
+    return processed_messages
+
+
+async def _generate(
+    model: str, messages: list[ChatCompletionMessageParam], completion_kwargs: dict[str, Any]
+) -> ModelResponse:
+    async with AsyncOpenAI(
+        api_key=settings.openai_api_key,
+        base_url=settings.openai_base_url,
+    ) as client:
+        # Use parse for structured output, create for regular
+        if completion_kwargs.get("response_format"):
+            raw_response = await client.chat.completions.with_raw_response.parse(  # type: ignore[var-annotated]
+                **completion_kwargs,
+            )
+        else:
+            raw_response = await client.chat.completions.with_raw_response.create(  # type: ignore[var-annotated]
+                **completion_kwargs
+            )
+
+        response = ModelResponse(raw_response.parse())  # type: ignore[arg-type]
+        response.set_model_options(completion_kwargs)
+        response.set_headers(dict(raw_response.headers.items()))  # type: ignore[arg-type]
+        return response
+
+
+async def _generate_with_retry(
+    model: str,
+    context: AIMessages,
+    messages: AIMessages,
+    options: ModelOptions,
+) -> ModelResponse:
+    """Core generation logic with exponential backoff retry."""
+    if not model:
+        raise ValueError("Model must be provided")
+    if not context and not messages:
+        raise ValueError("Either context or messages must be provided")
+
+    processed_messages = _process_messages(context, messages, options.system_prompt)
+    completion_kwargs: dict[str, Any] = {
+        "model": model,
+        "messages": processed_messages,
+        **options.to_openai_completion_kwargs(),
+    }
+
+    for attempt in range(options.retries):
+        try:
+            with Laminar.start_as_current_span(model, span_type="LLM", input=messages) as span:
+                response = await _generate(model, processed_messages, completion_kwargs)
+                span.set_attributes(response.get_laminar_metadata())
+                Laminar.set_span_output(response.content)
+                if not response.content:
+                    # disable cache in case of empty response
+                    completion_kwargs["extra_body"]["cache"] = {"no-cache": True}
+                    raise ValueError(f"Model {model} returned an empty response.")
+                return response
+        except (asyncio.TimeoutError, ValueError, Exception) as e:
+            logger.warning(
+                "LLM generation failed (attempt %d/%d): %s",
+                attempt + 1,
+                options.retries,
+                e,
+            )
+            if attempt == options.retries - 1:
+                raise LLMError("Exhausted all retry attempts for LLM generation.") from e
+
+        await asyncio.sleep(options.retry_delay_seconds)
+
+    raise LLMError("Unknown error occurred during LLM generation.")
+
+
+@trace(ignore_inputs=["context"])
+async def generate(
+    model: ModelName | str,
+    *,
+    context: AIMessages = AIMessages(),
+    messages: AIMessages | str,
+    options: ModelOptions = ModelOptions(),
+) -> ModelResponse:
+    """Generate response using a large or small model.
+
+    Args:
+        model: The model to use for generation
+        context: Messages to be cached (optional) - keyword only
+        messages: Regular messages that won't be cached - keyword only
+        options: Model options - keyword only
+
+    Returns:
+        Model response
+    """
+    if isinstance(messages, str):
+        messages = AIMessages([messages])
+
+    return await _generate_with_retry(model, context, messages, options)
+
+
+T = TypeVar("T", bound=BaseModel)
+
+
+@trace
+async def generate_structured(
+    model: ModelName,
+    response_format: type[T],
+    *,
+    context: AIMessages = AIMessages(),
+    messages: AIMessages | str,
+    options: ModelOptions = ModelOptions(),
+) -> StructuredModelResponse[T]:
+    """Generate structured response using Pydantic models.
+
+    Args:
+        model: The model to use for generation
+        response_format: A Pydantic model class
+        context: Messages to be cached (optional) - keyword only
+        messages: Regular messages that won't be cached - keyword only
+        options: Model options - keyword only
+
+    Returns:
+        A StructuredModelResponse containing the parsed Pydantic model instance
+    """
+    options.response_format = response_format
+
+    if isinstance(messages, str):
+        messages = AIMessages([messages])
+
+    # Call the internal generate function with structured output enabled
+    response = await _generate_with_retry(model, context, messages, options)
+
+    # Extract the parsed value from the response
+    parsed_value: T | None = None
+
+    # Check if response has choices and parsed content
+    if response.choices and hasattr(response.choices[0].message, "parsed"):
+        parsed: Any = response.choices[0].message.parsed  # type: ignore[attr-defined]
+
+        # If parsed is a dict, instantiate it as the response format class
+        if isinstance(parsed, dict):
+            parsed_value = response_format(**parsed)
+        # If it's already the right type, use it
+        elif isinstance(parsed, response_format):
+            parsed_value = parsed
+        else:
+            # Otherwise try to convert it
+            raise TypeError(
+                f"Unable to convert parsed response to {response_format.__name__}: "
+                f"got type {type(parsed).__name__}"  # type: ignore[reportUnknownArgumentType]
+            )
+
+    if parsed_value is None:
+        raise ValueError("No parsed content available from the model response")
+
+    # Create a StructuredModelResponse with the parsed value
+    return StructuredModelResponse[T](chat_completion=response, parsed_value=parsed_value)

ai_pipeline_core/llm/model_options.py

@@ -0,0 +1,39 @@
+from typing import Any, Literal
+
+from pydantic import BaseModel
+
+
+class ModelOptions(BaseModel):
+    system_prompt: str | None = None
+    search_context_size: Literal["low", "medium", "high"] | None = None
+    reasoning_effort: Literal["low", "medium", "high"] | None = None
+    retries: int = 3
+    retry_delay_seconds: int = 10
+    timeout: int = 300
+    service_tier: Literal["auto", "default", "flex", "scale", "priority"] | None = None
+    max_completion_tokens: int | None = None
+    response_format: type[BaseModel] | None = None
+
+    def to_openai_completion_kwargs(self) -> dict[str, Any]:
+        """Convert ModelOptions to OpenAI completion kwargs."""
+        kwargs: dict[str, Any] = {
+            "timeout": self.timeout,
+            "extra_body": {},
+        }
+
+        if self.max_completion_tokens:
+            kwargs["max_completion_tokens"] = self.max_completion_tokens
+
+        if self.reasoning_effort:
+            kwargs["reasoning_effort"] = self.reasoning_effort
+
+        if self.search_context_size:
+            kwargs["web_search_options"] = {"search_context_size": self.search_context_size}
+
+        if self.response_format:
+            kwargs["response_format"] = self.response_format
+
+        if self.service_tier:
+            kwargs["service_tier"] = self.service_tier
+
+        return kwargs

ai_pipeline_core/llm/model_response.py

@@ -0,0 +1,149 @@
+import copy
+from typing import Any, Generic, TypeVar
+
+from openai.types.chat import ChatCompletion, ParsedChatCompletion
+from pydantic import BaseModel, Field
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class ModelResponse(ChatCompletion):
+    """Response from an LLM without structured output."""
+
+    headers: dict[str, str] = Field(default_factory=dict)
+    model_options: dict[str, Any] = Field(default_factory=dict)
+
+    def __init__(self, chat_completion: ChatCompletion | None = None, **kwargs: Any) -> None:
+        """Initialize ModelResponse from a ChatCompletion."""
+        if chat_completion:
+            # Copy all attributes from the ChatCompletion instance
+            data = chat_completion.model_dump()
+            data["headers"] = {}  # Add default headers
+            super().__init__(**data)
+        else:
+            # Initialize from kwargs
+            if "headers" not in kwargs:
+                kwargs["headers"] = {}
+            super().__init__(**kwargs)
+
+    @property
+    def content(self) -> str:
+        """Get the text content of the response."""
+        return self.choices[0].message.content or ""
+
+    def set_model_options(self, options: dict[str, Any]) -> None:
+        """Set the model options."""
+        self.model_options = copy.deepcopy(options)
+        if "messages" in self.model_options:
+            del self.model_options["messages"]
+
+    def set_headers(self, headers: dict[str, str]) -> None:
+        """Set the response headers."""
+        self.headers = copy.deepcopy(headers)
+
+    def get_laminar_metadata(self) -> dict[str, str | int | float]:
+        """Extract metadata for Laminar observability logging."""
+        metadata: dict[str, str | int | float] = {}
+
+        litellm_id = self.headers.get("x-litellm-call-id")
+        cost = float(self.headers.get("x-litellm-response-cost") or 0)
+
+        # Add all x-litellm-* headers
+        for header, value in self.headers.items():
+            if header.startswith("x-litellm-"):
+                header_name = header.replace("x-litellm-", "").lower()
+                metadata[f"litellm.{header_name}"] = value
+
+        # Add base metadata
+        metadata.update(
+            {
+                "gen_ai.response.id": litellm_id or self.id,
+                "gen_ai.response.model": self.model,
+                "get_ai.system": "litellm",
+            }
+        )
+
+        # Add usage metadata if available
+        if self.usage:
+            metadata.update(
+                {
+                    "gen_ai.usage.prompt_tokens": self.usage.prompt_tokens,
+                    "gen_ai.usage.completion_tokens": self.usage.completion_tokens,
+                    "gen_ai.usage.total_tokens": self.usage.total_tokens,
+                }
+            )
+
+            # Check for cost in usage object
+            if hasattr(self.usage, "cost"):
+                # The 'cost' attribute is added by LiteLLM but not in OpenAI types
+                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
+
+            # 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
+
+        # 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,
+                }
+            )
+
+        if self.model_options:
+            for key, value in self.model_options.items():
+                metadata[f"model_options.{key}"] = str(value)
+
+        return metadata
+
+
+class StructuredModelResponse(ModelResponse, Generic[T]):
+    """Response from an LLM with structured output of type T."""
+
+    def __init__(
+        self,
+        chat_completion: ChatCompletion | None = None,
+        parsed_value: T | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize StructuredModelResponse with a parsed value.
+
+        Args:
+            chat_completion: The base chat completion
+            parsed_value: The parsed structured output
+            **kwargs: Additional arguments for ChatCompletion
+        """
+        super().__init__(chat_completion, **kwargs)
+        self._parsed_value: T | None = parsed_value
+
+        # Extract parsed value from ParsedChatCompletion if available
+        if chat_completion and isinstance(chat_completion, ParsedChatCompletion):
+            if chat_completion.choices:  # type: ignore[attr-defined]
+                message = chat_completion.choices[0].message  # type: ignore[attr-defined]
+                if hasattr(message, "parsed"):  # type: ignore
+                    self._parsed_value = message.parsed  # type: ignore[attr-defined]
+
+    @property
+    def parsed(self) -> T:
+        """Get the parsed structured output.
+
+        Returns:
+            The parsed value of type T.
+
+        Raises:
+            ValueError: If no parsed content is available.
+        """
+        if self._parsed_value is not None:
+            return self._parsed_value
+
+        raise ValueError(
+            "No parsed content available. This should not happen for StructuredModelResponse."
+        )

ai_pipeline_core/llm/model_types.py

@@ -0,0 +1,17 @@
+from typing import Literal, TypeAlias
+
+ModelName: TypeAlias = Literal[
+    # Core models
+    "gemini-2.5-pro",
+    "gpt-5",
+    "grok-4",
+    # Small models
+    "gemini-2.5-flash",
+    "gpt-5-mini",
+    "grok-3-mini",
+    # Search models
+    "gemini-2.5-flash-search",
+    "sonar-pro-search",
+    "gpt-4o-search",
+    "grok-3-mini-search",
+]

ai_pipeline_core/logging/__init__.py

@@ -0,0 +1,10 @@
+from .logging_config import LoggingConfig, get_pipeline_logger, setup_logging
+from .logging_mixin import LoggerMixin, StructuredLoggerMixin
+
+__all__ = [
+    "LoggerMixin",
+    "StructuredLoggerMixin",
+    "LoggingConfig",
+    "setup_logging",
+    "get_pipeline_logger",
+]