mistralai 1.10.0__py3-none-any.whl → 1.11.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (270) hide show
  1. mistralai/_hooks/tracing.py +28 -3
  2. mistralai/_version.py +3 -3
  3. mistralai/accesses.py +22 -12
  4. mistralai/agents.py +88 -44
  5. mistralai/basesdk.py +6 -0
  6. mistralai/chat.py +96 -40
  7. mistralai/classifiers.py +48 -23
  8. mistralai/conversations.py +186 -64
  9. mistralai/documents.py +72 -26
  10. mistralai/embeddings.py +24 -9
  11. mistralai/extra/README.md +1 -1
  12. mistralai/extra/mcp/auth.py +10 -11
  13. mistralai/extra/mcp/base.py +17 -16
  14. mistralai/extra/mcp/sse.py +13 -15
  15. mistralai/extra/mcp/stdio.py +5 -6
  16. mistralai/extra/observability/otel.py +47 -68
  17. mistralai/extra/run/context.py +33 -43
  18. mistralai/extra/run/result.py +29 -30
  19. mistralai/extra/run/tools.py +8 -9
  20. mistralai/extra/struct_chat.py +15 -8
  21. mistralai/extra/utils/response_format.py +5 -3
  22. mistralai/files.py +58 -24
  23. mistralai/fim.py +20 -12
  24. mistralai/httpclient.py +0 -1
  25. mistralai/jobs.py +65 -26
  26. mistralai/libraries.py +20 -10
  27. mistralai/mistral_agents.py +438 -30
  28. mistralai/mistral_jobs.py +62 -17
  29. mistralai/models/__init__.py +46 -1
  30. mistralai/models/agent.py +1 -1
  31. mistralai/models/agentconversation.py +1 -1
  32. mistralai/models/agenthandoffdoneevent.py +1 -1
  33. mistralai/models/agenthandoffentry.py +3 -2
  34. mistralai/models/agenthandoffstartedevent.py +1 -1
  35. mistralai/models/agents_api_v1_agents_get_versionop.py +21 -0
  36. mistralai/models/agents_api_v1_agents_list_versionsop.py +33 -0
  37. mistralai/models/agents_api_v1_agents_listop.py +5 -1
  38. mistralai/models/agents_api_v1_conversations_listop.py +1 -1
  39. mistralai/models/agentscompletionrequest.py +2 -5
  40. mistralai/models/agentscompletionstreamrequest.py +2 -5
  41. mistralai/models/archiveftmodelout.py +1 -1
  42. mistralai/models/assistantmessage.py +1 -1
  43. mistralai/models/audiochunk.py +1 -1
  44. mistralai/models/audioencoding.py +18 -0
  45. mistralai/models/audioformat.py +17 -0
  46. mistralai/models/basemodelcard.py +1 -1
  47. mistralai/models/batchjobin.py +18 -9
  48. mistralai/models/batchjobout.py +6 -1
  49. mistralai/models/batchjobsout.py +1 -1
  50. mistralai/models/batchrequest.py +48 -0
  51. mistralai/models/chatcompletionchoice.py +10 -5
  52. mistralai/models/chatcompletionrequest.py +2 -5
  53. mistralai/models/chatcompletionstreamrequest.py +2 -5
  54. mistralai/models/classificationrequest.py +37 -3
  55. mistralai/models/classifierdetailedjobout.py +4 -2
  56. mistralai/models/classifierftmodelout.py +3 -2
  57. mistralai/models/classifierjobout.py +4 -2
  58. mistralai/models/codeinterpretertool.py +1 -1
  59. mistralai/models/completiondetailedjobout.py +5 -2
  60. mistralai/models/completionftmodelout.py +3 -2
  61. mistralai/models/completionjobout.py +5 -2
  62. mistralai/models/completionresponsestreamchoice.py +9 -8
  63. mistralai/models/conversationappendrequest.py +4 -1
  64. mistralai/models/conversationappendstreamrequest.py +4 -1
  65. mistralai/models/conversationhistory.py +2 -1
  66. mistralai/models/conversationmessages.py +1 -1
  67. mistralai/models/conversationrequest.py +5 -1
  68. mistralai/models/conversationresponse.py +2 -1
  69. mistralai/models/conversationrestartrequest.py +4 -1
  70. mistralai/models/conversationrestartstreamrequest.py +4 -1
  71. mistralai/models/conversationstreamrequest.py +5 -1
  72. mistralai/models/documentlibrarytool.py +1 -1
  73. mistralai/models/documenturlchunk.py +1 -1
  74. mistralai/models/embeddingdtype.py +7 -1
  75. mistralai/models/embeddingrequest.py +11 -3
  76. mistralai/models/encodingformat.py +4 -1
  77. mistralai/models/entitytype.py +8 -1
  78. mistralai/models/filepurpose.py +8 -1
  79. mistralai/models/files_api_routes_list_filesop.py +4 -11
  80. mistralai/models/files_api_routes_upload_fileop.py +2 -6
  81. mistralai/models/fileschema.py +3 -5
  82. mistralai/models/finetuneablemodeltype.py +4 -1
  83. mistralai/models/ftclassifierlossfunction.py +4 -1
  84. mistralai/models/ftmodelcard.py +1 -1
  85. mistralai/models/functioncallentry.py +3 -2
  86. mistralai/models/functioncallevent.py +1 -1
  87. mistralai/models/functionresultentry.py +3 -2
  88. mistralai/models/functiontool.py +1 -1
  89. mistralai/models/githubrepositoryin.py +1 -1
  90. mistralai/models/githubrepositoryout.py +1 -1
  91. mistralai/models/httpvalidationerror.py +4 -2
  92. mistralai/models/imagegenerationtool.py +1 -1
  93. mistralai/models/imageurlchunk.py +1 -1
  94. mistralai/models/jobs_api_routes_batch_get_batch_jobop.py +40 -3
  95. mistralai/models/jobsout.py +1 -1
  96. mistralai/models/legacyjobmetadataout.py +1 -1
  97. mistralai/models/messageinputentry.py +9 -3
  98. mistralai/models/messageoutputentry.py +6 -3
  99. mistralai/models/messageoutputevent.py +4 -2
  100. mistralai/models/mistralerror.py +11 -7
  101. mistralai/models/mistralpromptmode.py +1 -1
  102. mistralai/models/modelconversation.py +1 -1
  103. mistralai/models/no_response_error.py +5 -1
  104. mistralai/models/ocrrequest.py +11 -1
  105. mistralai/models/ocrtableobject.py +4 -1
  106. mistralai/models/referencechunk.py +1 -1
  107. mistralai/models/requestsource.py +5 -1
  108. mistralai/models/responsedoneevent.py +1 -1
  109. mistralai/models/responseerrorevent.py +1 -1
  110. mistralai/models/responseformats.py +5 -1
  111. mistralai/models/responsestartedevent.py +1 -1
  112. mistralai/models/responsevalidationerror.py +2 -0
  113. mistralai/models/retrievefileout.py +3 -5
  114. mistralai/models/sampletype.py +7 -1
  115. mistralai/models/sdkerror.py +2 -0
  116. mistralai/models/shareenum.py +7 -1
  117. mistralai/models/sharingdelete.py +2 -4
  118. mistralai/models/sharingin.py +3 -5
  119. mistralai/models/source.py +8 -1
  120. mistralai/models/systemmessage.py +1 -1
  121. mistralai/models/textchunk.py +1 -1
  122. mistralai/models/thinkchunk.py +1 -1
  123. mistralai/models/timestampgranularity.py +1 -1
  124. mistralai/models/tool.py +2 -6
  125. mistralai/models/toolcall.py +2 -6
  126. mistralai/models/toolchoice.py +2 -6
  127. mistralai/models/toolchoiceenum.py +6 -1
  128. mistralai/models/toolexecutiondeltaevent.py +2 -1
  129. mistralai/models/toolexecutiondoneevent.py +2 -1
  130. mistralai/models/toolexecutionentry.py +4 -2
  131. mistralai/models/toolexecutionstartedevent.py +2 -1
  132. mistralai/models/toolfilechunk.py +13 -5
  133. mistralai/models/toolmessage.py +1 -1
  134. mistralai/models/toolreferencechunk.py +15 -5
  135. mistralai/models/tooltypes.py +1 -1
  136. mistralai/models/transcriptionsegmentchunk.py +1 -1
  137. mistralai/models/transcriptionstreamdone.py +1 -1
  138. mistralai/models/transcriptionstreamlanguage.py +1 -1
  139. mistralai/models/transcriptionstreamsegmentdelta.py +1 -1
  140. mistralai/models/transcriptionstreamtextdelta.py +1 -1
  141. mistralai/models/unarchiveftmodelout.py +1 -1
  142. mistralai/models/uploadfileout.py +3 -5
  143. mistralai/models/usermessage.py +1 -1
  144. mistralai/models/wandbintegration.py +1 -1
  145. mistralai/models/wandbintegrationout.py +1 -1
  146. mistralai/models/websearchpremiumtool.py +1 -1
  147. mistralai/models/websearchtool.py +1 -1
  148. mistralai/models_.py +24 -12
  149. mistralai/ocr.py +38 -10
  150. mistralai/sdk.py +2 -2
  151. mistralai/transcriptions.py +28 -12
  152. mistralai/types/basemodel.py +41 -3
  153. mistralai/utils/__init__.py +0 -3
  154. mistralai/utils/annotations.py +32 -8
  155. mistralai/utils/enums.py +60 -0
  156. mistralai/utils/forms.py +21 -10
  157. mistralai/utils/queryparams.py +14 -2
  158. mistralai/utils/requestbodies.py +3 -3
  159. mistralai/utils/retries.py +69 -5
  160. mistralai/utils/serializers.py +0 -20
  161. mistralai/utils/unmarshal_json_response.py +15 -1
  162. {mistralai-1.10.0.dist-info → mistralai-1.11.1.dist-info}/METADATA +144 -159
  163. mistralai-1.11.1.dist-info/RECORD +495 -0
  164. {mistralai-1.10.0.dist-info → mistralai-1.11.1.dist-info}/WHEEL +1 -1
  165. mistralai_azure/_version.py +3 -3
  166. mistralai_azure/basesdk.py +21 -5
  167. mistralai_azure/chat.py +82 -109
  168. mistralai_azure/httpclient.py +0 -1
  169. mistralai_azure/models/__init__.py +66 -4
  170. mistralai_azure/models/assistantmessage.py +1 -1
  171. mistralai_azure/models/chatcompletionchoice.py +10 -7
  172. mistralai_azure/models/chatcompletionrequest.py +24 -10
  173. mistralai_azure/models/chatcompletionstreamrequest.py +24 -10
  174. mistralai_azure/models/completionresponsestreamchoice.py +11 -7
  175. mistralai_azure/models/documenturlchunk.py +1 -1
  176. mistralai_azure/models/httpvalidationerror.py +15 -8
  177. mistralai_azure/models/imageurlchunk.py +1 -1
  178. mistralai_azure/models/mistralazureerror.py +30 -0
  179. mistralai_azure/models/mistralpromptmode.py +1 -1
  180. mistralai_azure/models/no_response_error.py +17 -0
  181. mistralai_azure/models/ocrpageobject.py +32 -5
  182. mistralai_azure/models/ocrrequest.py +20 -1
  183. mistralai_azure/models/ocrtableobject.py +34 -0
  184. mistralai_azure/models/prediction.py +4 -0
  185. mistralai_azure/models/referencechunk.py +1 -1
  186. mistralai_azure/models/responseformat.py +4 -2
  187. mistralai_azure/models/responseformats.py +5 -2
  188. mistralai_azure/models/responsevalidationerror.py +27 -0
  189. mistralai_azure/models/sdkerror.py +32 -14
  190. mistralai_azure/models/systemmessage.py +8 -4
  191. mistralai_azure/models/systemmessagecontentchunks.py +21 -0
  192. mistralai_azure/models/textchunk.py +1 -1
  193. mistralai_azure/models/thinkchunk.py +35 -0
  194. mistralai_azure/models/tool.py +2 -6
  195. mistralai_azure/models/toolcall.py +2 -6
  196. mistralai_azure/models/toolchoice.py +2 -6
  197. mistralai_azure/models/toolchoiceenum.py +6 -1
  198. mistralai_azure/models/toolmessage.py +1 -1
  199. mistralai_azure/models/tooltypes.py +1 -1
  200. mistralai_azure/models/usermessage.py +1 -1
  201. mistralai_azure/ocr.py +39 -40
  202. mistralai_azure/types/basemodel.py +41 -3
  203. mistralai_azure/utils/__init__.py +18 -8
  204. mistralai_azure/utils/annotations.py +32 -8
  205. mistralai_azure/utils/enums.py +60 -0
  206. mistralai_azure/utils/eventstreaming.py +10 -0
  207. mistralai_azure/utils/forms.py +21 -10
  208. mistralai_azure/utils/queryparams.py +14 -2
  209. mistralai_azure/utils/requestbodies.py +3 -3
  210. mistralai_azure/utils/retries.py +69 -5
  211. mistralai_azure/utils/serializers.py +3 -22
  212. mistralai_azure/utils/unmarshal_json_response.py +38 -0
  213. mistralai_gcp/_hooks/types.py +7 -0
  214. mistralai_gcp/_version.py +4 -4
  215. mistralai_gcp/basesdk.py +33 -25
  216. mistralai_gcp/chat.py +98 -109
  217. mistralai_gcp/fim.py +62 -85
  218. mistralai_gcp/httpclient.py +6 -17
  219. mistralai_gcp/models/__init__.py +321 -116
  220. mistralai_gcp/models/assistantmessage.py +2 -2
  221. mistralai_gcp/models/chatcompletionchoice.py +10 -7
  222. mistralai_gcp/models/chatcompletionrequest.py +38 -7
  223. mistralai_gcp/models/chatcompletionresponse.py +6 -6
  224. mistralai_gcp/models/chatcompletionstreamrequest.py +38 -7
  225. mistralai_gcp/models/completionresponsestreamchoice.py +12 -8
  226. mistralai_gcp/models/deltamessage.py +1 -1
  227. mistralai_gcp/models/fimcompletionrequest.py +9 -10
  228. mistralai_gcp/models/fimcompletionresponse.py +6 -6
  229. mistralai_gcp/models/fimcompletionstreamrequest.py +9 -10
  230. mistralai_gcp/models/httpvalidationerror.py +15 -8
  231. mistralai_gcp/models/imageurl.py +1 -1
  232. mistralai_gcp/models/imageurlchunk.py +1 -1
  233. mistralai_gcp/models/jsonschema.py +1 -1
  234. mistralai_gcp/models/mistralgcperror.py +30 -0
  235. mistralai_gcp/models/mistralpromptmode.py +8 -0
  236. mistralai_gcp/models/no_response_error.py +17 -0
  237. mistralai_gcp/models/prediction.py +4 -0
  238. mistralai_gcp/models/referencechunk.py +1 -1
  239. mistralai_gcp/models/responseformat.py +5 -3
  240. mistralai_gcp/models/responseformats.py +5 -2
  241. mistralai_gcp/models/responsevalidationerror.py +27 -0
  242. mistralai_gcp/models/sdkerror.py +32 -14
  243. mistralai_gcp/models/systemmessage.py +8 -4
  244. mistralai_gcp/models/systemmessagecontentchunks.py +21 -0
  245. mistralai_gcp/models/textchunk.py +1 -1
  246. mistralai_gcp/models/thinkchunk.py +35 -0
  247. mistralai_gcp/models/tool.py +2 -6
  248. mistralai_gcp/models/toolcall.py +2 -6
  249. mistralai_gcp/models/toolchoice.py +2 -6
  250. mistralai_gcp/models/toolchoiceenum.py +6 -1
  251. mistralai_gcp/models/toolmessage.py +2 -2
  252. mistralai_gcp/models/tooltypes.py +1 -1
  253. mistralai_gcp/models/usageinfo.py +71 -8
  254. mistralai_gcp/models/usermessage.py +2 -2
  255. mistralai_gcp/sdk.py +12 -10
  256. mistralai_gcp/sdkconfiguration.py +0 -7
  257. mistralai_gcp/types/basemodel.py +41 -3
  258. mistralai_gcp/utils/__init__.py +141 -46
  259. mistralai_gcp/utils/annotations.py +32 -8
  260. mistralai_gcp/utils/datetimes.py +23 -0
  261. mistralai_gcp/utils/enums.py +125 -25
  262. mistralai_gcp/utils/eventstreaming.py +10 -0
  263. mistralai_gcp/utils/forms.py +62 -30
  264. mistralai_gcp/utils/queryparams.py +14 -2
  265. mistralai_gcp/utils/requestbodies.py +3 -3
  266. mistralai_gcp/utils/retries.py +69 -5
  267. mistralai_gcp/utils/serializers.py +33 -23
  268. mistralai_gcp/utils/unmarshal_json_response.py +38 -0
  269. mistralai-1.10.0.dist-info/RECORD +0 -475
  270. {mistralai-1.10.0.dist-info → mistralai-1.11.1.dist-info}/licenses/LICENSE +0 -0
