feat: add warning log when targeting a model planned for deprecation (#111)

* add model deprecation warning

* use Optional

Author: GaspardBT

src/mistralai/async_client.py

@@ -1,7 +1,7 @@
 import asyncio
 import posixpath
 from json import JSONDecodeError
-from typing import Any, AsyncGenerator, Dict, List, Optional, Union
+from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Union
 
 from httpx import (
     AsyncClient,
@@ -101,6 +101,7 @@ async def _request(
         stream: bool = False,
         attempt: int = 1,
         data: Optional[Dict[str, Any]] = None,
+        check_model_deprecation_headers_callback: Optional[Callable] = None,
         **kwargs: Any,
     ) -> AsyncGenerator[Dict[str, Any], None]:
         accept_header = "text/event-stream" if stream else "application/json"
@@ -129,6 +130,8 @@ async def _request(
                     data=data,
                     **kwargs,
                 ) as response:
+                    if check_model_deprecation_headers_callback:
+                        check_model_deprecation_headers_callback(response.headers)
                     await self._check_streaming_response(response)
 
                     async for line in response.aiter_lines():
@@ -145,7 +148,8 @@ async def _request(
                     data=data,
                     **kwargs,
                 )
-
+                if check_model_deprecation_headers_callback:
+                    check_model_deprecation_headers_callback(response.headers)
                 yield await self._check_response(response)
 
         except ConnectError as e:
@@ -213,7 +217,12 @@ async def chat(
             response_format=response_format,
         )
 
-        single_response = self._request("post", request, "v1/chat/completions")
+        single_response = self._request(
+            "post",
+            request,
+            "v1/chat/completions",
+            check_model_deprecation_headers_callback=self._check_model_deprecation_header_callback_factory(model),
+        )
 
         async for response in single_response:
             return ChatCompletionResponse(**response)
@@ -267,7 +276,13 @@ async def chat_stream(
             tool_choice=tool_choice,
             response_format=response_format,
         )
-        async_response = self._request("post", request, "v1/chat/completions", stream=True)
+        async_response = self._request(
+            "post",
+            request,
+            "v1/chat/completions",
+            stream=True,
+            check_model_deprecation_headers_callback=self._check_model_deprecation_header_callback_factory(model),
+        )
 
         async for json_response in async_response:
             yield ChatCompletionStreamResponse(**json_response)
@@ -284,7 +299,12 @@ async def embeddings(self, model: str, input: Union[str, List[str]]) -> Embeddin
             EmbeddingResponse: A response object containing the embeddings.
         """
         request = {"model": model, "input": input}
-        single_response = self._request("post", request, "v1/embeddings")
+        single_response = self._request(
+            "post",
+            request,
+            "v1/embeddings",
+            check_model_deprecation_headers_callback=self._check_model_deprecation_header_callback_factory(model),
+        )
 
         async for response in single_response:
             return EmbeddingResponse(**response)
@@ -341,7 +361,12 @@ async def completion(
         request = self._make_completion_request(
             prompt, model, suffix, temperature, max_tokens, top_p, random_seed, stop
         )
-        single_response = self._request("post", request, "v1/fim/completions")
+        single_response = self._request(
+            "post",
+            request,
+            "v1/fim/completions",
+            check_model_deprecation_headers_callback=self._check_model_deprecation_header_callback_factory(model),
+        )
 
         async for response in single_response:
             return ChatCompletionResponse(**response)
@@ -376,9 +401,23 @@ async def completion_stream(
             Dict[str, Any]: a response object containing the generated text.
         """
         request = self._make_completion_request(
-            prompt, model, suffix, temperature, max_tokens, top_p, random_seed, stop, stream=True
+            prompt,
+            model,
+            suffix,
+            temperature,
+            max_tokens,
+            top_p,
+            random_seed,
+            stop,
+            stream=True,
+        )
+        async_response = self._request(
+            "post",
+            request,
+            "v1/fim/completions",
+            stream=True,
+            check_model_deprecation_headers_callback=self._check_model_deprecation_header_callback_factory(model),
         )
-        async_response = self._request("post", request, "v1/fim/completions", stream=True)
 
         async for json_response in async_response:
             yield ChatCompletionStreamResponse(**json_response)

src/mistralai/client.py

@@ -1,7 +1,7 @@
 import posixpath
 import time
 from json import JSONDecodeError
-from typing import Any, Dict, Iterable, Iterator, List, Optional, Union
+from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Union
 
 from httpx import Client, ConnectError, HTTPTransport, RequestError, Response
 
@@ -40,7 +40,9 @@ def __init__(
         super().__init__(endpoint, api_key, max_retries, timeout)
 
         self._client = Client(
-            follow_redirects=True, timeout=self._timeout, transport=HTTPTransport(retries=self._max_retries)
+            follow_redirects=True,
+            timeout=self._timeout,
+            transport=HTTPTransport(retries=self._max_retries),
         )
         self.files = FilesClient(self)
         self.jobs = JobsClient(self)
@@ -94,6 +96,7 @@ def _request(
         stream: bool = False,
         attempt: int = 1,
         data: Optional[Dict[str, Any]] = None,
+        check_model_deprecation_headers_callback: Optional[Callable] = None,
         **kwargs: Any,
     ) -> Iterator[Dict[str, Any]]:
         accept_header = "text/event-stream" if stream else "application/json"
@@ -122,6 +125,8 @@ def _request(
                     data=data,
                     **kwargs,
                 ) as response:
+                    if check_model_deprecation_headers_callback:
+                        check_model_deprecation_headers_callback(response.headers)
                     self._check_streaming_response(response)
 
                     for line in response.iter_lines():
@@ -138,7 +143,8 @@ def _request(
                     data=data,
                     **kwargs,
                 )
-
+                if check_model_deprecation_headers_callback:
+                    check_model_deprecation_headers_callback(response.headers)
                 yield self._check_response(response)
 
         except ConnectError as e:
@@ -207,7 +213,12 @@ def chat(
             response_format=response_format,
         )
 
-        single_response = self._request("post", request, "v1/chat/completions")
+        single_response = self._request(
+            "post",
+            request,
+            "v1/chat/completions",
+            check_model_deprecation_headers_callback=self._check_model_deprecation_header_callback_factory(model),
+        )
 
         for response in single_response:
             return ChatCompletionResponse(**response)
@@ -261,7 +272,13 @@ def chat_stream(
             response_format=response_format,
         )
 
-        response = self._request("post", request, "v1/chat/completions", stream=True)
+        response = self._request(
+            "post",
+            request,
+            "v1/chat/completions",
+            stream=True,
+            check_model_deprecation_headers_callback=self._check_model_deprecation_header_callback_factory(model),
+        )
 
         for json_streamed_response in response:
             yield ChatCompletionStreamResponse(**json_streamed_response)
@@ -278,7 +295,12 @@ def embeddings(self, model: str, input: Union[str, List[str]]) -> EmbeddingRespo
             EmbeddingResponse: A response object containing the embeddings.
         """
         request = {"model": model, "input": input}
-        singleton_response = self._request("post", request, "v1/embeddings")
+        singleton_response = self._request(
+            "post",
+            request,
+            "v1/embeddings",
+            check_model_deprecation_headers_callback=self._check_model_deprecation_header_callback_factory(model),
+        )
 
         for response in singleton_response:
             return EmbeddingResponse(**response)
@@ -337,7 +359,13 @@ def completion(
             prompt, model, suffix, temperature, max_tokens, top_p, random_seed, stop
         )
 
-        single_response = self._request("post", request, "v1/fim/completions", stream=False)
+        single_response = self._request(
+            "post",
+            request,
+            "v1/fim/completions",
+            stream=False,
+            check_model_deprecation_headers_callback=self._check_model_deprecation_header_callback_factory(model),
+        )
 
         for response in single_response:
             return ChatCompletionResponse(**response)
@@ -372,10 +400,24 @@ def completion_stream(
             Iterable[Dict[str, Any]]: a generator that yields response objects containing the generated text.
         """
         request = self._make_completion_request(
-            prompt, model, suffix, temperature, max_tokens, top_p, random_seed, stop, stream=True
+            prompt,
+            model,
+            suffix,
+            temperature,
+            max_tokens,
+            top_p,
+            random_seed,
+            stop,
+            stream=True,
         )
 
-        response = self._request("post", request, "v1/fim/completions", stream=True)
+        response = self._request(
+            "post",
+            request,
+            "v1/fim/completions",
+            stream=True,
+            check_model_deprecation_headers_callback=self._check_model_deprecation_header_callback_factory(model),
+        )
 
         for json_streamed_response in response:
             yield ChatCompletionStreamResponse(**json_streamed_response)

src/mistralai/client_base.py

@@ -1,14 +1,19 @@
 import logging
 import os
 from abc import ABC
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import orjson
-
-from mistralai.exceptions import (
-    MistralException,
+from httpx import Headers
+
+from mistralai.constants import HEADER_MODEL_DEPRECATION_TIMESTAMP
+from mistralai.exceptions import MistralException
+from mistralai.models.chat_completion import (
+    ChatMessage,
+    Function,
+    ResponseFormat,
+    ToolChoice,
 )
-from mistralai.models.chat_completion import ChatMessage, Function, ResponseFormat, ToolChoice
 
 CLIENT_VERSION = "0.4.1"
 
@@ -38,6 +43,14 @@ def __init__(
 
         self._version = CLIENT_VERSION
 
+    def _get_model(self, model: Optional[str] = None) -> str:
+        if model is not None:
+            return model
+        else:
+            if self._default_model is None:
+                raise MistralException(message="model must be provided")
+            return self._default_model
+
     def _parse_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
         parsed_tools: List[Dict[str, Any]] = []
         for tool in tools:
@@ -73,6 +86,22 @@ def _parse_messages(self, messages: List[Any]) -> List[Dict[str, Any]]:
 
         return parsed_messages
 
+    def _check_model_deprecation_header_callback_factory(self, model: Optional[str] = None) -> Callable:
+        model = self._get_model(model)
+
+        def _check_model_deprecation_header_callback(
+            headers: Headers,
+        ) -> None:
+            if HEADER_MODEL_DEPRECATION_TIMESTAMP in headers:
+                self._logger.warning(
+                    f"WARNING: The model {model} is deprecated "
+                    f"and will be removed on {headers[HEADER_MODEL_DEPRECATION_TIMESTAMP]}. "
+                    "Please refer to https://docs.mistral.ai/getting-started/models/#api-versioning "
+                    "for more information."
+                )
+
+        return _check_model_deprecation_header_callback
+
     def _make_completion_request(
         self,
         prompt: str,
@@ -95,16 +124,14 @@ def _make_completion_request(
         if stop is not None:
             request_data["stop"] = stop
 
-        if model is not None:
-            request_data["model"] = model
-        else:
-            if self._default_model is None:
-                raise MistralException(message="model must be provided")
-            request_data["model"] = self._default_model
+        request_data["model"] = self._get_model(model)
 
         request_data.update(
             self._build_sampling_params(
-                temperature=temperature, max_tokens=max_tokens, top_p=top_p, random_seed=random_seed
+                temperature=temperature,
+                max_tokens=max_tokens,
+                top_p=top_p,
+                random_seed=random_seed,
             )
         )
 
@@ -148,16 +175,14 @@ def _make_chat_request(
             "messages": self._parse_messages(messages),
         }
 
-        if model is not None:
-            request_data["model"] = model
-        else:
-            if self._default_model is None:
-                raise MistralException(message="model must be provided")
-            request_data["model"] = self._default_model
+        request_data["model"] = self._get_model(model)
 
         request_data.update(
             self._build_sampling_params(
-                temperature=temperature, max_tokens=max_tokens, top_p=top_p, random_seed=random_seed
+                temperature=temperature,
+                max_tokens=max_tokens,
+                top_p=top_p,
+                random_seed=random_seed,
             )
         )
 

src/mistralai/constants.py

@@ -1,3 +1,5 @@
 RETRY_STATUS_CODES = {429, 500, 502, 503, 504}
 
 ENDPOINT = "https://api.mistral.ai"
+
+HEADER_MODEL_DEPRECATION_TIMESTAMP = "x-model-deprecation-timestamp"