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

model_library/providers/openai.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import io
 import json
+import logging
 from typing import Any, Literal, Sequence, cast
 
 from openai import APIConnectionError, AsyncOpenAI
@@ -15,6 +16,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,
@@ -28,6 +30,7 @@ from model_library.base import (
     LLM,
     BatchResult,
     Citation,
+    FileBase,
     FileInput,
     FileWithBase64,
     FileWithId,
@@ -41,7 +44,8 @@ from model_library.base import (
     QueryResultCost,
     QueryResultExtras,
     QueryResultMetadata,
-    RawInputItem,
+    RawInput,
+    RawResponse,
     TextInput,
     ToolBody,
     ToolCall,
@@ -52,6 +56,7 @@ 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
@@ -257,7 +262,9 @@ class OpenAIModel(LLM):
         use_completions: bool = False,
     ):
         super().__init__(model_name, provider, config=config)
-        self.use_completions: bool = use_completions
+        self.use_completions: bool = (
+            use_completions  # TODO: do completions in a separate file
+        )
         self.deep_research = self.provider_config.deep_research
 
         # allow custom client to act as delegate (native)
@@ -269,6 +276,29 @@ class OpenAIModel(LLM):
             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,
@@ -276,63 +306,70 @@ 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:
+                        pass
+                        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
 
@@ -468,19 +505,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(
@@ -505,8 +536,11 @@ class OpenAIModel(LLM):
         if self.reasoning:
             del body["max_tokens"]
             body["max_completion_tokens"] = self.max_tokens
-            if self.reasoning_effort:
-                body["reasoning_effort"] = self.reasoning_effort
+
+        # some model endpoints (like `fireworks/deepseek-v3p2`)
+        # require explicitly setting reasoning effort to disable thinking
+        if self.reasoning_effort is not None:
+            body["reasoning_effort"] = self.reasoning_effort
 
         if self.supports_temperature:
             if self.temperature is not None:
@@ -516,6 +550,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()
@@ -628,7 +679,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,
         )
 
@@ -663,13 +714,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)
 
@@ -701,8 +756,8 @@ class OpenAIModel(LLM):
 
         if self.reasoning:
             body["reasoning"] = {"summary": "auto"}
-            if self.reasoning_effort:
-                body["reasoning"]["effort"] = self.reasoning_effort
+            if self.reasoning_effort is not None:
+                body["reasoning"]["effort"] = self.reasoning_effort  # type: ignore[reportArgumentType]
 
         if self.supports_temperature:
             if self.temperature is not None:
@@ -713,7 +768,6 @@ class OpenAIModel(LLM):
         _ = kwargs.pop("stream", None)
 
         body.update(kwargs)
-
         return body
 
     @override
@@ -722,6 +776,7 @@ class OpenAIModel(LLM):
         input: Sequence[InputItem],
         *,
         tools: list[ToolDefinition],
+        query_logger: logging.Logger,
         **kwargs: object,
     ) -> QueryResult:
         if self.use_completions:
@@ -780,13 +835,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])
@@ -809,7 +863,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:

model_library/providers/vals.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import io
 import json
+import logging
 import random
 import re
 import time
@@ -50,7 +51,7 @@ class DummyAIBatchMixin(LLMBatchMixin):
             "custom_id": custom_id,
             "method": "",
             "url": "",
-            "body": await self._root.create_body(input, tools=[], **kwargs),
+            "body": await self._root.build_body(input, tools=[], **kwargs),
         }
 
     @override
@@ -226,7 +227,8 @@ class DummyAIModel(LLM):
     ) -> FileWithId:
         raise NotImplementedError()
 
-    async def create_body(
+    @override
+    async def build_body(
         self,
         input: Sequence[InputItem],
         *,
@@ -271,9 +273,10 @@ class DummyAIModel(LLM):
         input: Sequence[InputItem],
         *,
         tools: list[ToolDefinition],
+        query_logger: logging.Logger,
         **kwargs: object,
     ) -> QueryResult:
-        body = await self.create_body(input, tools=tools, **kwargs)
+        body = await self.build_body(input, tools=tools, **kwargs)
 
         fail_rate = FAIL_RATE