model-library 0.1.6__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
@@ -16,6 +18,7 @@ from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam
 from openai.types.create_embedding_response import CreateEmbeddingResponse
 from openai.types.moderation_create_response import ModerationCreateResponse
 from openai.types.responses import (
+    ResponseFunctionToolCall,
     ResponseOutputItem,
     ResponseOutputText,
     ResponseStreamEvent,
@@ -29,6 +32,8 @@ from model_library.base import (
     LLM,
     BatchResult,
     Citation,
+    DelegateConfig,
+    FileBase,
     FileInput,
     FileWithBase64,
     FileWithId,
@@ -42,7 +47,9 @@ from model_library.base import (
     QueryResultCost,
     QueryResultExtras,
     QueryResultMetadata,
-    RawInputItem,
+    RateLimit,
+    RawInput,
+    RawResponse,
     TextInput,
     ToolBody,
     ToolCall,
@@ -53,9 +60,11 @@ from model_library.exceptions import (
     ImmediateRetryException,
     MaxOutputTokensExceededError,
     ModelNoOutputError,
+    NoMatchingToolCallError,
 )
 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
 
 
@@ -230,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,
@@ -254,22 +271,48 @@ class OpenAIModel(LLM):
         provider: str = "openai",
         *,
         config: LLMConfig | None = None,
-        custom_client: AsyncOpenAI | None = None,
         use_completions: bool = False,
+        delegate_config: DelegateConfig | None = None,
     ):
+        self.use_completions: bool = (
+            use_completions  # TODO: do completions in a separate file
+        )
+        self.delegate_config = delegate_config
+
         super().__init__(model_name, provider, config=config)
-        self.use_completions: bool = use_completions
-        self.deep_research = self.provider_config.deep_research
 
-        # allow custom client to act as delegate (native)
-        self._delegate_client: AsyncOpenAI | None = custom_client
+        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
         )
 
+    async def get_tool_call_ids(self, input: Sequence[InputItem]) -> list[str]:
+        raw_responses = [x for x in input if isinstance(x, RawResponse)]
+        tool_call_ids: list[str] = []
+
+        if self.use_completions:
+            calls = [
+                y
+                for x in raw_responses
+                if isinstance(x.response, ChatCompletionMessage)
+                and x.response.tool_calls
+                for y in x.response.tool_calls
+            ]
+            tool_call_ids.extend([x.id for x in calls if x.id])
+        else:
+            calls = [
+                y
+                for x in raw_responses
+                for y in x.response
+                if isinstance(y, ResponseFunctionToolCall)
+            ]
+            tool_call_ids.extend([x.id for x in calls if x.id])
+        return tool_call_ids
+
     @override
     async def parse_input(
         self,
@@ -277,63 +320,69 @@ class OpenAIModel(LLM):
         **kwargs: Any,
     ) -> list[dict[str, Any] | Any]:
         new_input: list[dict[str, Any] | Any] = []
+
         content_user: list[dict[str, Any]] = []
+
+        def flush_content_user():
+            if content_user:
+                # NOTE: must make new object as we clear()
+                new_input.append({"role": "user", "content": content_user.copy()})
+                content_user.clear()
+
+        tool_call_ids = await self.get_tool_call_ids(input)
+
         for item in input:
+            if isinstance(item, TextInput):
+                if self.use_completions:
+                    text_key = "text"
+                else:
+                    text_key = "input_text"
+                content_user.append({"type": text_key, "text": item.text})
+                continue
+
+            if isinstance(item, FileBase):
+                match item.type:
+                    case "image":
+                        parsed = await self.parse_image(item)
+                    case "file":
+                        parsed = await self.parse_file(item)
+                content_user.append(parsed)
+                continue
+
+            # non content user item
+            flush_content_user()
+
             match item:
-                case TextInput():
+                case ToolResult():
+                    if item.tool_call.id not in tool_call_ids:
+                        raise NoMatchingToolCallError()
+
                     if self.use_completions:
-                        content_user.append({"type": "text", "text": item.text})
+                        new_input.append(
+                            {
+                                "role": "tool",
+                                "tool_call_id": item.tool_call.id,
+                                "content": item.result,
+                            }
+                        )
                     else:
-                        content_user.append({"type": "input_text", "text": item.text})
-                case FileWithBase64() | FileWithUrl() | FileWithId():
-                    match item.type:
-                        case "image":
-                            content_user.append(await self.parse_image(item))
-                        case "file":
-                            content_user.append(await self.parse_file(item))
-                case _:
-                    if content_user:
-                        new_input.append({"role": "user", "content": content_user})
-                        content_user = []
-                    match item:
-                        case ToolResult():
-                            if not (
-                                not isinstance(x, dict)
-                                and x.type == "function_call"
-                                and x.call_id == item.tool_call.call_id
-                                for x in new_input
-                            ):
-                                raise Exception(
-                                    "Tool call result provided with no matching tool call"
-                                )
-                            if self.use_completions:
-                                new_input.append(
-                                    {
-                                        "role": "tool",
-                                        "tool_call_id": item.tool_call.id,
-                                        "content": item.result,
-                                    }
-                                )
-                            else:
-                                new_input.append(
-                                    {
-                                        "type": "function_call_output",
-                                        "call_id": item.tool_call.call_id,
-                                        "output": item.result,
-                                    }
-                                )
-                        case dict():  # RawInputItem
-                            item = cast(RawInputItem, item)
-                            new_input.append(item)
-                        case _:  # RawResponse
-                            if self.use_completions:
-                                item = cast(ChatCompletionMessageToolCall, item)
-                            else:
-                                item = cast(ResponseOutputItem, item)
-                            new_input.append(item)
-
-        if content_user:
-            new_input.append({"role": "user", "content": content_user})
+                        new_input.append(
+                            {
+                                "type": "function_call_output",
+                                "call_id": item.tool_call.call_id,
+                                "output": item.result,
+                            }
+                        )
+                case RawResponse():
+                    if self.use_completions:
+                        new_input.append(item.response)
+                    else:
+                        new_input.extend(item.response)
+                case RawInput():
+                    new_input.append(item.input)
+
+        # in case content user item is the last item
+        flush_content_user()
 
         return new_input
 
@@ -469,19 +518,13 @@ class OpenAIModel(LLM):
             file_id=response.id,
         )
 
