mistralai 1.5.1__py3-none-any.whl → 1.5.2rc1__py3-none-any.whl

mistralai_gcp/chat.py

@@ -3,7 +3,7 @@
 from .basesdk import BaseSDK
 from mistralai_gcp import models, utils
 from mistralai_gcp._hooks import HookContext
-from mistralai_gcp.types import Nullable, OptionalNullable, UNSET
+from mistralai_gcp.types import OptionalNullable, UNSET
 from mistralai_gcp.utils import eventstreaming
 from typing import Any, List, Mapping, Optional, Union
 
@@ -14,7 +14,7 @@ class Chat(BaseSDK):
     def stream(
         self,
         *,
-        model: Nullable[str],
+        model: str,
         messages: Union[List[models.Messages], List[models.MessagesTypedDict]],
         temperature: OptionalNullable[float] = UNSET,
         top_p: Optional[float] = None,
@@ -37,6 +37,9 @@ class Chat(BaseSDK):
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
         n: OptionalNullable[int] = UNSET,
+        prediction: Optional[
+            Union[models.Prediction, models.PredictionTypedDict]
+        ] = None,
         retries: OptionalNullable[utils.RetryConfig] = UNSET,
         server_url: Optional[str] = None,
         timeout_ms: Optional[int] = None,
@@ -60,6 +63,7 @@ class Chat(BaseSDK):
         :param presence_penalty: presence_penalty determines how much the model penalizes the repetition of words or phrases. A higher presence penalty encourages the model to use a wider variety of words and phrases, making the output more diverse and creative.
         :param frequency_penalty: frequency_penalty penalizes the repetition of words based on their frequency in the generated text. A higher frequency penalty discourages the model from repeating words that have already appeared frequently in the output, promoting diversity and reducing repetition.
         :param n: Number of completions to return for each request, input tokens are only billed once.
+        :param prediction:
         :param retries: Override the default retry configuration for this method
         :param server_url: Override the default server URL for this method
         :param timeout_ms: Override the default request timeout configuration for this method in milliseconds
@@ -72,6 +76,8 @@ class Chat(BaseSDK):
 
         if server_url is not None:
             base_url = server_url
+        else:
+            base_url = self._get_url(base_url, url_variables)
 
         request = models.ChatCompletionStreamRequest(
             model=model,
@@ -92,6 +98,9 @@ class Chat(BaseSDK):
             presence_penalty=presence_penalty,
             frequency_penalty=frequency_penalty,
             n=n,
+            prediction=utils.get_pydantic_model(
+                prediction, Optional[models.Prediction]
+            ),
         )
 
         req = self._build_request(
@@ -123,6 +132,7 @@ class Chat(BaseSDK):
 
         http_res = self.do_request(
             hook_ctx=HookContext(
+                base_url=base_url or "",
                 operation_id="stream_chat",
                 oauth2_scopes=[],
                 security_source=self.sdk_configuration.security,
@@ -133,7 +143,7 @@ class Chat(BaseSDK):
             retry_config=retry_config,
         )
 
-        data: Any = None
+        response_data: Any = None
         if utils.match_response(http_res, "200", "text/event-stream"):
             return eventstreaming.EventStream(
                 http_res,
@@ -142,9 +152,16 @@ class Chat(BaseSDK):
             )
         if utils.match_response(http_res, "422", "application/json"):
             http_res_text = utils.stream_to_text(http_res)
-            data = utils.unmarshal_json(http_res_text, models.HTTPValidationErrorData)
-            raise models.HTTPValidationError(data=data)
-        if utils.match_response(http_res, ["4XX", "5XX"], "*"):
+            response_data = utils.unmarshal_json(
+                http_res_text, models.HTTPValidationErrorData
+            )
+            raise models.HTTPValidationError(data=response_data)
+        if utils.match_response(http_res, "4XX", "*"):
+            http_res_text = utils.stream_to_text(http_res)
+            raise models.SDKError(
+                "API error occurred", http_res.status_code, http_res_text, http_res
+            )
+        if utils.match_response(http_res, "5XX", "*"):
             http_res_text = utils.stream_to_text(http_res)
             raise models.SDKError(
                 "API error occurred", http_res.status_code, http_res_text, http_res
@@ -162,7 +179,7 @@ class Chat(BaseSDK):
     async def stream_async(
         self,
         *,
-        model: Nullable[str],
+        model: str,
         messages: Union[List[models.Messages], List[models.MessagesTypedDict]],
         temperature: OptionalNullable[float] = UNSET,
         top_p: Optional[float] = None,
@@ -185,6 +202,9 @@ class Chat(BaseSDK):
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
         n: OptionalNullable[int] = UNSET,
+        prediction: Optional[
+            Union[models.Prediction, models.PredictionTypedDict]
+        ] = None,
         retries: OptionalNullable[utils.RetryConfig] = UNSET,
         server_url: Optional[str] = None,
         timeout_ms: Optional[int] = None,
@@ -208,6 +228,7 @@ class Chat(BaseSDK):
         :param presence_penalty: presence_penalty determines how much the model penalizes the repetition of words or phrases. A higher presence penalty encourages the model to use a wider variety of words and phrases, making the output more diverse and creative.
         :param frequency_penalty: frequency_penalty penalizes the repetition of words based on their frequency in the generated text. A higher frequency penalty discourages the model from repeating words that have already appeared frequently in the output, promoting diversity and reducing repetition.
         :param n: Number of completions to return for each request, input tokens are only billed once.
+        :param prediction:
         :param retries: Override the default retry configuration for this method
         :param server_url: Override the default server URL for this method
         :param timeout_ms: Override the default request timeout configuration for this method in milliseconds
@@ -220,6 +241,8 @@ class Chat(BaseSDK):
 
         if server_url is not None:
             base_url = server_url
+        else:
+            base_url = self._get_url(base_url, url_variables)
 
         request = models.ChatCompletionStreamRequest(
             model=model,
@@ -240,6 +263,9 @@ class Chat(BaseSDK):
             presence_penalty=presence_penalty,
             frequency_penalty=frequency_penalty,
             n=n,
+            prediction=utils.get_pydantic_model(
+                prediction, Optional[models.Prediction]
+            ),
         )
 
         req = self._build_request_async(
@@ -271,6 +297,7 @@ class Chat(BaseSDK):
 
         http_res = await self.do_request_async(
             hook_ctx=HookContext(
+                base_url=base_url or "",
                 operation_id="stream_chat",
                 oauth2_scopes=[],
                 security_source=self.sdk_configuration.security,
@@ -281,7 +308,7 @@ class Chat(BaseSDK):
             retry_config=retry_config,
         )
 
-        data: Any = None
+        response_data: Any = None
         if utils.match_response(http_res, "200", "text/event-stream"):
             return eventstreaming.EventStreamAsync(
                 http_res,
@@ -290,9 +317,16 @@ class Chat(BaseSDK):
             )
         if utils.match_response(http_res, "422", "application/json"):
             http_res_text = await utils.stream_to_text_async(http_res)
-            data = utils.unmarshal_json(http_res_text, models.HTTPValidationErrorData)
-            raise models.HTTPValidationError(data=data)
-        if utils.match_response(http_res, ["4XX", "5XX"], "*"):
+            response_data = utils.unmarshal_json(
+                http_res_text, models.HTTPValidationErrorData
+            )
+            raise models.HTTPValidationError(data=response_data)
+        if utils.match_response(http_res, "4XX", "*"):
+            http_res_text = await utils.stream_to_text_async(http_res)
+            raise models.SDKError(
+                "API error occurred", http_res.status_code, http_res_text, http_res
+            )
+        if utils.match_response(http_res, "5XX", "*"):
             http_res_text = await utils.stream_to_text_async(http_res)
             raise models.SDKError(
                 "API error occurred", http_res.status_code, http_res_text, http_res
@@ -310,7 +344,7 @@ class Chat(BaseSDK):
     def complete(
         self,
         *,
-        model: Nullable[str],
+        model: str,
         messages: Union[
             List[models.ChatCompletionRequestMessages],
             List[models.ChatCompletionRequestMessagesTypedDict],
@@ -341,6 +375,9 @@ class Chat(BaseSDK):
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
         n: OptionalNullable[int] = UNSET,
+        prediction: Optional[
+            Union[models.Prediction, models.PredictionTypedDict]
+        ] = None,
         retries: OptionalNullable[utils.RetryConfig] = UNSET,
         server_url: Optional[str] = None,
         timeout_ms: Optional[int] = None,
@@ -362,6 +399,7 @@ class Chat(BaseSDK):
         :param presence_penalty: presence_penalty determines how much the model penalizes the repetition of words or phrases. A higher presence penalty encourages the model to use a wider variety of words and phrases, making the output more diverse and creative.
         :param frequency_penalty: frequency_penalty penalizes the repetition of words based on their frequency in the generated text. A higher frequency penalty discourages the model from repeating words that have already appeared frequently in the output, promoting diversity and reducing repetition.
         :param n: Number of completions to return for each request, input tokens are only billed once.
+        :param prediction:
         :param retries: Override the default retry configuration for this method
         :param server_url: Override the default server URL for this method
         :param timeout_ms: Override the default request timeout configuration for this method in milliseconds
@@ -374,6 +412,8 @@ class Chat(BaseSDK):
 
         if server_url is not None:
             base_url = server_url
+        else:
+            base_url = self._get_url(base_url, url_variables)
 
         request = models.ChatCompletionRequest(
             model=model,
@@ -396,6 +436,9 @@ class Chat(BaseSDK):
             presence_penalty=presence_penalty,
             frequency_penalty=frequency_penalty,
             n=n,
+            prediction=utils.get_pydantic_model(
+                prediction, Optional[models.Prediction]
+            ),
         )
 
         req = self._build_request(
@@ -427,6 +470,7 @@ class Chat(BaseSDK):
 
         http_res = self.do_request(
             hook_ctx=HookContext(
+                base_url=base_url or "",
                 operation_id="chat_completion_v1_chat_completions_post",
                 oauth2_scopes=[],
                 security_source=self.sdk_configuration.security,
@@ -436,15 +480,22 @@ class Chat(BaseSDK):
             retry_config=retry_config,
         )
 
-        data: Any = None
+        response_data: Any = None
         if utils.match_response(http_res, "200", "application/json"):
             return utils.unmarshal_json(
                 http_res.text, Optional[models.ChatCompletionResponse]
             )
         if utils.match_response(http_res, "422", "application/json"):
-            data = utils.unmarshal_json(http_res.text, models.HTTPValidationErrorData)
-            raise models.HTTPValidationError(data=data)
-        if utils.match_response(http_res, ["4XX", "5XX"], "*"):
+            response_data = utils.unmarshal_json(
+                http_res.text, models.HTTPValidationErrorData
+            )
+            raise models.HTTPValidationError(data=response_data)
+        if utils.match_response(http_res, "4XX", "*"):
+            http_res_text = utils.stream_to_text(http_res)
+            raise models.SDKError(
+                "API error occurred", http_res.status_code, http_res_text, http_res
+            )
+        if utils.match_response(http_res, "5XX", "*"):
             http_res_text = utils.stream_to_text(http_res)
             raise models.SDKError(
                 "API error occurred", http_res.status_code, http_res_text, http_res
@@ -462,7 +513,7 @@ class Chat(BaseSDK):
     async def complete_async(
         self,
         *,
-        model: Nullable[str],
+        model: str,
         messages: Union[
             List[models.ChatCompletionRequestMessages],
             List[models.ChatCompletionRequestMessagesTypedDict],
@@ -493,6 +544,9 @@ class Chat(BaseSDK):
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
         n: OptionalNullable[int] = UNSET,
+        prediction: Optional[
+            Union[models.Prediction, models.PredictionTypedDict]
+        ] = None,
         retries: OptionalNullable[utils.RetryConfig] = UNSET,
         server_url: Optional[str] = None,
         timeout_ms: Optional[int] = None,
@@ -514,6 +568,7 @@ class Chat(BaseSDK):
         :param presence_penalty: presence_penalty determines how much the model penalizes the repetition of words or phrases. A higher presence penalty encourages the model to use a wider variety of words and phrases, making the output more diverse and creative.
         :param frequency_penalty: frequency_penalty penalizes the repetition of words based on their frequency in the generated text. A higher frequency penalty discourages the model from repeating words that have already appeared frequently in the output, promoting diversity and reducing repetition.
         :param n: Number of completions to return for each request, input tokens are only billed once.
+        :param prediction:
         :param retries: Override the default retry configuration for this method
         :param server_url: Override the default server URL for this method
         :param timeout_ms: Override the default request timeout configuration for this method in milliseconds
@@ -526,6 +581,8 @@ class Chat(BaseSDK):
 
         if server_url is not None:
             base_url = server_url
+        else:
+            base_url = self._get_url(base_url, url_variables)
 
         request = models.ChatCompletionRequest(
             model=model,
@@ -548,6 +605,9 @@ class Chat(BaseSDK):
             presence_penalty=presence_penalty,
             frequency_penalty=frequency_penalty,
             n=n,
+            prediction=utils.get_pydantic_model(
+                prediction, Optional[models.Prediction]
+            ),
         )
 
         req = self._build_request_async(
@@ -579,6 +639,7 @@ class Chat(BaseSDK):
 
         http_res = await self.do_request_async(
             hook_ctx=HookContext(
+                base_url=base_url or "",
                 operation_id="chat_completion_v1_chat_completions_post",
                 oauth2_scopes=[],
                 security_source=self.sdk_configuration.security,
@@ -588,15 +649,22 @@ class Chat(BaseSDK):
             retry_config=retry_config,
         )
 
-        data: Any = None
+        response_data: Any = None
         if utils.match_response(http_res, "200", "application/json"):
             return utils.unmarshal_json(
                 http_res.text, Optional[models.ChatCompletionResponse]
             )
         if utils.match_response(http_res, "422", "application/json"):
-            data = utils.unmarshal_json(http_res.text, models.HTTPValidationErrorData)
-            raise models.HTTPValidationError(data=data)
-        if utils.match_response(http_res, ["4XX", "5XX"], "*"):
+            response_data = utils.unmarshal_json(
+                http_res.text, models.HTTPValidationErrorData
+            )
+            raise models.HTTPValidationError(data=response_data)
+        if utils.match_response(http_res, "4XX", "*"):
+            http_res_text = await utils.stream_to_text_async(http_res)
+            raise models.SDKError(
+                "API error occurred", http_res.status_code, http_res_text, http_res
+            )
+        if utils.match_response(http_res, "5XX", "*"):
             http_res_text = await utils.stream_to_text_async(http_res)
             raise models.SDKError(
                 "API error occurred", http_res.status_code, http_res_text, http_res

mistralai_gcp/fim.py

@@ -3,7 +3,7 @@
 from .basesdk import BaseSDK
 from mistralai_gcp import models, utils
 from mistralai_gcp._hooks import HookContext
-from mistralai_gcp.types import Nullable, OptionalNullable, UNSET
+from mistralai_gcp.types import OptionalNullable, UNSET
 from mistralai_gcp.utils import eventstreaming
 from typing import Any, Mapping, Optional, Union
 
@@ -14,7 +14,7 @@ class Fim(BaseSDK):
     def stream(
         self,
         *,
-        model: Nullable[str],
+        model: str,
         prompt: str,
         temperature: OptionalNullable[float] = UNSET,
         top_p: Optional[float] = 1,
@@ -60,6 +60,8 @@ class Fim(BaseSDK):
 
         if server_url is not None:
             base_url = server_url
+        else:
+            base_url = self._get_url(base_url, url_variables)
 
         request = models.FIMCompletionStreamRequest(
             model=model,
@@ -103,6 +105,7 @@ class Fim(BaseSDK):
 
         http_res = self.do_request(
             hook_ctx=HookContext(
+                base_url=base_url or "",
                 operation_id="stream_fim",
                 oauth2_scopes=[],
                 security_source=self.sdk_configuration.security,
@@ -113,7 +116,7 @@ class Fim(BaseSDK):
             retry_config=retry_config,
         )
 
-        data: Any = None
+        response_data: Any = None
         if utils.match_response(http_res, "200", "text/event-stream"):
             return eventstreaming.EventStream(
                 http_res,
@@ -122,9 +125,16 @@ class Fim(BaseSDK):
             )
         if utils.match_response(http_res, "422", "application/json"):
             http_res_text = utils.stream_to_text(http_res)
-            data = utils.unmarshal_json(http_res_text, models.HTTPValidationErrorData)
-            raise models.HTTPValidationError(data=data)
-        if utils.match_response(http_res, ["4XX", "5XX"], "*"):
+            response_data = utils.unmarshal_json(
+                http_res_text, models.HTTPValidationErrorData
+            )
+            raise models.HTTPValidationError(data=response_data)
+        if utils.match_response(http_res, "4XX", "*"):
+            http_res_text = utils.stream_to_text(http_res)
+            raise models.SDKError(
+                "API error occurred", http_res.status_code, http_res_text, http_res
+            )
+        if utils.match_response(http_res, "5XX", "*"):
             http_res_text = utils.stream_to_text(http_res)
             raise models.SDKError(
                 "API error occurred", http_res.status_code, http_res_text, http_res
@@ -142,7 +152,7 @@ class Fim(BaseSDK):
     async def stream_async(
         self,
         *,
-        model: Nullable[str],
+        model: str,
         prompt: str,
         temperature: OptionalNullable[float] = UNSET,
         top_p: Optional[float] = 1,
@@ -188,6 +198,8 @@ class Fim(BaseSDK):
 
         if server_url is not None:
             base_url = server_url
+        else:
+            base_url = self._get_url(base_url, url_variables)
 
         request = models.FIMCompletionStreamRequest(
             model=model,
@@ -231,6 +243,7 @@ class Fim(BaseSDK):
 
         http_res = await self.do_request_async(
             hook_ctx=HookContext(
+                base_url=base_url or "",
                 operation_id="stream_fim",
                 oauth2_scopes=[],
                 security_source=self.sdk_configuration.security,
@@ -241,7 +254,7 @@ class Fim(BaseSDK):
             retry_config=retry_config,
         )
 
-        data: Any = None
+        response_data: Any = None
         if utils.match_response(http_res, "200", "text/event-stream"):
             return eventstreaming.EventStreamAsync(
                 http_res,
@@ -250,9 +263,16 @@ class Fim(BaseSDK):
             )
         if utils.match_response(http_res, "422", "application/json"):
             http_res_text = await utils.stream_to_text_async(http_res)
-            data = utils.unmarshal_json(http_res_text, models.HTTPValidationErrorData)
-            raise models.HTTPValidationError(data=data)
-        if utils.match_response(http_res, ["4XX", "5XX"], "*"):
+            response_data = utils.unmarshal_json(
+                http_res_text, models.HTTPValidationErrorData
+            )
+            raise models.HTTPValidationError(data=response_data)
+        if utils.match_response(http_res, "4XX", "*"):
+            http_res_text = await utils.stream_to_text_async(http_res)
+            raise models.SDKError(
+                "API error occurred", http_res.status_code, http_res_text, http_res
+            )
+        if utils.match_response(http_res, "5XX", "*"):
             http_res_text = await utils.stream_to_text_async(http_res)
             raise models.SDKError(
                 "API error occurred", http_res.status_code, http_res_text, http_res
@@ -270,7 +290,7 @@ class Fim(BaseSDK):
     def complete(
         self,
         *,
-        model: Nullable[str],
+        model: str,
         prompt: str,
         temperature: OptionalNullable[float] = UNSET,
         top_p: Optional[float] = 1,
@@ -316,6 +336,8 @@ class Fim(BaseSDK):
 
         if server_url is not None:
             base_url = server_url
+        else:
+            base_url = self._get_url(base_url, url_variables)
 
         request = models.FIMCompletionRequest(
             model=model,
@@ -359,6 +381,7 @@ class Fim(BaseSDK):
 
         http_res = self.do_request(
             hook_ctx=HookContext(
+                base_url=base_url or "",
                 operation_id="fim_completion_v1_fim_completions_post",
                 oauth2_scopes=[],
                 security_source=self.sdk_configuration.security,
@@ -368,15 +391,22 @@ class Fim(BaseSDK):
             retry_config=retry_config,
         )
 
-        data: Any = None
+        response_data: Any = None
         if utils.match_response(http_res, "200", "application/json"):
             return utils.unmarshal_json(
                 http_res.text, Optional[models.FIMCompletionResponse]
             )
         if utils.match_response(http_res, "422", "application/json"):
-            data = utils.unmarshal_json(http_res.text, models.HTTPValidationErrorData)
-            raise models.HTTPValidationError(data=data)
-        if utils.match_response(http_res, ["4XX", "5XX"], "*"):
+            response_data = utils.unmarshal_json(
+                http_res.text, models.HTTPValidationErrorData
+            )
+            raise models.HTTPValidationError(data=response_data)
+        if utils.match_response(http_res, "4XX", "*"):
+            http_res_text = utils.stream_to_text(http_res)
+            raise models.SDKError(
+                "API error occurred", http_res.status_code, http_res_text, http_res
+            )
+        if utils.match_response(http_res, "5XX", "*"):
             http_res_text = utils.stream_to_text(http_res)
             raise models.SDKError(
                 "API error occurred", http_res.status_code, http_res_text, http_res
@@ -394,7 +424,7 @@ class Fim(BaseSDK):
     async def complete_async(
         self,
         *,
-        model: Nullable[str],
+        model: str,
         prompt: str,
         temperature: OptionalNullable[float] = UNSET,
         top_p: Optional[float] = 1,
@@ -440,6 +470,8 @@ class Fim(BaseSDK):
 
         if server_url is not None:
             base_url = server_url
+        else:
+            base_url = self._get_url(base_url, url_variables)
 
         request = models.FIMCompletionRequest(
             model=model,
@@ -483,6 +515,7 @@ class Fim(BaseSDK):
 
         http_res = await self.do_request_async(
             hook_ctx=HookContext(
+                base_url=base_url or "",
                 operation_id="fim_completion_v1_fim_completions_post",
                 oauth2_scopes=[],
                 security_source=self.sdk_configuration.security,
@@ -492,15 +525,22 @@ class Fim(BaseSDK):
             retry_config=retry_config,
         )
 
-        data: Any = None
+        response_data: Any = None
         if utils.match_response(http_res, "200", "application/json"):
             return utils.unmarshal_json(
                 http_res.text, Optional[models.FIMCompletionResponse]
             )
         if utils.match_response(http_res, "422", "application/json"):
-            data = utils.unmarshal_json(http_res.text, models.HTTPValidationErrorData)
-            raise models.HTTPValidationError(data=data)
-        if utils.match_response(http_res, ["4XX", "5XX"], "*"):
+            response_data = utils.unmarshal_json(
+                http_res.text, models.HTTPValidationErrorData
+            )
+            raise models.HTTPValidationError(data=response_data)
+        if utils.match_response(http_res, "4XX", "*"):
+            http_res_text = await utils.stream_to_text_async(http_res)
+            raise models.SDKError(
+                "API error occurred", http_res.status_code, http_res_text, http_res
+            )
+        if utils.match_response(http_res, "5XX", "*"):
             http_res_text = await utils.stream_to_text_async(http_res)
             raise models.SDKError(
                 "API error occurred", http_res.status_code, http_res_text, http_res

mistralai_gcp/httpclient.py

@@ -1,6 +1,8 @@
 """Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT."""
 
 # pyright: reportReturnType = false
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
 from typing_extensions import Protocol, runtime_checkable
 import httpx
 from typing import Any, Optional, Union
@@ -82,3 +84,53 @@ class AsyncHttpClient(Protocol):
 
     async def aclose(self) -> None:
         pass
+
+
+class ClientOwner(Protocol):
+    client: Union[HttpClient, None]
+    async_client: Union[AsyncHttpClient, None]
+
+
+def close_clients(
+    owner: ClientOwner,
+    sync_client: Union[HttpClient, None],
+    sync_client_supplied: bool,
+    async_client: Union[AsyncHttpClient, None],
+    async_client_supplied: bool,
+) -> None:
+    """
+    A finalizer function that is meant to be used with weakref.finalize to close
+    httpx clients used by an SDK so that underlying resources can be garbage
+    collected.
+    """
+
+    # Unset the client/async_client properties so there are no more references
+    # to them from the owning SDK instance and they can be reaped.
+    owner.client = None
+    owner.async_client = None
+
+    if sync_client is not None and not sync_client_supplied:
+        try:
+            sync_client.close()
+        except Exception:
+            pass
+
+    if async_client is not None and not async_client_supplied:
+        is_async = False
+        try:
+            asyncio.get_running_loop()
+            is_async = True
+        except RuntimeError:
+            pass
+
+        try:
+            # If this function is called in an async loop then start another
+            # loop in a separate thread to close the async http client.
+            if is_async:
+                with ThreadPoolExecutor(max_workers=1) as executor:
+                    future = executor.submit(asyncio.run, async_client.aclose())
+                    future.result()
+            else:
+                asyncio.run(async_client.aclose())
+        except Exception:
+            pass

mistralai_gcp/models/__init__.py

@@ -67,6 +67,8 @@ from .functioncall import (
 )
 from .functionname import FunctionName, FunctionNameTypedDict
 from .httpvalidationerror import HTTPValidationError, HTTPValidationErrorData
+from .jsonschema import JSONSchema, JSONSchemaTypedDict
+from .prediction import Prediction, PredictionTypedDict
 from .referencechunk import ReferenceChunk, ReferenceChunkType, ReferenceChunkTypedDict
 from .responseformat import ResponseFormat, ResponseFormatTypedDict
 from .responseformats import ResponseFormats
@@ -107,6 +109,7 @@ from .validationerror import (
     ValidationErrorTypedDict,
 )
 
+
 __all__ = [
     "Arguments",
     "ArgumentsTypedDict",
@@ -163,10 +166,14 @@ __all__ = [
     "FunctionTypedDict",
     "HTTPValidationError",
     "HTTPValidationErrorData",
+    "JSONSchema",
+    "JSONSchemaTypedDict",
     "Loc",
     "LocTypedDict",
     "Messages",
     "MessagesTypedDict",
+    "Prediction",
+    "PredictionTypedDict",
     "ReferenceChunk",
     "ReferenceChunkType",
     "ReferenceChunkTypedDict",

mistralai_gcp/models/assistantmessage.py

@@ -32,6 +32,7 @@ class AssistantMessageTypedDict(TypedDict):
     content: NotRequired[Nullable[AssistantMessageContentTypedDict]]
     tool_calls: NotRequired[Nullable[List[ToolCallTypedDict]]]
     prefix: NotRequired[bool]
+    r"""Set this to `true` when adding an assistant message as prefix to condition the model response. The role of the prefix message is to force the model to start its answer by the content of the message."""
     role: NotRequired[AssistantMessageRole]
 
 
@@ -41,6 +42,7 @@ class AssistantMessage(BaseModel):
     tool_calls: OptionalNullable[List[ToolCall]] = UNSET
 
     prefix: Optional[bool] = False
+    r"""Set this to `true` when adding an assistant message as prefix to condition the model response. The role of the prefix message is to force the model to start its answer by the content of the message."""
 
     role: Optional[AssistantMessageRole] = "assistant"