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

model_library/providers/mistral.py

@@ -1,6 +1,5 @@
 import io
 import logging
-import time
 from collections.abc import Sequence
 from typing import Any, Literal
 
@@ -13,14 +12,16 @@ from typing_extensions import override
 from model_library import model_library_settings
 from model_library.base import (
     LLM,
+    FileBase,
     FileInput,
     FileWithBase64,
     FileWithId,
-    FileWithUrl,
     InputItem,
     LLMConfig,
     QueryResult,
     QueryResultMetadata,
+    RawInput,
+    RawResponse,
     TextInput,
     ToolBody,
     ToolCall,
@@ -69,27 +70,30 @@ class MistralModel(LLM):
         content_user: list[dict[str, Any]] = []
 
         def flush_content_user():
-            nonlocal content_user
-
             if content_user:
-                new_input.append({"role": "user", "content": content_user})
-                content_user = []
+                # NOTE: must make new object as we clear()
+                new_input.append({"role": "user", "content": content_user.copy()})
+                content_user.clear()
 
         for item in input:
+            if isinstance(item, TextInput):
+                content_user.append({"type": "text", "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():
-                    content_user.append({"type": "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 AssistantMessage():
-                    flush_content_user()
-                    new_input.append(item)
                 case ToolResult():
-                    flush_content_user()
                     new_input.append(
                         {
                             "role": "tool",
@@ -98,9 +102,12 @@ class MistralModel(LLM):
                             "tool_call_id": item.tool_call.id,
                         }
                     )
-                case _:
-                    raise BadInputError("Unsupported input type")
+                case RawResponse():
+                    new_input.append(item.response)
+                case RawInput():
+                    new_input.append(item.input)
 
+        # in case content user item is the last item
         flush_content_user()
 
         return new_input
@@ -167,14 +174,13 @@ class MistralModel(LLM):
         raise NotImplementedError()
 
     @override
-    async def _query_impl(
+    async def build_body(
         self,
         input: Sequence[InputItem],
         *,
         tools: list[ToolDefinition],
-        query_logger: logging.Logger,
         **kwargs: object,
-    ) -> QueryResult:
+    ) -> dict[str, Any]:
         # mistral supports max 8 images, merge extra images into the 8th image
         input = trim_images(input, max_images=8)
 
@@ -205,8 +211,18 @@ class MistralModel(LLM):
                 body["top_p"] = self.top_p
 
         body.update(kwargs)
+        return body
 
-        start = time.time()
+    @override
+    async def _query_impl(
+        self,
+        input: Sequence[InputItem],
+        *,
+        tools: list[ToolDefinition],
+        query_logger: logging.Logger,
+        **kwargs: object,
+    ) -> QueryResult:
+        body = await self.build_body(input, tools=tools, **kwargs)
 
         response: EventStreamAsync[
             CompletionEvent
@@ -247,8 +263,6 @@ class MistralModel(LLM):
                     in_tokens += data.usage.prompt_tokens or 0
                     out_tokens += data.usage.completion_tokens or 0
 
-            self.logger.info(f"Finished in: {time.time() - start}")
-
         except Exception as e:
             self.logger.error(f"Error: {e}", exc_info=True)
             raise e
@@ -302,7 +316,7 @@ class MistralModel(LLM):
         return QueryResult(
             output_text=text,
             reasoning=reasoning or None,
-            history=[*input, message],
+            history=[*input, RawResponse(response=message)],
             tool_calls=tool_calls,
             metadata=QueryResultMetadata(
                 in_tokens=in_tokens,

model_library/providers/openai.py

@@ -16,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,
@@ -29,6 +30,7 @@ from model_library.base import (
     LLM,
     BatchResult,
     Citation,
+    FileBase,
     FileInput,
     FileWithBase64,
     FileWithId,
@@ -42,7 +44,8 @@ from model_library.base import (
     QueryResultCost,
     QueryResultExtras,
     QueryResultMetadata,
-    RawInputItem,
+    RawInput,
+    RawResponse,
     TextInput,
     ToolBody,
     ToolCall,
@@ -53,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
@@ -258,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)
@@ -270,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,
@@ -277,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
 
@@ -469,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(
@@ -520,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()
@@ -632,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,
         )
 
@@ -667,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)
 
@@ -717,7 +768,6 @@ class OpenAIModel(LLM):
         _ = kwargs.pop("stream", None)
 
         body.update(kwargs)
-
         return body
 
     @override
@@ -785,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])
@@ -814,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

@@ -51,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
@@ -227,7 +227,8 @@ class DummyAIModel(LLM):
     ) -> FileWithId:
         raise NotImplementedError()
 
-    async def create_body(
+    @override
+    async def build_body(
         self,
         input: Sequence[InputItem],
         *,
@@ -275,7 +276,7 @@ class DummyAIModel(LLM):
         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