model-library 0.1.7__py3-none-any.whl → 0.1.8__py3-none-any.whl

model_library/providers/openai.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+import datetime
 import io
 import json
 import logging
+import time
 from typing import Any, Literal, Sequence, cast
 
 from openai import APIConnectionError, AsyncOpenAI
@@ -30,6 +32,7 @@ from model_library.base import (
     LLM,
     BatchResult,
     Citation,
+    DelegateConfig,
     FileBase,
     FileInput,
     FileWithBase64,
@@ -44,6 +47,7 @@ from model_library.base import (
     QueryResultCost,
     QueryResultExtras,
     QueryResultMetadata,
+    RateLimit,
     RawInput,
     RawResponse,
     TextInput,
@@ -60,6 +64,7 @@ from model_library.exceptions import (
 )
 from model_library.model_utils import get_reasoning_in_tag
 from model_library.register_models import register_provider
+from model_library.retriers.base import BaseRetrier
 from model_library.utils import create_openai_client_with_defaults
 
 
@@ -234,23 +239,31 @@ class OpenAIBatchMixin(LLMBatchMixin):
 
 class OpenAIConfig(ProviderConfig):
     deep_research: bool = False
+    verbosity: Literal["low", "medium", "high"] | None = None
 
 
 @register_provider("openai")
 class OpenAIModel(LLM):
     provider_config = OpenAIConfig()
 
-    _client: AsyncOpenAI | None = None
+    @override
+    def _get_default_api_key(self) -> str:
+        if self.delegate_config:
+            return self.delegate_config.api_key.get_secret_value()
+        return model_library_settings.OPENAI_API_KEY
 
     @override
-    def get_client(self) -> AsyncOpenAI:
-        if self._delegate_client:
-            return self._delegate_client
-        if not OpenAIModel._client:
-            OpenAIModel._client = create_openai_client_with_defaults(
-                api_key=model_library_settings.OPENAI_API_KEY
+    def get_client(self, api_key: str | None = None) -> AsyncOpenAI:
+        if not self.has_client():
+            assert api_key
+            client = create_openai_client_with_defaults(
+                base_url=self.delegate_config.base_url
+                if self.delegate_config
+                else None,
+                api_key=api_key,
             )
-        return OpenAIModel._client
+            self.assign_client(client)
+        return super().get_client()
 
     def __init__(
         self,
@@ -258,20 +271,21 @@ class OpenAIModel(LLM):
         provider: str = "openai",
         *,
         config: LLMConfig | None = None,
-        custom_client: AsyncOpenAI | None = None,
         use_completions: bool = False,
+        delegate_config: DelegateConfig | None = None,
     ):
-        super().__init__(model_name, provider, config=config)
         self.use_completions: bool = (
             use_completions  # TODO: do completions in a separate file
         )
-        self.deep_research = self.provider_config.deep_research
+        self.delegate_config = delegate_config
 
-        # allow custom client to act as delegate (native)
-        self._delegate_client: AsyncOpenAI | None = custom_client
+        super().__init__(model_name, provider, config=config)
+
+        self.deep_research = self.provider_config.deep_research
+        self.verbosity = self.provider_config.verbosity
 
         # batch client
-        self.supports_batch: bool = self.supports_batch and not custom_client
+        self.supports_batch: bool = self.supports_batch and not self.delegate_config
         self.batch: LLMBatchMixin | None = (
             OpenAIBatchMixin(self) if self.supports_batch else None
         )
@@ -361,7 +375,6 @@ class OpenAIModel(LLM):
                         )
                 case RawResponse():
                     if self.use_completions:
-                        pass
                         new_input.append(item.response)
                     else:
                         new_input.extend(item.response)
@@ -522,18 +535,20 @@ class OpenAIModel(LLM):
 
         body: dict[str, Any] = {
             "model": self.model_name,
-            "max_tokens": self.max_tokens,
             "messages": parsed_input,
             # enable usage data in streaming responses
             "stream_options": {"include_usage": True},
         }
 
+        if self.max_tokens:
+            body["max_tokens"] = self.max_tokens
+
         if self.supports_tools:
             parsed_tools = await self.parse_tools(tools)
             if parsed_tools:
                 body["tools"] = parsed_tools
 
-        if self.reasoning:
+        if self.reasoning and self.max_tokens:
             del body["max_tokens"]
             body["max_completion_tokens"] = self.max_tokens
 
@@ -687,7 +702,7 @@ class OpenAIModel(LLM):
         self, tools: Sequence[ToolDefinition], **kwargs: object
     ) -> None:
         min_tokens = 30_000
-        if self.max_tokens < min_tokens:
+        if not self.max_tokens or self.max_tokens < min_tokens:
             self.logger.warning(
                 f"Recommended to set max_tokens >= {min_tokens} for deep research models"
             )
@@ -745,10 +760,12 @@ class OpenAIModel(LLM):
 
         body: dict[str, Any] = {
             "model": self.model_name,
-            "max_output_tokens": self.max_tokens,
             "input": parsed_input,
         }
 
+        if self.max_tokens:
+            body["max_output_tokens"] = self.max_tokens
+
         if parsed_tools:
             body["tools"] = parsed_tools
         else:
@@ -759,6 +776,9 @@ class OpenAIModel(LLM):
             if self.reasoning_effort is not None:
                 body["reasoning"]["effort"] = self.reasoning_effort  # type: ignore[reportArgumentType]
 
+        if self.verbosity is not None:
+            body["text"] = {"format": {"type": "text"}, "verbosity": self.verbosity}
+
         if self.supports_temperature:
             if self.temperature is not None:
                 body["temperature"] = self.temperature
@@ -883,6 +903,61 @@ class OpenAIModel(LLM):
 
         return result
 
+    @override
+    async def get_rate_limit(self) -> RateLimit | None:
+        headers = {}
+
+        try:
+            # NOTE: with_streaming_response doesn't seem to always work
+            if self.use_completions:
+                response = (
+                    await self.get_client().chat.completions.with_raw_response.create(
+                        max_completion_tokens=16,
+                        model=self.model_name,
+                        messages=[
+                            {
+                                "role": "user",
+                                "content": "Ping",
+                            }
+                        ],
+                        stream=True,
+                    )
+                )
+            else:
+                response = await self.get_client().responses.with_raw_response.create(
+                    max_output_tokens=16,
+                    input="Ping",
+                    model=self.model_name,
+                )
+            headers = response.headers
+
+            server_time_str = headers.get("date")
+            if server_time_str:
+                server_time = datetime.datetime.strptime(
+                    server_time_str, "%a, %d %b %Y %H:%M:%S GMT"
+                ).replace(tzinfo=datetime.timezone.utc)
+                timestamp = server_time.timestamp()
+            else:
+                timestamp = time.time()
+
+            # NOTE: for openai, max_tokens is used to reject requests if the amount of tokens left is less than the max_tokens
+
+            # we calculate estimated_tokens as (character_count / 4) + max_tokens. Note that OpenAI's rate limiter doesn't tokenize the request using the model's specific tokenizer but relies on a character count-based heuristic.
+
+            return RateLimit(
+                raw=headers,
+                unix_timestamp=timestamp,
+                request_limit=headers.get("x-ratelimit-limit-requests", None)
+                or headers.get("x-ratelimit-limit", None),
+                request_remaining=headers.get("x-ratelimit-remaining-requests", None)
+                or headers.get("x-ratelimit-remaining"),
+                token_limit=int(headers["x-ratelimit-limit-tokens"]),
+                token_remaining=int(headers["x-ratelimit-remaining-tokens"]),
+            )
+        except Exception as e:
+            self.logger.warning(f"Failed to get rate limit: {e}")
+            return None
+
     @override
     async def query_json(
         self,
@@ -906,7 +981,9 @@ class OpenAIModel(LLM):
             except APIConnectionError:
                 raise ImmediateRetryException("Failed to connect to OpenAI")
 
-        response = await LLM.immediate_retry_wrapper(func=_query, logger=self.logger)
+        response = await BaseRetrier.immediate_retry_wrapper(
+            func=_query, logger=self.logger
+        )
 
         parsed: PydanticT | None = response.output_parsed
         if parsed is None:
@@ -937,7 +1014,7 @@ class OpenAIModel(LLM):
 
             return response.data[0].embedding
 
-        return await LLM.immediate_retry_wrapper(
+        return await BaseRetrier.immediate_retry_wrapper(
             func=_get_embedding, logger=self.logger
         )
 
@@ -952,7 +1029,7 @@ class OpenAIModel(LLM):
             except Exception as e:
                 raise Exception("Failed to query OpenAI's Moderation endpoint") from e
 
-        return await LLM.immediate_retry_wrapper(
+        return await BaseRetrier.immediate_retry_wrapper(
             func=_moderate_content, logger=self.logger
         )
 

model_library/providers/openrouter.py

@@ -0,0 +1,34 @@
+from typing import Literal
+
+from pydantic import SecretStr
+
+from model_library import model_library_settings
+from model_library.base import (
+    DelegateConfig,
+    DelegateOnly,
+    LLMConfig,
+)
+from model_library.register_models import register_provider
+
+
+@register_provider("openrouter")
+class OpenRouterModel(DelegateOnly):
+    def __init__(
+        self,
+        model_name: str,
+        provider: Literal["openrouter"] = "openrouter",
+        *,
+        config: LLMConfig | None = None,
+    ):
+        super().__init__(model_name, provider, config=config)
+
+        # https://openrouter.ai/docs/guides/community/openai-sdk
+        self.init_delegate(
+            config=config,
+            delegate_config=DelegateConfig(
+                base_url="https://openrouter.ai/api/v1",
+                api_key=SecretStr(model_library_settings.OPENROUTER_API_KEY),
+            ),
+            use_completions=True,
+            delegate_provider="openai",
+        )

model_library/providers/perplexity.py

@@ -1,13 +1,14 @@
 from typing import Literal
 
+from pydantic import SecretStr
+
 from model_library import model_library_settings
 from model_library.base import (
+    DelegateConfig,
     DelegateOnly,
     LLMConfig,
 )
-from model_library.providers.openai import OpenAIModel
 from model_library.register_models import register_provider
-from model_library.utils import create_openai_client_with_defaults
 
 
 @register_provider("perplexity")
@@ -22,13 +23,12 @@ class PerplexityModel(DelegateOnly):
         super().__init__(model_name, provider, config=config)
 
         # https://docs.perplexity.ai/guides/chat-completions-guide
-        self.delegate = OpenAIModel(
-            model_name=self.model_name,
-            provider=self.provider,
+        self.init_delegate(
             config=config,
-            custom_client=create_openai_client_with_defaults(
-                api_key=model_library_settings.PERPLEXITY_API_KEY,
+            delegate_config=DelegateConfig(
                 base_url="https://api.perplexity.ai",
+                api_key=SecretStr(model_library_settings.PERPLEXITY_API_KEY),
             ),
             use_completions=True,
+            delegate_provider="openai",
         )

model_library/providers/together.py

@@ -1,18 +1,18 @@
 from typing import Literal
 
+from pydantic import SecretStr
 from typing_extensions import override
 
 from model_library import model_library_settings
 from model_library.base import (
+    DelegateConfig,
     DelegateOnly,
     LLMConfig,
     ProviderConfig,
     QueryResultCost,
     QueryResultMetadata,
 )
-from model_library.providers.openai import OpenAIModel
 from model_library.register_models import register_provider
-from model_library.utils import create_openai_client_with_defaults
 
 
 class TogetherConfig(ProviderConfig):
@@ -32,15 +32,14 @@ class TogetherModel(DelegateOnly):
     ):
         super().__init__(model_name, provider, config=config)
         # https://docs.together.ai/docs/openai-api-compatibility
-        self.delegate = OpenAIModel(
-            model_name=self.model_name,
-            provider=self.provider,
+        self.init_delegate(
             config=config,
-            custom_client=create_openai_client_with_defaults(
-                api_key=model_library_settings.TOGETHER_API_KEY,
-                base_url="https://api.together.xyz/v1",
+            delegate_config=DelegateConfig(
+                base_url="https://api.together.xyz/v1/",
+                api_key=SecretStr(model_library_settings.TOGETHER_API_KEY),
             ),
             use_completions=True,
+            delegate_provider="openai",
         )
 
     @override

model_library/providers/vals.py

@@ -151,13 +151,17 @@ class DummyAIBatchMixin(LLMBatchMixin):
 class DummyAIModel(LLM):
     _client: Redis | None = None
 
-    @override
-    def get_client(self) -> Redis:
-        if not DummyAIModel._client:
-            DummyAIModel._client = redis.from_url(  # pyright: ignore[reportUnknownMemberType]
+    def _get_default_api_key(self) -> str:
+        return model_library_settings.REDIS_URL
+
+    def get_client(self, api_key: str | None = None) -> Redis:
+        if not self.has_client():
+            assert api_key
+            client = redis.from_url(  # pyright: ignore[reportUnknownMemberType]
                 model_library_settings.REDIS_URL, decode_responses=True
             )
-        return DummyAIModel._client
+            self.assign_client(client)
+        return super().get_client()
 
     def __init__(
         self,
@@ -238,12 +242,14 @@ class DummyAIModel(LLM):
         messages = await self.parse_input(input)
         body: dict[str, Any] = {
             "model": self.model_name,
-            "max_tokens": self.max_tokens,
             "seed": 0,
             "messages": messages,
             "tools": await self.parse_tools(tools),
         }
 
+        if self.max_tokens:
+            body["max_tokens"] = self.max_tokens
+
         if self.supports_temperature:
             if self.temperature is not None:
                 body["temperature"] = self.temperature

model_library/providers/xai.py

@@ -2,7 +2,7 @@ import io
 import logging
 from typing import Any, Literal, Sequence
 
-import grpc
+from pydantic import SecretStr
 from typing_extensions import override
 from xai_sdk import AsyncClient
 from xai_sdk.aio.chat import Chat
@@ -14,6 +14,7 @@ from xai_sdk.proto.v6.chat_pb2 import Message, Tool
 from model_library import model_library_settings
 from model_library.base import (
     LLM,
+    DelegateConfig,
     FileBase,
     FileInput,
     FileWithBase64,
@@ -36,24 +37,26 @@ from model_library.exceptions import (
     MaxOutputTokensExceededError,
     ModelNoOutputError,
     NoMatchingToolCallError,
-    RateLimitException,
 )
 from model_library.providers.openai import OpenAIModel
 from model_library.register_models import register_provider
-from model_library.utils import create_openai_client_with_defaults
 
 
 @register_provider("grok")
 class XAIModel(LLM):
-    _client: AsyncClient | None = None
+    @override
+    def _get_default_api_key(self) -> str:
+        return model_library_settings.XAI_API_KEY
 
     @override
-    def get_client(self) -> AsyncClient:
-        if not XAIModel._client:
-            XAIModel._client = AsyncClient(
-                api_key=model_library_settings.XAI_API_KEY,
+    def get_client(self, api_key: str | None = None) -> AsyncClient:
+        if not self.has_client():
+            assert api_key
+            client = AsyncClient(
+                api_key=api_key,
             )
-        return XAIModel._client
+            self.assign_client(client)
+        return super().get_client()
 
     @override
     def __init__(
@@ -73,13 +76,13 @@ class XAIModel(LLM):
                 model_name=self.model_name,
                 provider=provider,
                 config=config,
-                custom_client=create_openai_client_with_defaults(
-                    api_key=model_library_settings.XAI_API_KEY,
+                delegate_config=DelegateConfig(
                     base_url=(
                         "https://us-west-1.api.x.ai/v1"
                         if "grok-3-mini-reasoning" in self.model_name
                         else "https://api.x.ai/v1"
                     ),
+                    api_key=SecretStr(model_library_settings.XAI_API_KEY),
                 ),
                 use_completions=True,
             )
@@ -210,12 +213,14 @@ class XAIModel(LLM):
             messages.append(system(str(kwargs.pop("system_prompt"))))
 
         body: dict[str, Any] = {
-            "max_tokens": self.max_tokens,
             "model": self.model_name,
             "tools": await self.parse_tools(tools),
             "messages": messages,
         }
 
+        if self.max_tokens:
+            body["max_tokens"] = self.max_tokens
+
         if self.supports_temperature:
             if self.temperature is not None:
                 body["temperature"] = self.temperature
@@ -253,38 +258,35 @@ class XAIModel(LLM):
 
         body = await self.build_body(input, tools=tools, **kwargs)
 
-        try:
-            chat: Chat = self.get_client().chat.create(**body)
-
-            latest_response: Response | None = None
-            async for response, _ in chat.stream():
-                latest_response = response
-
-            if not latest_response:
-                raise ModelNoOutputError("Model failed to produce a response")
-
-            tool_calls: list[ToolCall] = []
-            if (
-                latest_response.finish_reason == "REASON_TOOL_CALLS"
-                and latest_response.tool_calls
-            ):
-                for tool_call in latest_response.tool_calls:
-                    tool_calls.append(
-                        ToolCall(
-                            id=tool_call.id,
-                            name=tool_call.function.name,
-                            args=tool_call.function.arguments,
-                        )
+        chat: Chat = self.get_client().chat.create(**body)
+
+        latest_response: Response | None = None
+        async for response, _ in chat.stream():
+            latest_response = response
+
+        if not latest_response:
+            raise ModelNoOutputError("Model failed to produce a response")
+
+        tool_calls: list[ToolCall] = []
+        if (
+            latest_response.finish_reason == "REASON_TOOL_CALLS"
+            and latest_response.tool_calls
+        ):
+            for tool_call in latest_response.tool_calls:
+                tool_calls.append(
+                    ToolCall(
+                        id=tool_call.id,
+                        name=tool_call.function.name,
+                        args=tool_call.function.arguments,
                     )
+                )
 
-            if (
-                latest_response.finish_reason == "REASON_MAX_LEN"
-                and not latest_response.content
-                and not latest_response.reasoning_content
-            ):
-                raise MaxOutputTokensExceededError()
-        except grpc.RpcError as e:
-            raise RateLimitException(e.details())
+        if (
+            latest_response.finish_reason == "REASON_MAX_LEN"
+            and not latest_response.content
+            and not latest_response.reasoning_content
+        ):
+            raise MaxOutputTokensExceededError()
 
         return QueryResult(
             output_text=latest_response.content,
@@ -310,6 +312,9 @@ class XAIModel(LLM):
         tools: list[ToolDefinition] = [],
         **kwargs: object,
     ) -> int:
+        if not input and not history:
+            return 0
+
         string_input = await self.stringify_input(input, history=history, tools=tools)
         self.logger.debug(string_input)
 

model_library/providers/zai.py

@@ -1,17 +1,36 @@
-from typing import Literal
+from typing import Any, Literal
+
+from pydantic import SecretStr
+from typing_extensions import override
 
 from model_library import model_library_settings
 from model_library.base import (
+    DelegateConfig,
     DelegateOnly,
     LLMConfig,
+    ProviderConfig,
 )
-from model_library.providers.openai import OpenAIModel
 from model_library.register_models import register_provider
-from model_library.utils import create_openai_client_with_defaults
+
+
+class ZAIConfig(ProviderConfig):
+    """Configuration for ZAI (GLM) models.
+
+    Attributes:
+        clear_thinking: When disabled, reasoning content from previous turns is
+            preserved in context. This is useful for multi-turn conversations where
+            you want the model to maintain coherent reasoning across turns.
+            Enabled by default on the standard API endpoint.
+            See: https://docs.z.ai/guides/capabilities/thinking-mode
+    """
+
+    clear_thinking: bool = True
 
 
 @register_provider("zai")
 class ZAIModel(DelegateOnly):
+    provider_config = ZAIConfig()
+
     def __init__(
         self,
         model_name: str,
@@ -21,14 +40,25 @@ class ZAIModel(DelegateOnly):
     ):
         super().__init__(model_name, provider, config=config)
 
+        self.clear_thinking = self.provider_config.clear_thinking
+
         # https://docs.z.ai/guides/develop/openai/python
-        self.delegate = OpenAIModel(
-            model_name=self.model_name,
-            provider=self.provider,
+        self.init_delegate(
             config=config,
-            custom_client=create_openai_client_with_defaults(
-                api_key=model_library_settings.ZAI_API_KEY,
+            delegate_config=DelegateConfig(
                 base_url="https://open.bigmodel.cn/api/paas/v4/",
+                api_key=SecretStr(model_library_settings.ZAI_API_KEY),
             ),
             use_completions=True,
+            delegate_provider="openai",
         )
+
+    @override
+    def _get_extra_body(self) -> dict[str, Any]:
+        """Build extra body parameters for GLM-specific features."""
+        return {
+            "thinking": {
+                "type": "enabled" if self.reasoning else "disabled",
+                "clear_thinking": self.clear_thinking,
+            }
+        }