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

model_library/providers/mistral.py

@@ -1,10 +1,15 @@
 import io
 import logging
-import time
 from collections.abc import Sequence
 from typing import Any, Literal
 
-from mistralai import AssistantMessage, ContentChunk, Mistral, TextChunk, ThinkChunk
+from mistralai import (
+    AssistantMessage,
+    ContentChunk,
+    Mistral,
+    TextChunk,
+    ThinkChunk,
+)
 from mistralai.models.completionevent import CompletionEvent
 from mistralai.models.toolcall import ToolCall as MistralToolCall
 from mistralai.utils.eventstreaming import EventStreamAsync
@@ -13,14 +18,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,
@@ -39,16 +46,20 @@ from model_library.utils import default_httpx_client
 
 @register_provider("mistralai")
 class MistralModel(LLM):
-    _client: Mistral | None = None
+    @override
+    def _get_default_api_key(self) -> str:
+        return model_library_settings.MISTRAL_API_KEY
 
     @override
-    def get_client(self) -> Mistral:
-        if not MistralModel._client:
-            MistralModel._client = Mistral(
-                api_key=model_library_settings.MISTRAL_API_KEY,
+    def get_client(self, api_key: str | None = None) -> Mistral:
+        if not self.has_client():
+            assert api_key
+            client = Mistral(
+                api_key=api_key,
                 async_client=default_httpx_client(),
             )
-        return MistralModel._client
+            self.assign_client(client)
+        return super().get_client()
 
     def __init__(
         self,
@@ -69,27 +80,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 +112,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 +184,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)
 
@@ -192,12 +208,14 @@ class MistralModel(LLM):
 
         body: dict[str, Any] = {
             "model": self.model_name,
-            "max_tokens": self.max_tokens,
             "messages": messages,
             "prompt_mode": "reasoning" if self.reasoning else None,
             "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
@@ -205,8 +223,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 +275,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 +328,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,