model-library 0.1.0__py3-none-any.whl

model_library/base.py

@@ -0,0 +1,814 @@
+import io
+import logging
+import time
+import uuid
+from abc import ABC, abstractmethod
+from collections.abc import Awaitable
+from pprint import pformat
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    Callable,
+    Literal,
+    Mapping,
+    Sequence,
+    TypeVar,
+    cast,
+)
+
+from pydantic import computed_field, field_validator, model_serializer
+from pydantic.fields import Field
+from pydantic.main import BaseModel
+from typing_extensions import override
+
+from model_library.exceptions import (
+    ImmediateRetryException,
+    retry_llm_call,
+)
+from model_library.utils import sum_optional, truncate_str
+
+PydanticT = TypeVar("PydanticT", bound=BaseModel)
+
+DEFAULT_MAX_TOKENS = 2048
+DEFAULT_TEMPERATURE = 0.7
+DEFAULT_TOP_P = 1
+
+if TYPE_CHECKING:
+    from model_library.providers.openai import OpenAIModel
+
+"""
+--- FILES ---
+"""
+
+
+class FileBase(BaseModel):
+    type: Literal["image", "file"]
+    name: str
+    mime: str
+
+    @override
+    def __repr__(self):
+        attrs = vars(self).copy()
+        if "base64" in attrs:
+            attrs["base64"] = truncate_str(attrs["base64"])
+        return f"{self.__class__.__name__}(\n{pformat(attrs, indent=2)}\n)"
+
+
+class FileWithBase64(FileBase):
+    append_type: Literal["base64"] = "base64"
+    base64: str
+
+
+class FileWithUrl(FileBase):
+    append_type: Literal["url"] = "url"
+    url: str
+
+
+class FileWithId(FileBase):
+    append_type: Literal["file_id"] = "file_id"
+    file_id: str
+
+
+FileInput = Annotated[
+    FileWithBase64 | FileWithUrl | FileWithId,
+    Field(discriminator="append_type"),
+]
+
+
+"""
+--- TOOLS ---
+"""
+
+
+class ToolBody(BaseModel):
+    name: str
+    description: str
+    properties: dict[str, Any]
+    required: list[str]
+    kwargs: dict[str, Any] = {}
+
+
+class ToolDefinition(BaseModel):
+    name: str  # acts as a key
+    body: ToolBody | Any
+
+
+class ToolCall(BaseModel):
+    id: str
+    call_id: str | None = None
+    name: str
+    args: dict[str, Any] | str
+
+
+"""
+--- INPUT ---
+"""
+
+RawResponse = Any
+
+
+class ToolInput(BaseModel):
+    tools: list[ToolDefinition] = []
+
+
+class ToolResult(BaseModel):
+    tool_call: ToolCall
+    result: Any
+
+
+class TextInput(BaseModel):
+    text: str
+
+
+RawInputItem = dict[
+    str, Any
+]  # to pass in, for example, a mock convertsation with {"role": "user", "content": "Hello"}
+
+
+InputItem = (
+    TextInput | FileInput | ToolResult | RawInputItem | RawResponse
+)  # input item can either be a prompt, a file (image or file), a tool call result, raw input, or a previous response
+
+
+"""
+--- OUTPUT ---
+"""
+
+
+class Citation(BaseModel):
+    type: str | None = None
+    title: str | None = None
+    url: str | None = None
+    start_index: int | None = None
+    end_index: int | None = None
+    file_id: str | None = None
+    filename: str | None = None
+    index: int | None = None
+    container_id: str | None = None
+
+
+class QueryResultExtras(BaseModel):
+    citations: list[Citation] = Field(default_factory=list)
+
+
+class QueryResultCost(BaseModel):
+    """
+    Cost information for a query
+    Includes total cost and a structured breakdown.
+    """
+
+    input: float
+    output: float
+    reasoning: float | None = None
+    cache_read: float | None = None
+    cache_write: float | None = None
+
+    @computed_field
+    @property
+    def total(self) -> float:
+        return sum(
+            filter(
+                None,
+                [
+                    self.input,
+                    self.output,
+                    self.reasoning,
+                    self.cache_read,
+                    self.cache_write,
+                ],
+            )
+        )
+
+    @override
+    def __repr__(self):
+        use_cents = self.total < 1
+
+        def format_cost(value: float | None):
+            if value is None:
+                return None
+            return f"{value * 100:.3f} cents" if use_cents else f"${value:.2f}"
+
+        return (
+            f"{format_cost(self.total)} "
+            + f"(uncached input: {format_cost(self.input)} | output: {format_cost(self.output)} | reasoning: {format_cost(self.reasoning)} | cache_read: {format_cost(self.cache_read)} | cache_write: {format_cost(self.cache_write)})"
+        )
+
+
+class QueryResultMetadata(BaseModel):
+    """
+    Metadata for a query: token usage and timing.
+
+    """
+
+    cost: QueryResultCost | None = None  # set post query
+    duration_seconds: float | None = None  # set post query
+    in_tokens: int = 0
+    out_tokens: int = 0
+    reasoning_tokens: int | None = None
+    cache_read_tokens: int | None = None
+    cache_write_tokens: int | None = None
+
+    @property
+    def default_duration_seconds(self) -> float:
+        return self.duration_seconds or 0
+
+    def __add__(self, other: "QueryResultMetadata") -> "QueryResultMetadata":
+        return QueryResultMetadata(
+            in_tokens=self.in_tokens + other.in_tokens,
+            out_tokens=self.out_tokens + other.out_tokens,
+            reasoning_tokens=sum_optional(
+                self.reasoning_tokens, other.reasoning_tokens
+            ),
+            cache_read_tokens=sum_optional(
+                self.cache_read_tokens, other.cache_read_tokens
+            ),
+            cache_write_tokens=sum_optional(
+                self.cache_write_tokens, other.cache_write_tokens
+            ),
+            duration_seconds=self.default_duration_seconds
+            + other.default_duration_seconds,
+        )
+
+    @override
+    def __repr__(self):
+        attrs = vars(self).copy()
+        return f"{self.__class__.__name__}(\n{pformat(attrs, indent=2, sort_dicts=False)}\n)"
+
+
+class QueryResult(BaseModel):
+    """
+    Result of a query
+    Contains the text, reasoning, metadata, tool calls, and history
+    """
+
+    output_text: str | None = None
+    reasoning: str | None = None
+    metadata: QueryResultMetadata = Field(default_factory=QueryResultMetadata)
+    tool_calls: list[ToolCall] = Field(default_factory=list)
+    history: list[InputItem] = Field(default_factory=list)
+    extras: QueryResultExtras = Field(default_factory=QueryResultExtras)
+    raw: dict[str, Any] = Field(default_factory=dict)
+
+    @property
+    def output_text_str(self) -> str:
+        return self.output_text or ""
+
+    @field_validator("reasoning", mode="before")
+    def default_reasoning(cls, v: str | None):
+        return None if not v else v  # make reasoning None if empty
+
+    @property
+    def search_results(self) -> Any | None:
+        """Expose provider-supplied search metadata without additional processing."""
+        raw_dict = cast(dict[str, Any], getattr(self, "raw", {}))
+        raw_candidate = raw_dict.get("search_results")
+        if raw_candidate is not None:
+            return raw_candidate
+
+        return _get_from_history(self.history, "search_results")
+
+    @override
+    def __repr__(self):
+        attrs = vars(self).copy()
+        ordered_attrs = {
+            "output_text": truncate_str(attrs.pop("output_text", None), 400),
+            "reasoning": truncate_str(attrs.pop("reasoning", None), 400),
+            "metadata": attrs.pop("metadata", None),
+        }
+        if self.tool_calls:
+            ordered_attrs["tool_calls"] = self.tool_calls
+        return f"{self.__class__.__name__}(\n{pformat(ordered_attrs, indent=2, sort_dicts=False)}\n)"
+
+
+def _get_from_history(history: Sequence[InputItem], key: str) -> Any | None:
+    for item in reversed(history):
+        value = getattr(item, key, None)
+        if value is not None:
+            return value
+
+        extra = getattr(item, "model_extra", None)
+        if isinstance(extra, Mapping):
+            value = cast(Mapping[str, Any], extra).get(key)
+            if value is not None:
+                return value
+
+    return None
+
+
+class ProviderConfig(BaseModel):
+    """Base class for provider-specific configs. Do not use directly."""
+
+    @model_serializer(mode="plain")
+    def serialize_actual(self):
+        return self.__dict__
+
+
+class LLMConfig(BaseModel):
+    max_tokens: int = DEFAULT_MAX_TOKENS
+    temperature: float | None = None
+    top_p: float | None = None
+    reasoning: bool = False
+    reasoning_effort: str | None = None
+    supports_images: bool = False
+    supports_files: bool = False
+    supports_videos: bool = False
+    supports_batch: bool = False
+    supports_temperature: bool = True
+    supports_tools: bool = False
+    native: bool = True
+    provider_config: ProviderConfig | None = None
+    registry_key: str | None = None
+
+
+RetrierType = Callable[[Callable[..., Awaitable[Any]]], Callable[..., Awaitable[Any]]]
+
+R = TypeVar("R")  # return type
+
+
+class LLM(ABC):
+    """
+    Base class for all LLMs
+    LLM call errors should be raised as exceptions
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        provider: str,
+        *,
+        config: LLMConfig | None = None,
+    ):
+        self.instance_id = uuid.uuid4().hex[:8]
+
+        self.provider: str = provider
+        self.model_name: str = model_name
+
+        config = config or LLMConfig()
+        self._registry_key = config.registry_key
+
+        if config.provider_config:
+            self.provider_config = config.provider_config
+
+        self.max_tokens: int = config.max_tokens
+        self.temperature: float | None = config.temperature
+        self.top_p: float | None = config.top_p
+
+        self.reasoning: bool = config.reasoning
+        self.reasoning_effort: str | None = config.reasoning_effort
+
+        self.supports_files: bool = config.supports_files
+        self.supports_videos: bool = config.supports_videos
+        self.supports_images: bool = config.supports_images
+        self.supports_batch: bool = config.supports_batch
+        self.supports_temperature: bool = config.supports_temperature
+        self.supports_tools: bool = config.supports_tools
+
+        self.native: bool = config.native
+        self.delegate: "OpenAIModel | None" = None
+        self.batch: LLMBatchMixin | None = None
+
+        self.logger: logging.Logger = logging.getLogger(
+            f"llm.{provider}.{model_name}<instance={self.instance_id}>"
+        )
+        self.custom_retrier: Callable[..., RetrierType] | None = retry_llm_call
+
+    @override
+    def __repr__(self):
+        attrs = vars(self).copy()
+        attrs.pop("logger", None)
+        attrs.pop("custom_retrier", None)
+        attrs.pop("_key", None)
+        return f"{self.__class__.__name__}(\n{pformat(attrs, indent=2, sort_dicts=False)}\n)"
+
+    @abstractmethod
+    def get_client(self) -> object:
+        """Return the instance of the appropriate SDK client."""
+        ...
+
+    @staticmethod
+    async def timer_wrapper(func: Callable[[], Awaitable[R]]) -> tuple[R, float]:
+        """
+        Time the query
+        """
+        start = time.perf_counter()
+        result = await func()
+        return result, round(time.perf_counter() - start, 4)
+
+    @staticmethod
+    async def immediate_retry_wrapper(
+        func: Callable[[], Awaitable[R]],
+        logger: logging.Logger,
+    ) -> R:
+        """
+        Retry the query immediately
+        """
+        MAX_IMMEDIATE_RETRIES = 10
+        retries = 0
+        while True:
+            try:
+                return await func()
+            except ImmediateRetryException as e:
+                if retries >= MAX_IMMEDIATE_RETRIES:
+                    logger.error(f"Query reached max immediate retries {retries}: {e}")
+                    raise Exception(
+                        f"Query reached max immediate retries {retries}: {e}"
+                    ) from e
+                retries += 1
+
+                logger.warning(
+                    f"Query retried immediately {retries}/{MAX_IMMEDIATE_RETRIES}: {e}"
+                )
+
+    @staticmethod
+    async def backoff_retry_wrapper(
+        func: Callable[..., Awaitable[R]],
+        backoff_retrier: RetrierType | None,
+    ) -> R:
+        """
+        Retry the query with backoff
+        """
+        if not backoff_retrier:
+            return await func()
+        return await backoff_retrier(func)()
+
+    async def delegate_query(
+        self,
+        input: Sequence[InputItem],
+        *,
+        tools: list[ToolDefinition] = [],
+        **kwargs: object,
+    ) -> QueryResult:
+        if not self.delegate:
+            raise Exception("Delegate not set")
+        return await self.delegate._query_impl(input, tools=tools, **kwargs)  # pyright: ignore[reportPrivateUsage]
+
+    async def query(
+        self,
+        input: Sequence[InputItem] | str,
+        *,
+        history: Sequence[InputItem] = [],
+        tools: list[ToolDefinition] = [],
+        # for backwards compatibility
+        files: list[FileInput] = [],
+        images: list[FileInput] = [],
+        **kwargs: object,
+    ) -> QueryResult:
+        """
+        Query the model
+        Join input with history
+        Log, Time, and Retry
+        """
+        # format str input
+        if isinstance(input, str):
+            input = [TextInput(text=input)]
+
+        # prepends files and images to input
+        input = [*files, *images, *input]
+
+        # format input info
+        item_info = f"--- input ({len(input)}): {get_pretty_input_types(input)}\n"
+        if history:
+            item_info += (
+                f"--- history({len(history)}): {get_pretty_input_types(history)}\n"
+            )
+
+        # format tool info
+        tool_results = [t for t in input if isinstance(t, ToolResult)]
+        tool_names = [tool.name for tool in tools or []]
+
+        tool_info = (
+            f"--- tools ({len(tools)}): {tool_names}\n"
+            + f"--- tool results ({len(tool_results)}): "
+            + f"{[{tool.tool_call.name: truncate_str(str(tool.result))} for tool in tool_results]}\n"
+            if tools
+            else ""
+        )
+
+        short_kwargs = {k: truncate_str(repr(v)) for k, v in kwargs.items()}
+
+        # join input with history
+        input = [*history, *input]
+
+        # unique logger for the query
+        query_id = uuid.uuid4().hex[:14]
+        query_logger = logging.getLogger(f"{self.logger.name}<query={query_id}>")
+
+        query_logger.info(
+            "Query started:\n" + item_info + tool_info + f"--- kwargs: {short_kwargs}\n"
+        )
+
+        async def query_func() -> QueryResult:
+            return await self._query_impl(input, tools=tools, **kwargs)
+
+        async def timed_query() -> tuple[QueryResult, float]:
+            return await LLM.timer_wrapper(query_func)
+
+        async def immediate_retry() -> tuple[QueryResult, float]:
+            return await LLM.immediate_retry_wrapper(timed_query, query_logger)
+
+        async def backoff_retry() -> tuple[QueryResult, float]:
+            backoff_retrier = (
+                self.custom_retrier(query_logger) if self.custom_retrier else None
+            )
+            return await LLM.backoff_retry_wrapper(immediate_retry, backoff_retrier)
+
+        output, duration = await backoff_retry()
+        output.metadata.duration_seconds = duration
+        output.metadata.cost = await self._calculate_cost(output.metadata)
+
+        query_logger.info(f"Query completed: {repr(output)}")
+
+        return output
+
+    async def query_json(
+        self,
+        input: Sequence[InputItem],
+        pydantic_model: type[PydanticT],
+        **kwargs: object,
+    ) -> PydanticT:
+        """Query the model with JSON response format using Pydantic model.
+
+        This is a convenience method that is not implemented for all providers.
+        Only OpenAI and Google providers currently support this method.
+
+        Args:
+            input: Input items (text, files, etc.)
+            pydantic_model: Pydantic model class defining the expected response structure
+            **kwargs: Additional arguments passed to the query method
+
+        Returns:
+            Instance of the pydantic_model with the model's response
+
+        Raises:
+            NotImplementedError: If the provider does not support structured JSON output
+        """
+        raise NotImplementedError(
+            f"query_json is not implemented for {self.__class__.__name__}. "
+            f"Only OpenAI and Google providers currently support this method."
+        )
+
+    async def _calculate_cost(
+        self,
+        metadata: QueryResultMetadata,
+        batch: bool = False,
+        bill_reasoning: bool = True,
+    ) -> QueryResultCost | None:
+        """Calculate cost for a query"""
+        from model_library.registry_utils import get_model_cost
+
+        if not self._registry_key:
+            self.logger.warning("Model has no registry key, skipping cost calculation")
+            return None
+
+        costs = get_model_cost(self._registry_key)
+        if not costs:
+            return None
+
+        MILLION = 1_000_000
+
+        # base input and output
+        if costs.input is None or costs.output is None:
+            raise Exception("Base costs not set")
+        input_cost = costs.input
+        output_cost = costs.output
+
+        # apply fixed values or discounts/markup
+        # applied before other price changes
+        cache_read_cost, cache_write_cost = None, None
+        if metadata.cache_read_tokens or metadata.cache_write_tokens:
+            if not costs.cache:
+                raise Exception("Cache costs not set")
+            cache_read_cost, cache_write_cost = costs.cache.get_costs(
+                input_cost, output_cost
+            )
+
+        # costs for long context
+        total_in = metadata.in_tokens + (metadata.cache_read_tokens or 0)
+        if costs.context and total_in > costs.context.threshold:
+            input_cost, output_cost = costs.context.get_costs(
+                input_cost,
+                output_cost,
+                total_in,
+            )
+            if costs.context.cache:
+                cache_read_cost, cache_write_cost = costs.context.cache.get_costs(
+                    input_cost, output_cost
+                )
+
+        # costs for batching
+        if batch:
+            if not costs.batch:
+                raise Exception("Batch costs not set")
+            input_cost, output_cost = costs.batch.get_costs(input_cost, output_cost)
+
+        return QueryResultCost(
+            input=input_cost * metadata.in_tokens / MILLION,
+            output=output_cost * metadata.out_tokens / MILLION,
+            reasoning=output_cost * metadata.reasoning_tokens / MILLION
+            if metadata.reasoning_tokens is not None and bill_reasoning
+            else None,
+            cache_read=cache_read_cost * metadata.cache_read_tokens / MILLION
+            if metadata.cache_read_tokens is not None and cache_read_cost
+            else None,
+            cache_write=cache_write_cost * metadata.cache_write_tokens / MILLION
+            if metadata.cache_write_tokens is not None and cache_write_cost
+            else None,
+        )
+
+    @abstractmethod
+    async def _query_impl(
+        self,
+        input: Sequence[InputItem],
+        *,
+        tools: list[ToolDefinition],
+        **kwargs: object,  # TODO: pass in query logger
+    ) -> QueryResult:
+        """
+        Query the model with input
+        Input can consist on text, images, files, or model specific raw responses
+        Optionally pass in tools
+        Kwargs will be passed to the model call (apart from exceptions like system_prompt)
+        Images and files should be preprocessed according to what the model supports:
+            - base64
+            - url
+            - file_id
+        """
+        ...
+
+    @abstractmethod
+    async def parse_input(
+        self,
+        input: Sequence[InputItem],
+        **kwargs: Any,
+    ) -> Any:
+        """
+        Parses input into the appropriate format for the model
+        Handles prompts, images, and files
+        Handles history and tool call results
+        Calls
+            - parse_image
+            - parse_file
+        """
+        ...
+
+    @abstractmethod
+    async def parse_image(self, image: FileInput) -> Any:
+        """Parse an image into the appropriate format for the model"""
+        ...
+
+    @abstractmethod
+    async def parse_file(self, file: FileInput) -> Any:
+        """Parse a file into the appropriate format for the model"""
+        ...
+
+    @abstractmethod
+    async def parse_tools(self, tools: list[ToolDefinition]) -> Any:
+        """Parse tools into the appropriate format for the model"""
+        ...
+
+    @abstractmethod
+    async def upload_file(
+        self,
+        name: str,
+        mime: str,
+        bytes: io.BytesIO,
+        type: Literal["image", "file"] = "file",
+    ) -> FileWithId:
+        """Upload a file to the model provider"""
+        ...
+
+
+class BatchResult(BaseModel):
+    custom_id: str
+    output: QueryResult
+    error_message: str | None = None
+
+
+class LLMBatchMixin(ABC):
+    @abstractmethod
+    async def create_batch_query_request(
+        self,
+        custom_id: str,
+        input: Sequence[InputItem],
+        **kwargs: object,
+    ) -> dict[str, Any]:
+        """Return a single query request
+
+        The batch api sends out a batch of query requests to various endpoints.
+
+        For example OpenAI sends can send requests to /v1/responses or /v1/chat/completions endpoints.
+
+        This method creates a query request for methods such methods
+        """
+        ...
+
+    @abstractmethod
+    async def batch_query(
+        self,
+        batch_name: str,
+        requests: list[dict[str, Any]],
+    ) -> str:
+        """
+        Batch query the model
+        Returns:
+            str: batch_id
+        Raises:
+            Exception: If failed to batch query
+        """
+        ...
+
+    @abstractmethod
+    async def get_batch_results(self, batch_id: str) -> list[BatchResult]:
+        """
+        Returns results for batch
+        Raises:
+            Exception: If failed to get results
+        """
+        ...
+
+    @abstractmethod
+    async def get_batch_progress(self, batch_id: str) -> int:
+        """
+        Returns number of completed requests for batch
+        Raises:
+            Exception: If failed to get progress
+        """
+        ...
+
+    @abstractmethod
+    async def cancel_batch_request(self, batch_id: str) -> None:
+        """
+        Cancels batch
+        Raises:
+            Exception: If failed to cancel
+        """
+        ...
+
+    @abstractmethod
+    async def get_batch_status(
+        self,
+        batch_id: str,
+    ) -> str:
+        """
+        Returns batch status
+        Raises:
+            Exception: If failed to get status
+        """
+        ...
+
+    @classmethod
+    @abstractmethod
+    def is_batch_status_completed(
+        cls,
+        batch_status: str,
+    ) -> bool:
+        """
+        Returns if batch status is completed
+
+        A completed state is any state that is final and not in-progress
+        Example: failed | cancelled | expired | completed
+
+        An incompleted state is any state that is not completed
+        Example: in_progress | pending | running
+        """
+        ...
+
+    @classmethod
+    @abstractmethod
+    def is_batch_status_failed(
+        cls,
+        batch_status: str,
+    ) -> bool:
+        """Returns if batch status is failed"""
+        ...
+
+    @classmethod
+    @abstractmethod
+    def is_batch_status_cancelled(
+        cls,
+        batch_status: str,
+    ) -> bool:
+        """Returns if batch status is cancelled"""
+        ...
+
+
+def get_pretty_input_types(input: Sequence["InputItem"]) -> str:
+    # for logging
+    def process_item(item: "InputItem"):
+        match item:
+            case TextInput():
+                return truncate_str(repr(item))
+            case FileBase():  # FileInput
+                return repr(item)
+            case ToolResult():
+                return repr(item)
+            case dict():
+                item = cast(RawInputItem, item)
+                return repr(item)
+            case _:
+                # RawResponse
+                return repr(item)
+
+    processed_items = [f"  {process_item(item)}" for item in input]
+    return "\n" + "\n".join(processed_items) if processed_items else ""