-    async def _query_completions(
+    async def _build_body_completions(
         self,
         input: Sequence[InputItem],
         *,
         tools: list[ToolDefinition],
         **kwargs: object,
-    ) -> QueryResult:
-        """
-        Completions endpoint
-        Generally not used for openai models
-        Used by some providers using openai as a delegate
-        """
-
+    ) -> dict[str, Any]:
         parsed_input: list[dict[str, Any] | ChatCompletionMessage] = []
         if "system_prompt" in kwargs:
             parsed_input.append(
@@ -492,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
 
@@ -520,6 +565,23 @@ class OpenAIModel(LLM):
 
         body.update(kwargs)
 
+        return body
+
+    async def _query_completions(
+        self,
+        input: Sequence[InputItem],
+        *,
+        tools: list[ToolDefinition],
+        **kwargs: object,
+    ) -> QueryResult:
+        """
+        Completions endpoint
+        Generally not used for openai models
+        Used by providers using openai as a delegate
+        """
+
+        body = await self.build_body(input, tools=tools, **kwargs)
+
         output_text: str = ""
         reasoning_text: str = ""
         metadata: QueryResultMetadata = QueryResultMetadata()
@@ -632,7 +694,7 @@ class OpenAIModel(LLM):
             output_text=output_text,
             reasoning=reasoning_text,
             tool_calls=tool_calls,
-            history=[*input, final_message],
+            history=[*input, RawResponse(response=final_message)],
             metadata=metadata,
         )
 
@@ -640,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"
             )
@@ -667,13 +729,17 @@ class OpenAIModel(LLM):
         if not valid:
             raise Exception("Deep research models require web search tools")
 
+    @override
     async def build_body(
         self,
         input: Sequence[InputItem],
         *,
-        tools: Sequence[ToolDefinition],
+        tools: list[ToolDefinition],
         **kwargs: object,
     ) -> dict[str, Any]:
+        if self.use_completions:
+            return await self._build_body_completions(input, tools=tools, **kwargs)
+
         if self.deep_research:
             await self._check_deep_research_args(tools, **kwargs)
 
@@ -694,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:
@@ -708,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
@@ -717,7 +788,6 @@ class OpenAIModel(LLM):
         _ = kwargs.pop("stream", None)
 
         body.update(kwargs)
-
         return body
 
     @override
@@ -785,13 +855,12 @@ class OpenAIModel(LLM):
         citations: list[Citation] = []
         reasoning = None
         for output in response.output:
-            if self.deep_research:
-                if output.type == "message":
-                    for content in output.content:
-                        if not isinstance(content, ResponseOutputText):
-                            continue
-                        for citation in content.annotations:
-                            citations.append(Citation(**citation.model_dump()))
+            if output.type == "message":
+                for content in output.content:
+                    if not isinstance(content, ResponseOutputText):
+                        continue
+                    for citation in content.annotations:
+                        citations.append(Citation(**citation.model_dump()))
 
             if output.type == "reasoning":
                 reasoning = " ".join([i.text for i in output.summary])
@@ -814,7 +883,7 @@ class OpenAIModel(LLM):
             output_text=response.output_text,
             reasoning=reasoning,
             tool_calls=tool_calls,
-            history=[*input, *response.output],
+            history=[*input, RawResponse(response=response.output)],
             extras=QueryResultExtras(citations=citations),
         )
         if response.usage:
@@ -834,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,
@@ -857,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:
@@ -888,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
         )
 
@@ -903,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