model-library 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

model_library/base/delegate_only.py

@@ -0,0 +1,94 @@
+import io
+from typing import Any, Literal, Sequence
+
+from typing_extensions import override
+
+from model_library.base import (
+    LLM,
+    FileInput,
+    FileWithId,
+    InputItem,
+    LLMConfig,
+    QueryResult,
+    ToolDefinition,
+)
+
+
+class DelegateOnlyException(Exception):
+    """
+    Raised when native model functionality is performed on a
+    delegate-only model.
+    """
+
+    DEFAULT_MESSAGE: str = "This model supports only delegate-only functionality. Only the query() method should be used."
+
+    def __init__(self, message: str | None = None):
+        super().__init__(message or DelegateOnlyException.DEFAULT_MESSAGE)
+
+
+class DelegateOnly(LLM):
+    @override
+    def get_client(self) -> None:
+        raise DelegateOnlyException()
+
+    def __init__(
+        self,
+        model_name: str,
+        provider: str,
+        *,
+        config: LLMConfig | None = None,
+    ):
+        config = config or LLMConfig()
+        config.native = False
+        super().__init__(model_name, provider, config=config)
+
+    @override
+    async def _query_impl(
+        self,
+        input: Sequence[InputItem],
+        *,
+        tools: list[ToolDefinition],
+        **kwargs: object,
+    ) -> QueryResult:
+        assert self.delegate
+
+        return await self.delegate_query(input, tools=tools, **kwargs)
+
+    @override
+    async def parse_input(
+        self,
+        input: Sequence[InputItem],
+        **kwargs: Any,
+    ) -> Any:
+        raise DelegateOnlyException()
+
+    @override
+    async def parse_image(
+        self,
+        image: FileInput,
+    ) -> Any:
+        raise DelegateOnlyException()
+
+    @override
+    async def parse_file(
+        self,
+        file: FileInput,
+    ) -> Any:
+        raise DelegateOnlyException()
+
+    @override
+    async def parse_tools(
+        self,
+        tools: list[ToolDefinition],
+    ) -> Any:
+        raise DelegateOnlyException()
+
+    @override
+    async def upload_file(
+        self,
+        name: str,
+        mime: str,
+        bytes: io.BytesIO,
+        type: Literal["image", "file"] = "file",
+    ) -> FileWithId:
+        raise DelegateOnlyException()

model_library/base/input.py

@@ -0,0 +1,100 @@
+from pprint import pformat
+from typing import Annotated, Any, Literal
+
+from pydantic import BaseModel, Field
+from typing_extensions import override
+
+from model_library.utils import truncate_str
+
+"""
+--- 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

model_library/base/output.py

@@ -0,0 +1,229 @@
+"""
+--- OUTPUT ---
+"""
+
+from pprint import pformat
+from typing import Any, Mapping, Sequence, cast
+
+from pydantic import BaseModel, Field, computed_field, field_validator
+from typing_extensions import override
+
+from model_library.base.input import InputItem, ToolCall
+from model_library.base.utils import (
+    sum_optional,
+)
+from model_library.utils import truncate_str
+
+
+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,
+                ],
+            )
+        )
+
+    @computed_field
+    @property
+    def total_input(self) -> float:
+        return sum(
+            filter(
+                None,
+                [
+                    self.input,
+                    self.cache_read,
+                    self.cache_write,
+                ],
+            )
+        )
+
+    @computed_field
+    @property
+    def total_output(self) -> float:
+        return sum(
+            filter(
+                None,
+                [
+                    self.output,
+                    self.reasoning,
+                ],
+            )
+        )
+
+    @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
+
+    @computed_field
+    @property
+    def total_input_tokens(self) -> int:
+        return sum(
+            filter(
+                None,
+                [
+                    self.in_tokens,
+                    self.cache_read_tokens,
+                    self.cache_write_tokens,
+                ],
+            )
+        )
+
+    @computed_field
+    @property
+    def total_output_tokens(self) -> int:
+        return sum(
+            filter(
+                None,
+                [
+                    self.out_tokens,
+                    self.reasoning_tokens,
+                ],
+            )
+        )
+
+    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

model_library/base/utils.py

@@ -0,0 +1,43 @@
+from typing import Sequence, cast
+
+from model_library.base.input import (
+    FileBase,
+    InputItem,
+    RawInputItem,
+    TextInput,
+    ToolResult,
+)
+from model_library.utils import truncate_str
+
+
+def sum_optional(a: int | None, b: int | None) -> int | None:
+    """Sum two optional integers, returning None if both are None.
+
+    Preserves None to indicate "unknown/not provided" when both inputs are None,
+    otherwise treats None as 0 for summation.
+    """
+    if a is None and b is None:
+        return None
+    return (a or 0) + (b or 0)
+
+
+def get_pretty_input_types(input: Sequence["InputItem"], verbose: bool = False) -> str:
+    # for logging
+    def process_item(item: "InputItem"):
+        match item:
+            case TextInput():
+                item_str = repr(item)
+                return item_str if verbose else truncate_str(item_str)
+            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 ""

model_library/config/ai21labs_models.yaml

@@ -19,6 +19,7 @@ ai21labs-models:
       supports_temperature: true
     default_parameters:
       temperature: 0.4
+      max_output_tokens: 4096
 
   ai21labs/jamba-large-1.7:
     label: Jamba 1.7 Large