model-library 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

model_library/__init__.py

@@ -1,8 +1,11 @@
-from model_library.settings import ModelLibrarySettings
 from model_library.base import LLM, LLMConfig
+from model_library.logging import set_logging
+from model_library.settings import ModelLibrarySettings
 
 model_library_settings: ModelLibrarySettings = ModelLibrarySettings()
 
+set_logging()
+
 
 def model(model_str: str, override_config: LLMConfig | None = None) -> LLM:
     from model_library.registry_utils import get_registry_model
@@ -10,14 +13,15 @@ def model(model_str: str, override_config: LLMConfig | None = None) -> LLM:
     return get_registry_model(model_str, override_config)
 
 
-def raw_model(model_str: str, override_config: LLMConfig | None = None) -> LLM:
+def raw_model(model_str: str, config: LLMConfig | None = None) -> LLM:
     from model_library.registry_utils import get_raw_model
 
-    return get_raw_model(model_str, config=override_config)
+    return get_raw_model(model_str, config=config)
 
 
 __all__ = [
     "model_library_settings",
     "model",
     "raw_model",
+    "set_logging",
 ]

model_library/base/__init__.py

@@ -0,0 +1,7 @@
+# ruff: noqa: F403,F401
+from model_library.base.base import *
+from model_library.base.batch import *
+from model_library.base.delegate_only import *
+from model_library.base.input import *
+from model_library.base.output import *
+from model_library.base.utils import *

model_library/{base.py → base/base.py}

@@ -7,293 +7,46 @@ 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 import model_serializer
 from pydantic.main import BaseModel
 from typing_extensions import override
 
+from model_library.base.batch import (
+    LLMBatchMixin,
+)
+from model_library.base.input import (
+    FileInput,
+    FileWithId,
+    InputItem,
+    TextInput,
+    ToolDefinition,
+    ToolResult,
+)
+from model_library.base.output import (
+    QueryResult,
+    QueryResultCost,
+    QueryResultMetadata,
+)
+from model_library.base.utils import (
+    get_pretty_input_types,
+)
 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
+from model_library.utils import truncate_str
 
 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
+PydanticT = TypeVar("PydanticT", bound=BaseModel)
 
 
 class ProviderConfig(BaseModel):
@@ -304,6 +57,9 @@ class ProviderConfig(BaseModel):
         return self.__dict__
 
 
+DEFAULT_MAX_TOKENS = 2048
+
+
 class LLMConfig(BaseModel):
     max_tokens: int = DEFAULT_MAX_TOKENS
     temperature: float | None = None
@@ -347,9 +103,6 @@ class LLM(ABC):
         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
@@ -368,6 +121,12 @@ class LLM(ABC):
         self.delegate: "OpenAIModel | None" = None
         self.batch: LLMBatchMixin | None = None
 
+        if config.provider_config:
+            if isinstance(
+                config.provider_config, type(getattr(self, "provider_config"))
+            ):
+                self.provider_config = config.provider_config
+
         self.logger: logging.Logger = logging.getLogger(
             f"llm.{provider}.{model_name}<instance={self.instance_id}>"
         )
@@ -521,33 +280,6 @@ class LLM(ABC):
 
         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,
@@ -678,137 +410,29 @@ class LLM(ABC):
         """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(
+    async def query_json(
         self,
-        custom_id: str,
         input: Sequence[InputItem],
+        pydantic_model: type[PydanticT],
         **kwargs: object,
-    ) -> dict[str, Any]:
-        """Return a single query request
-
-        The batch api sends out a batch of query requests to various endpoints.
+    ) -> PydanticT:
+        """Query the model with JSON response format using Pydantic model.
 
-        For example OpenAI sends can send requests to /v1/responses or /v1/chat/completions endpoints.
+        This is a convenience method that is not implemented for all providers.
+        Only OpenAI and Google providers currently support this method.
 
-        This method creates a query request for methods such methods
-        """
-        ...
+        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
 
-    @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
-        """
-        ...
+            Instance of the pydantic_model with the model's response
 
-    @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
+            NotImplementedError: If the provider does not support structured JSON output
         """
-        ...
-
-    @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 ""
+        raise NotImplementedError(
+            f"query_json is not implemented for {self.__class__.__name__}. "
+            f"Only OpenAI and Google providers currently support this method."
+        )

model_library/base/batch.py

@@ -0,0 +1,121 @@
+from abc import ABC, abstractmethod
+from typing import Any, Sequence
+
+from pydantic import BaseModel
+
+from model_library.base.input import InputItem
+from model_library.base.output import QueryResult
+
+
+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"""
+        ...