mistralai/utils/enums.py CHANGED
@@ -2,6 +2,10 @@
2
2
 
3
3
  import enum
4
4
  import sys
5
+ from typing import Any
6
+
7
+ from pydantic_core import core_schema
8
+
5
9
 
6
10
  class OpenEnumMeta(enum.EnumMeta):
7
11
  # The __call__ method `boundary` kwarg was added in 3.11 and must be present
@@ -72,3 +76,59 @@ class OpenEnumMeta(enum.EnumMeta):
72
76
  )
73
77
  except ValueError:
74
78
  return value
79
+
80
+ def __new__(mcs, name, bases, namespace, **kwargs):
81
+ cls = super().__new__(mcs, name, bases, namespace, **kwargs)
82
+
83
+ # Add __get_pydantic_core_schema__ to make open enums work correctly
84
+ # in union discrimination. In strict mode (used by Pydantic for unions),
85
+ # only known enum values match. In lax mode, unknown values are accepted.
86
+ def __get_pydantic_core_schema__(
87
+ cls_inner: Any, _source_type: Any, _handler: Any
88
+ ) -> core_schema.CoreSchema:
89
+ # Create a validator that only accepts known enum values (for strict mode)
90
+ def validate_strict(v: Any) -> Any:
91
+ if isinstance(v, cls_inner):
92
+ return v
93
+ # Use the parent EnumMeta's __call__ which raises ValueError for unknown values
94
+ return enum.EnumMeta.__call__(cls_inner, v)
95
+
96
+ # Create a lax validator that accepts unknown values
97
+ def validate_lax(v: Any) -> Any:
98
+ if isinstance(v, cls_inner):
99
+ return v
100
+ try:
101
+ return enum.EnumMeta.__call__(cls_inner, v)
102
+ except ValueError:
103
+ # Return the raw value for unknown enum values
104
+ return v
105
+
106
+ # Determine the base type schema (str or int)
107
+ is_int_enum = False
108
+ for base in cls_inner.__mro__:
109
+ if base is int:
110
+ is_int_enum = True
111
+ break
112
+ if base is str:
113
+ break
114
+
115
+ base_schema = (
116
+ core_schema.int_schema()
117
+ if is_int_enum
118
+ else core_schema.str_schema()
119
+ )
120
+
121
+ # Use lax_or_strict_schema:
122
+ # - strict mode: only known enum values match (raises ValueError for unknown)
123
+ # - lax mode: accept any value, return enum member or raw value
124
+ return core_schema.lax_or_strict_schema(
125
+ lax_schema=core_schema.chain_schema(
126
+ [base_schema, core_schema.no_info_plain_validator_function(validate_lax)]
127
+ ),
128
+ strict_schema=core_schema.chain_schema(
129
+ [base_schema, core_schema.no_info_plain_validator_function(validate_strict)]
130
+ ),
131
+ )
132
+
133
+ setattr(cls, "__get_pydantic_core_schema__", classmethod(__get_pydantic_core_schema__))
134
+ return cls
mistralai/utils/forms.py CHANGED
@@ -142,16 +142,21 @@ def serialize_multipart_form(
142
142
  if field_metadata.file:
143
143
  if isinstance(val, List):
144
144
  # Handle array of files
145
+ array_field_name = f_name + "[]"
145
146
  for file_obj in val:
146
147
  if not _is_set(file_obj):
147
148
  continue
148
-
149
- file_name, content, content_type = _extract_file_properties(file_obj)
149
+
150
+ file_name, content, content_type = _extract_file_properties(
151
+ file_obj
152
+ )
150
153
 
151
154
  if content_type is not None:
152
- files.append((f_name + "[]", (file_name, content, content_type)))
155
+ files.append(
156
+ (array_field_name, (file_name, content, content_type))
157
+ )
153
158
  else:
154
- files.append((f_name + "[]", (file_name, content)))
159
+ files.append((array_field_name, (file_name, content)))
155
160
  else:
156
161
  # Handle single file
157
162
  file_name, content, content_type = _extract_file_properties(val)
@@ -161,11 +166,16 @@ def serialize_multipart_form(
161
166
  else:
162
167
  files.append((f_name, (file_name, content)))
163
168
  elif field_metadata.json:
164
- files.append((f_name, (
165
- None,
166
- marshal_json(val, request_field_types[name]),
167
- "application/json",
168
- )))
169
+ files.append(
170
+ (
171
+ f_name,
172
+ (
173
+ None,
174
+ marshal_json(val, request_field_types[name]),
175
+ "application/json",
176
+ ),
177
+ )
178
+ )
169
179
  else:
170
180
  if isinstance(val, List):
171
181
  values = []
@@ -175,7 +185,8 @@ def serialize_multipart_form(
175
185
  continue
176
186
  values.append(_val_to_string(value))
177
187
 
178
- form[f_name + "[]"] = values
188
+ array_field_name = f_name + "[]"
189
+ form[array_field_name] = values
179
190
  else:
180
191
  form[f_name] = _val_to_string(val)
181
192
  return media_type, form, files
@@ -27,12 +27,13 @@ from .forms import _populate_form
27
27
  def get_query_params(
28
28
  query_params: Any,
29
29
  gbls: Optional[Any] = None,
30
+ allow_empty_value: Optional[List[str]] = None,
30
31
  ) -> Dict[str, List[str]]:
31
32
  params: Dict[str, List[str]] = {}
32
33
 
33
- globals_already_populated = _populate_query_params(query_params, gbls, params, [])
34
+ globals_already_populated = _populate_query_params(query_params, gbls, params, [], allow_empty_value)
34
35
  if _is_set(gbls):
35
- _populate_query_params(gbls, None, params, globals_already_populated)
36
+ _populate_query_params(gbls, None, params, globals_already_populated, allow_empty_value)
36
37
 
37
38
  return params
38
39
 
@@ -42,6 +43,7 @@ def _populate_query_params(
42
43
  gbls: Any,
43
44
  query_param_values: Dict[str, List[str]],
44
45
  skip_fields: List[str],
46
+ allow_empty_value: Optional[List[str]] = None,
45
47
  ) -> List[str]:
46
48
  globals_already_populated: List[str] = []
47
49
 
@@ -69,6 +71,16 @@ def _populate_query_params(
69
71
  globals_already_populated.append(name)
70
72
 
71
73
  f_name = field.alias if field.alias is not None else name
74
+
75
+ allow_empty_set = set(allow_empty_value or [])
76
+ should_include_empty = f_name in allow_empty_set and (
77
+ value is None or value == [] or value == ""
78
+ )
79
+
80
+ if should_include_empty:
81
+ query_param_values[f_name] = [""]
82
+ continue
83
+
72
84
  serialization = metadata.serialization
73
85
  if serialization is not None:
74
86
  serialized_parms = _get_serialized_params(
@@ -44,15 +44,15 @@ def serialize_request_body(
44
44
 
45
45
  serialized_request_body = SerializedRequestBody(media_type)
46
46
 
47
- if re.match(r"(application|text)\/.*?\+*json.*", media_type) is not None:
47
+ if re.match(r"^(application|text)\/([^+]+\+)*json.*", media_type) is not None:
48
48
  serialized_request_body.content = marshal_json(request_body, request_body_type)
49
- elif re.match(r"multipart\/.*", media_type) is not None:
49
+ elif re.match(r"^multipart\/.*", media_type) is not None:
50
50
  (
51
51
  serialized_request_body.media_type,
52
52
  serialized_request_body.data,
53
53
  serialized_request_body.files,
54
54
  ) = serialize_multipart_form(media_type, request_body)
55
- elif re.match(r"application\/x-www-form-urlencoded.*", media_type) is not None:
55
+ elif re.match(r"^application\/x-www-form-urlencoded.*", media_type) is not None:
56
56
  serialized_request_body.data = serialize_form_data(request_body)
57
57
  elif isinstance(request_body, (bytes, bytearray, io.BytesIO, io.BufferedReader)):
58
58
  serialized_request_body.content = request_body
@@ -3,7 +3,9 @@
3
3
  import asyncio
4
4
  import random
5
5
  import time
6
- from typing import List
6
+ from datetime import datetime
7
+ from email.utils import parsedate_to_datetime
8
+ from typing import List, Optional
7
9
 
8
10
  import httpx
9
11
 
@@ -51,9 +53,11 @@ class Retries:
51
53
 
52
54
  class TemporaryError(Exception):
53
55
  response: httpx.Response
56
+ retry_after: Optional[int]
54
57
 
55
58
  def __init__(self, response: httpx.Response):
56
59
  self.response = response
60
+ self.retry_after = _parse_retry_after_header(response)
57
61
 
58
62
 
59
63
  class PermanentError(Exception):
@@ -63,6 +67,62 @@ class PermanentError(Exception):
63
67
  self.inner = inner
64
68
 
65
69
 
70
+ def _parse_retry_after_header(response: httpx.Response) -> Optional[int]:
71
+ """Parse Retry-After header from response.
72
+
73
+ Returns:
74
+ Retry interval in milliseconds, or None if header is missing or invalid.
75
+ """
76
+ retry_after_header = response.headers.get("retry-after")
77
+ if not retry_after_header:
78
+ return None
79
+
80
+ try:
81
+ seconds = float(retry_after_header)
82
+ return round(seconds * 1000)
83
+ except ValueError:
84
+ pass
85
+
86
+ try:
87
+ retry_date = parsedate_to_datetime(retry_after_header)
88
+ delta = (retry_date - datetime.now(retry_date.tzinfo)).total_seconds()
89
+ return round(max(0, delta) * 1000)
90
+ except (ValueError, TypeError):
91
+ pass
92
+
93
+ return None
94
+
95
+
96
+ def _get_sleep_interval(
97
+ exception: Exception,
98
+ initial_interval: int,
99
+ max_interval: int,
100
+ exponent: float,
101
+ retries: int,
102
+ ) -> float:
103
+ """Get sleep interval for retry with exponential backoff.
104
+
105
+ Args:
106
+ exception: The exception that triggered the retry.
107
+ initial_interval: Initial retry interval in milliseconds.
108
+ max_interval: Maximum retry interval in milliseconds.
109
+ exponent: Base for exponential backoff calculation.
110
+ retries: Current retry attempt count.
111
+
112
+ Returns:
113
+ Sleep interval in seconds.
114
+ """
115
+ if (
116
+ isinstance(exception, TemporaryError)
117
+ and exception.retry_after is not None
118
+ and exception.retry_after > 0
119
+ ):
120
+ return exception.retry_after / 1000
121
+
122
+ sleep = (initial_interval / 1000) * exponent**retries + random.uniform(0, 1)
123
+ return min(sleep, max_interval / 1000)
124
+
125
+
66
126
  def retry(func, retries: Retries):
67
127
  if retries.config.strategy == "backoff":
68
128
 
@@ -183,8 +243,10 @@ def retry_with_backoff(
183
243
  return exception.response
184
244
 
185
245
  raise
186
- sleep = (initial_interval / 1000) * exponent**retries + random.uniform(0, 1)
187
- sleep = min(sleep, max_interval / 1000)
246
+
247
+ sleep = _get_sleep_interval(
248
+ exception, initial_interval, max_interval, exponent, retries
249
+ )
188
250
  time.sleep(sleep)
189
251
  retries += 1
190
252
 
@@ -211,7 +273,9 @@ async def retry_with_backoff_async(
211
273
  return exception.response
212
274
 
213
275
  raise
214
- sleep = (initial_interval / 1000) * exponent**retries + random.uniform(0, 1)
215
- sleep = min(sleep, max_interval / 1000)
276
+
277
+ sleep = _get_sleep_interval(
278
+ exception, initial_interval, max_interval, exponent, retries
279
+ )
216
280
  await asyncio.sleep(sleep)
217
281
  retries += 1
@@ -102,26 +102,6 @@ def validate_int(b):
102
102
  return int(b)
103
103
 
104
104
 
105
- def validate_open_enum(is_int: bool):
106
- def validate(e):
107
- if e is None:
108
- return None
109
-
110
- if isinstance(e, Unset):
111
- return e
112
-
113
- if is_int:
114
- if not isinstance(e, int):
115
- raise ValueError("Expected int")
116
- else:
117
- if not isinstance(e, str):
118
- raise ValueError("Expected string")
119
-
120
- return e
121
-
122
- return validate
123
-
124
-
125
105
  def validate_const(v):
126
106
  def validate(c):
127
107
  # Optional[T] is a Union[T, None]
@@ -1,12 +1,26 @@
1
1
  """Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT."""
2
2
 
3
- from typing import Any, Optional
3
+ from typing import Any, Optional, Type, TypeVar, overload
4
4
 
5
5
  import httpx
6
6
 
7
7
  from .serializers import unmarshal_json
8
8
  from mistralai import models
9
9
 
10
+ T = TypeVar("T")
11
+
12
+
13
+ @overload
14
+ def unmarshal_json_response(
15
+ typ: Type[T], http_res: httpx.Response, body: Optional[str] = None
16
+ ) -> T: ...
17
+
18
+
19
+ @overload
20
+ def unmarshal_json_response(
21
+ typ: Any, http_res: httpx.Response, body: Optional[str] = None
22
+ ) -> Any: ...
23
+
10
24
 
11
25
  def unmarshal_json_response(
12
26
  typ: Any, http_res: httpx.Response, body: Optional[str] = None