mixedbread 0.1.0a1__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 (94) hide show
  1. mixedbread/__init__.py +96 -0
  2. mixedbread/_base_client.py +2045 -0
  3. mixedbread/_client.py +867 -0
  4. mixedbread/_compat.py +219 -0
  5. mixedbread/_constants.py +14 -0
  6. mixedbread/_exceptions.py +108 -0
  7. mixedbread/_files.py +123 -0
  8. mixedbread/_models.py +791 -0
  9. mixedbread/_qs.py +150 -0
  10. mixedbread/_resource.py +43 -0
  11. mixedbread/_response.py +826 -0
  12. mixedbread/_streaming.py +333 -0
  13. mixedbread/_types.py +217 -0
  14. mixedbread/_utils/__init__.py +57 -0
  15. mixedbread/_utils/_logs.py +25 -0
  16. mixedbread/_utils/_proxy.py +62 -0
  17. mixedbread/_utils/_reflection.py +42 -0
  18. mixedbread/_utils/_streams.py +12 -0
  19. mixedbread/_utils/_sync.py +71 -0
  20. mixedbread/_utils/_transform.py +392 -0
  21. mixedbread/_utils/_typing.py +149 -0
  22. mixedbread/_utils/_utils.py +414 -0
  23. mixedbread/_version.py +4 -0
  24. mixedbread/lib/.keep +4 -0
  25. mixedbread/lib/polling.py +123 -0
  26. mixedbread/py.typed +0 -0
  27. mixedbread/resources/__init__.py +89 -0
  28. mixedbread/resources/chat/__init__.py +33 -0
  29. mixedbread/resources/chat/chat.py +102 -0
  30. mixedbread/resources/chat/completions.py +164 -0
  31. mixedbread/resources/document_ai/__init__.py +47 -0
  32. mixedbread/resources/document_ai/document_ai.py +134 -0
  33. mixedbread/resources/document_ai/extract/__init__.py +33 -0
  34. mixedbread/resources/document_ai/extract/extract.py +452 -0
  35. mixedbread/resources/document_ai/extract/schema.py +356 -0
  36. mixedbread/resources/document_ai/parse.py +480 -0
  37. mixedbread/resources/embeddings.py +241 -0
  38. mixedbread/resources/files.py +706 -0
  39. mixedbread/resources/rerankings.py +226 -0
  40. mixedbread/resources/vector_stores/__init__.py +33 -0
  41. mixedbread/resources/vector_stores/files.py +712 -0
  42. mixedbread/resources/vector_stores/vector_stores.py +965 -0
  43. mixedbread/types/__init__.py +30 -0
  44. mixedbread/types/chat/__init__.py +3 -0
  45. mixedbread/types/client_embed_params.py +55 -0
  46. mixedbread/types/client_rerank_params.py +27 -0
  47. mixedbread/types/document_ai/__init__.py +12 -0
  48. mixedbread/types/document_ai/extract/__init__.py +10 -0
  49. mixedbread/types/document_ai/extract/created_json_schema.py +11 -0
  50. mixedbread/types/document_ai/extract/enhanced_json_schema.py +11 -0
  51. mixedbread/types/document_ai/extract/schema_create_params.py +12 -0
  52. mixedbread/types/document_ai/extract/schema_enhance_params.py +12 -0
  53. mixedbread/types/document_ai/extract/schema_validate_params.py +12 -0
  54. mixedbread/types/document_ai/extract/validated_json_schema.py +18 -0
  55. mixedbread/types/document_ai/extract_content_params.py +15 -0
  56. mixedbread/types/document_ai/extract_create_job_params.py +15 -0
  57. mixedbread/types/document_ai/extract_create_job_response.py +83 -0
  58. mixedbread/types/document_ai/extract_retrieve_job_response.py +83 -0
  59. mixedbread/types/document_ai/parse_create_job_params.py +38 -0
  60. mixedbread/types/document_ai/parse_create_job_response.py +161 -0
  61. mixedbread/types/document_ai/parse_retrieve_job_response.py +161 -0
  62. mixedbread/types/document_ai/result.py +11 -0
  63. mixedbread/types/embed_response.py +97 -0
  64. mixedbread/types/embedding_create_params.py +55 -0
  65. mixedbread/types/embedding_create_response.py +97 -0
  66. mixedbread/types/file_create_params.py +14 -0
  67. mixedbread/types/file_deleted.py +19 -0
  68. mixedbread/types/file_list_params.py +15 -0
  69. mixedbread/types/file_list_response.py +30 -0
  70. mixedbread/types/file_object.py +30 -0
  71. mixedbread/types/file_update_params.py +14 -0
  72. mixedbread/types/info_response.py +12 -0
  73. mixedbread/types/rerank_response.py +77 -0
  74. mixedbread/types/reranking_create_params.py +27 -0
  75. mixedbread/types/reranking_create_response.py +77 -0
  76. mixedbread/types/search_response.py +103 -0
  77. mixedbread/types/vector_store.py +69 -0
  78. mixedbread/types/vector_store_create_params.py +33 -0
  79. mixedbread/types/vector_store_delete_response.py +19 -0
  80. mixedbread/types/vector_store_list_params.py +15 -0
  81. mixedbread/types/vector_store_list_response.py +30 -0
  82. mixedbread/types/vector_store_question_answering_params.py +235 -0
  83. mixedbread/types/vector_store_search_params.py +202 -0
  84. mixedbread/types/vector_store_update_params.py +30 -0
  85. mixedbread/types/vector_stores/__init__.py +9 -0
  86. mixedbread/types/vector_stores/file_create_params.py +16 -0
  87. mixedbread/types/vector_stores/file_delete_response.py +19 -0
  88. mixedbread/types/vector_stores/file_list_params.py +15 -0
  89. mixedbread/types/vector_stores/file_list_response.py +30 -0
  90. mixedbread/types/vector_stores/vector_store_file.py +38 -0
  91. mixedbread-0.1.0a1.dist-info/METADATA +367 -0
  92. mixedbread-0.1.0a1.dist-info/RECORD +94 -0
  93. mixedbread-0.1.0a1.dist-info/WHEEL +4 -0
  94. mixedbread-0.1.0a1.dist-info/licenses/LICENSE +201 -0
mixedbread/_client.py ADDED
@@ -0,0 +1,867 @@
1
+ # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
2
+
3
+ from __future__ import annotations
4
+
5
+ import os
6
+ from typing import Any, Dict, List, Union, Mapping, Optional, cast
7
+ from typing_extensions import Self, Literal, override
8
+
9
+ import httpx
10
+
11
+ from . import _exceptions
12
+ from ._qs import Querystring
13
+ from .types import client_embed_params, client_rerank_params
14
+ from ._types import (
15
+ NOT_GIVEN,
16
+ Body,
17
+ Omit,
18
+ Query,
19
+ Headers,
20
+ Timeout,
21
+ NotGiven,
22
+ Transport,
23
+ ProxiesTypes,
24
+ RequestOptions,
25
+ )
26
+ from ._utils import (
27
+ is_given,
28
+ maybe_transform,
29
+ get_async_library,
30
+ async_maybe_transform,
31
+ )
32
+ from ._version import __version__
33
+ from ._response import (
34
+ to_raw_response_wrapper,
35
+ to_streamed_response_wrapper,
36
+ async_to_raw_response_wrapper,
37
+ async_to_streamed_response_wrapper,
38
+ )
39
+ from .resources import files, embeddings, rerankings
40
+ from ._streaming import Stream as Stream, AsyncStream as AsyncStream
41
+ from ._exceptions import APIStatusError, MixedbreadError
42
+ from ._base_client import (
43
+ DEFAULT_MAX_RETRIES,
44
+ SyncAPIClient,
45
+ AsyncAPIClient,
46
+ make_request_options,
47
+ )
48
+ from .resources.chat import chat
49
+ from .types.info_response import InfoResponse
50
+ from .types.embed_response import EmbedResponse
51
+ from .resources.document_ai import document_ai
52
+ from .types.rerank_response import RerankResponse
53
+ from .resources.vector_stores import vector_stores
54
+
55
+ __all__ = [
56
+ "ENVIRONMENTS",
57
+ "Timeout",
58
+ "Transport",
59
+ "ProxiesTypes",
60
+ "RequestOptions",
61
+ "Mixedbread",
62
+ "AsyncMixedbread",
63
+ "Client",
64
+ "AsyncClient",
65
+ ]
66
+
67
+ ENVIRONMENTS: Dict[str, str] = {
68
+ "production": "https://api.mixedbread.ai",
69
+ "environment_1": "http://127.0.0.1:8000",
70
+ }
71
+
72
+
73
+ class Mixedbread(SyncAPIClient):
74
+ document_ai: document_ai.DocumentAIResource
75
+ embeddings: embeddings.EmbeddingsResource
76
+ rerankings: rerankings.RerankingsResource
77
+ files: files.FilesResource
78
+ vector_stores: vector_stores.VectorStoresResource
79
+ chat: chat.ChatResource
80
+ with_raw_response: MixedbreadWithRawResponse
81
+ with_streaming_response: MixedbreadWithStreamedResponse
82
+
83
+ # client options
84
+ api_key: str
85
+
86
+ _environment: Literal["production", "environment_1"] | NotGiven
87
+
88
+ def __init__(
89
+ self,
90
+ *,
91
+ api_key: str | None = None,
92
+ environment: Literal["production", "environment_1"] | NotGiven = NOT_GIVEN,
93
+ base_url: str | httpx.URL | None | NotGiven = NOT_GIVEN,
94
+ timeout: Union[float, Timeout, None, NotGiven] = NOT_GIVEN,
95
+ max_retries: int = DEFAULT_MAX_RETRIES,
96
+ default_headers: Mapping[str, str] | None = None,
97
+ default_query: Mapping[str, object] | None = None,
98
+ # Configure a custom httpx client.
99
+ # We provide a `DefaultHttpxClient` class that you can pass to retain the default values we use for `limits`, `timeout` & `follow_redirects`.
100
+ # See the [httpx documentation](https://www.python-httpx.org/api/#client) for more details.
101
+ http_client: httpx.Client | None = None,
102
+ # Enable or disable schema validation for data returned by the API.
103
+ # When enabled an error APIResponseValidationError is raised
104
+ # if the API responds with invalid data for the expected schema.
105
+ #
106
+ # This parameter may be removed or changed in the future.
107
+ # If you rely on this feature, please open a GitHub issue
108
+ # outlining your use-case to help us decide if it should be
109
+ # part of our public interface in the future.
110
+ _strict_response_validation: bool = False,
111
+ ) -> None:
112
+ """Construct a new synchronous mixedbread client instance.
113
+
114
+ This automatically infers the `api_key` argument from the `MXBAI_API_KEY` environment variable if it is not provided.
115
+ """
116
+ if api_key is None:
117
+ api_key = os.environ.get("MXBAI_API_KEY")
118
+ if api_key is None:
119
+ raise MixedbreadError(
120
+ "The api_key client option must be set either by passing api_key to the client or by setting the MXBAI_API_KEY environment variable"
121
+ )
122
+ self.api_key = api_key
123
+
124
+ self._environment = environment
125
+
126
+ base_url_env = os.environ.get("MIXEDBREAD_BASE_URL")
127
+ if is_given(base_url) and base_url is not None:
128
+ # cast required because mypy doesn't understand the type narrowing
129
+ base_url = cast("str | httpx.URL", base_url) # pyright: ignore[reportUnnecessaryCast]
130
+ elif is_given(environment):
131
+ if base_url_env and base_url is not None:
132
+ raise ValueError(
133
+ "Ambiguous URL; The `MIXEDBREAD_BASE_URL` env var and the `environment` argument are given. If you want to use the environment, you must pass base_url=None",
134
+ )
135
+
136
+ try:
137
+ base_url = ENVIRONMENTS[environment]
138
+ except KeyError as exc:
139
+ raise ValueError(f"Unknown environment: {environment}") from exc
140
+ elif base_url_env is not None:
141
+ base_url = base_url_env
142
+ else:
143
+ self._environment = environment = "production"
144
+
145
+ try:
146
+ base_url = ENVIRONMENTS[environment]
147
+ except KeyError as exc:
148
+ raise ValueError(f"Unknown environment: {environment}") from exc
149
+
150
+ super().__init__(
151
+ version=__version__,
152
+ base_url=base_url,
153
+ max_retries=max_retries,
154
+ timeout=timeout,
155
+ http_client=http_client,
156
+ custom_headers=default_headers,
157
+ custom_query=default_query,
158
+ _strict_response_validation=_strict_response_validation,
159
+ )
160
+
161
+ self.document_ai = document_ai.DocumentAIResource(self)
162
+ self.embeddings = embeddings.EmbeddingsResource(self)
163
+ self.rerankings = rerankings.RerankingsResource(self)
164
+ self.files = files.FilesResource(self)
165
+ self.vector_stores = vector_stores.VectorStoresResource(self)
166
+ self.chat = chat.ChatResource(self)
167
+ self.with_raw_response = MixedbreadWithRawResponse(self)
168
+ self.with_streaming_response = MixedbreadWithStreamedResponse(self)
169
+
170
+ @property
171
+ @override
172
+ def qs(self) -> Querystring:
173
+ return Querystring(array_format="comma")
174
+
175
+ @property
176
+ @override
177
+ def auth_headers(self) -> dict[str, str]:
178
+ api_key = self.api_key
179
+ return {"Authorization": api_key}
180
+
181
+ @property
182
+ @override
183
+ def default_headers(self) -> dict[str, str | Omit]:
184
+ return {
185
+ **super().default_headers,
186
+ "X-Stainless-Async": "false",
187
+ **self._custom_headers,
188
+ }
189
+
190
+ def copy(
191
+ self,
192
+ *,
193
+ api_key: str | None = None,
194
+ environment: Literal["production", "environment_1"] | None = None,
195
+ base_url: str | httpx.URL | None = None,
196
+ timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
197
+ http_client: httpx.Client | None = None,
198
+ max_retries: int | NotGiven = NOT_GIVEN,
199
+ default_headers: Mapping[str, str] | None = None,
200
+ set_default_headers: Mapping[str, str] | None = None,
201
+ default_query: Mapping[str, object] | None = None,
202
+ set_default_query: Mapping[str, object] | None = None,
203
+ _extra_kwargs: Mapping[str, Any] = {},
204
+ ) -> Self:
205
+ """
206
+ Create a new client instance re-using the same options given to the current client with optional overriding.
207
+ """
208
+ if default_headers is not None and set_default_headers is not None:
209
+ raise ValueError("The `default_headers` and `set_default_headers` arguments are mutually exclusive")
210
+
211
+ if default_query is not None and set_default_query is not None:
212
+ raise ValueError("The `default_query` and `set_default_query` arguments are mutually exclusive")
213
+
214
+ headers = self._custom_headers
215
+ if default_headers is not None:
216
+ headers = {**headers, **default_headers}
217
+ elif set_default_headers is not None:
218
+ headers = set_default_headers
219
+
220
+ params = self._custom_query
221
+ if default_query is not None:
222
+ params = {**params, **default_query}
223
+ elif set_default_query is not None:
224
+ params = set_default_query
225
+
226
+ http_client = http_client or self._client
227
+ return self.__class__(
228
+ api_key=api_key or self.api_key,
229
+ base_url=base_url or self.base_url,
230
+ environment=environment or self._environment,
231
+ timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
232
+ http_client=http_client,
233
+ max_retries=max_retries if is_given(max_retries) else self.max_retries,
234
+ default_headers=headers,
235
+ default_query=params,
236
+ **_extra_kwargs,
237
+ )
238
+
239
+ # Alias for `copy` for nicer inline usage, e.g.
240
+ # client.with_options(timeout=10).foo.create(...)
241
+ with_options = copy
242
+
243
+ def embed(
244
+ self,
245
+ *,
246
+ input: client_embed_params.Input,
247
+ model: str,
248
+ dimensions: Optional[int] | NotGiven = NOT_GIVEN,
249
+ encoding_format: Union[
250
+ Literal["float", "float16", "base64", "binary", "ubinary", "int8", "uint8"],
251
+ List[Literal["float", "float16", "base64", "binary", "ubinary", "int8", "uint8"]],
252
+ ]
253
+ | NotGiven = NOT_GIVEN,
254
+ normalized: bool | NotGiven = NOT_GIVEN,
255
+ prompt: Optional[str] | NotGiven = NOT_GIVEN,
256
+ # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
257
+ # The extra values given here take precedence over values defined on the client or passed to this method.
258
+ extra_headers: Headers | None = None,
259
+ extra_query: Query | None = None,
260
+ extra_body: Body | None = None,
261
+ timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
262
+ ) -> EmbedResponse:
263
+ """
264
+ Create embeddings for text or images using the specified model, encoding format,
265
+ and normalization.
266
+
267
+ Args: params: The parameters for creating embeddings.
268
+
269
+ Returns: EmbeddingCreateResponse: The response containing the embeddings.
270
+
271
+ Args:
272
+ input: The input to create embeddings for.
273
+
274
+ model: The model to use for creating embeddings.
275
+
276
+ dimensions: The number of dimensions to use for the embeddings.
277
+
278
+ encoding_format: The encoding format of the embeddings.
279
+
280
+ normalized: Whether to normalize the embeddings.
281
+
282
+ prompt: The prompt to use for the embedding creation.
283
+
284
+ extra_headers: Send extra headers
285
+
286
+ extra_query: Add additional query parameters to the request
287
+
288
+ extra_body: Add additional JSON properties to the request
289
+
290
+ timeout: Override the client-level default timeout for this request, in seconds
291
+ """
292
+ return self.post(
293
+ "/v1/embeddings",
294
+ body=maybe_transform(
295
+ {
296
+ "input": input,
297
+ "model": model,
298
+ "dimensions": dimensions,
299
+ "encoding_format": encoding_format,
300
+ "normalized": normalized,
301
+ "prompt": prompt,
302
+ },
303
+ client_embed_params.ClientEmbedParams,
304
+ ),
305
+ options=make_request_options(
306
+ extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
307
+ ),
308
+ cast_to=EmbedResponse,
309
+ )
310
+
311
+ def info(
312
+ self,
313
+ *,
314
+ # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
315
+ # The extra values given here take precedence over values defined on the client or passed to this method.
316
+ extra_headers: Headers | None = None,
317
+ extra_query: Query | None = None,
318
+ extra_body: Body | None = None,
319
+ timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
320
+ ) -> InfoResponse:
321
+ """
322
+ Returns service information, including name and version.
323
+
324
+ Returns: InfoResponse: A response containing the service name and version.
325
+ """
326
+ return self.get(
327
+ "/",
328
+ options=make_request_options(
329
+ extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
330
+ ),
331
+ cast_to=InfoResponse,
332
+ )
333
+
334
+ def rerank(
335
+ self,
336
+ *,
337
+ input: object,
338
+ query: str,
339
+ model: str | NotGiven = NOT_GIVEN,
340
+ rank_fields: Optional[List[str]] | NotGiven = NOT_GIVEN,
341
+ return_input: bool | NotGiven = NOT_GIVEN,
342
+ top_k: int | NotGiven = NOT_GIVEN,
343
+ # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
344
+ # The extra values given here take precedence over values defined on the client or passed to this method.
345
+ extra_headers: Headers | None = None,
346
+ extra_query: Query | None = None,
347
+ extra_body: Body | None = None,
348
+ timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
349
+ ) -> RerankResponse:
350
+ """
351
+ Rerank different kind of documents for a given query.
352
+
353
+ Args: params: RerankingCreateParams: The parameters for reranking.
354
+
355
+ Returns: RerankingCreateResponse: The reranked documents for the input query.
356
+
357
+ Args:
358
+ query: The query to rerank the documents.
359
+
360
+ model: The model to use for reranking documents.
361
+
362
+ rank_fields: The fields of the documents to rank.
363
+
364
+ return_input: Whether to return the documents.
365
+
366
+ top_k: The number of documents to return.
367
+
368
+ extra_headers: Send extra headers
369
+
370
+ extra_query: Add additional query parameters to the request
371
+
372
+ extra_body: Add additional JSON properties to the request
373
+
374
+ timeout: Override the client-level default timeout for this request, in seconds
375
+ """
376
+ return self.post(
377
+ "/v1/reranking",
378
+ body=maybe_transform(
379
+ {
380
+ "input": input,
381
+ "query": query,
382
+ "model": model,
383
+ "rank_fields": rank_fields,
384
+ "return_input": return_input,
385
+ "top_k": top_k,
386
+ },
387
+ client_rerank_params.ClientRerankParams,
388
+ ),
389
+ options=make_request_options(
390
+ extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
391
+ ),
392
+ cast_to=RerankResponse,
393
+ )
394
+
395
+ @override
396
+ def _make_status_error(
397
+ self,
398
+ err_msg: str,
399
+ *,
400
+ body: object,
401
+ response: httpx.Response,
402
+ ) -> APIStatusError:
403
+ if response.status_code == 400:
404
+ return _exceptions.BadRequestError(err_msg, response=response, body=body)
405
+
406
+ if response.status_code == 401:
407
+ return _exceptions.AuthenticationError(err_msg, response=response, body=body)
408
+
409
+ if response.status_code == 403:
410
+ return _exceptions.PermissionDeniedError(err_msg, response=response, body=body)
411
+
412
+ if response.status_code == 404:
413
+ return _exceptions.NotFoundError(err_msg, response=response, body=body)
414
+
415
+ if response.status_code == 409:
416
+ return _exceptions.ConflictError(err_msg, response=response, body=body)
417
+
418
+ if response.status_code == 422:
419
+ return _exceptions.UnprocessableEntityError(err_msg, response=response, body=body)
420
+
421
+ if response.status_code == 429:
422
+ return _exceptions.RateLimitError(err_msg, response=response, body=body)
423
+
424
+ if response.status_code >= 500:
425
+ return _exceptions.InternalServerError(err_msg, response=response, body=body)
426
+ return APIStatusError(err_msg, response=response, body=body)
427
+
428
+
429
+ class AsyncMixedbread(AsyncAPIClient):
430
+ document_ai: document_ai.AsyncDocumentAIResource
431
+ embeddings: embeddings.AsyncEmbeddingsResource
432
+ rerankings: rerankings.AsyncRerankingsResource
433
+ files: files.AsyncFilesResource
434
+ vector_stores: vector_stores.AsyncVectorStoresResource
435
+ chat: chat.AsyncChatResource
436
+ with_raw_response: AsyncMixedbreadWithRawResponse
437
+ with_streaming_response: AsyncMixedbreadWithStreamedResponse
438
+
439
+ # client options
440
+ api_key: str
441
+
442
+ _environment: Literal["production", "environment_1"] | NotGiven
443
+
444
+ def __init__(
445
+ self,
446
+ *,
447
+ api_key: str | None = None,
448
+ environment: Literal["production", "environment_1"] | NotGiven = NOT_GIVEN,
449
+ base_url: str | httpx.URL | None | NotGiven = NOT_GIVEN,
450
+ timeout: Union[float, Timeout, None, NotGiven] = NOT_GIVEN,
451
+ max_retries: int = DEFAULT_MAX_RETRIES,
452
+ default_headers: Mapping[str, str] | None = None,
453
+ default_query: Mapping[str, object] | None = None,
454
+ # Configure a custom httpx client.
455
+ # We provide a `DefaultAsyncHttpxClient` class that you can pass to retain the default values we use for `limits`, `timeout` & `follow_redirects`.
456
+ # See the [httpx documentation](https://www.python-httpx.org/api/#asyncclient) for more details.
457
+ http_client: httpx.AsyncClient | None = None,
458
+ # Enable or disable schema validation for data returned by the API.
459
+ # When enabled an error APIResponseValidationError is raised
460
+ # if the API responds with invalid data for the expected schema.
461
+ #
462
+ # This parameter may be removed or changed in the future.
463
+ # If you rely on this feature, please open a GitHub issue
464
+ # outlining your use-case to help us decide if it should be
465
+ # part of our public interface in the future.
466
+ _strict_response_validation: bool = False,
467
+ ) -> None:
468
+ """Construct a new async mixedbread client instance.
469
+
470
+ This automatically infers the `api_key` argument from the `MXBAI_API_KEY` environment variable if it is not provided.
471
+ """
472
+ if api_key is None:
473
+ api_key = os.environ.get("MXBAI_API_KEY")
474
+ if api_key is None:
475
+ raise MixedbreadError(
476
+ "The api_key client option must be set either by passing api_key to the client or by setting the MXBAI_API_KEY environment variable"
477
+ )
478
+ self.api_key = api_key
479
+
480
+ self._environment = environment
481
+
482
+ base_url_env = os.environ.get("MIXEDBREAD_BASE_URL")
483
+ if is_given(base_url) and base_url is not None:
484
+ # cast required because mypy doesn't understand the type narrowing
485
+ base_url = cast("str | httpx.URL", base_url) # pyright: ignore[reportUnnecessaryCast]
486
+ elif is_given(environment):
487
+ if base_url_env and base_url is not None:
488
+ raise ValueError(
489
+ "Ambiguous URL; The `MIXEDBREAD_BASE_URL` env var and the `environment` argument are given. If you want to use the environment, you must pass base_url=None",
490
+ )
491
+
492
+ try:
493
+ base_url = ENVIRONMENTS[environment]
494
+ except KeyError as exc:
495
+ raise ValueError(f"Unknown environment: {environment}") from exc
496
+ elif base_url_env is not None:
497
+ base_url = base_url_env
498
+ else:
499
+ self._environment = environment = "production"
500
+
501
+ try:
502
+ base_url = ENVIRONMENTS[environment]
503
+ except KeyError as exc:
504
+ raise ValueError(f"Unknown environment: {environment}") from exc
505
+
506
+ super().__init__(
507
+ version=__version__,
508
+ base_url=base_url,
509
+ max_retries=max_retries,
510
+ timeout=timeout,
511
+ http_client=http_client,
512
+ custom_headers=default_headers,
513
+ custom_query=default_query,
514
+ _strict_response_validation=_strict_response_validation,
515
+ )
516
+
517
+ self.document_ai = document_ai.AsyncDocumentAIResource(self)
518
+ self.embeddings = embeddings.AsyncEmbeddingsResource(self)
519
+ self.rerankings = rerankings.AsyncRerankingsResource(self)
520
+ self.files = files.AsyncFilesResource(self)
521
+ self.vector_stores = vector_stores.AsyncVectorStoresResource(self)
522
+ self.chat = chat.AsyncChatResource(self)
523
+ self.with_raw_response = AsyncMixedbreadWithRawResponse(self)
524
+ self.with_streaming_response = AsyncMixedbreadWithStreamedResponse(self)
525
+
526
+ @property
527
+ @override
528
+ def qs(self) -> Querystring:
529
+ return Querystring(array_format="comma")
530
+
531
+ @property
532
+ @override
533
+ def auth_headers(self) -> dict[str, str]:
534
+ api_key = self.api_key
535
+ return {"Authorization": api_key}
536
+
537
+ @property
538
+ @override
539
+ def default_headers(self) -> dict[str, str | Omit]:
540
+ return {
541
+ **super().default_headers,
542
+ "X-Stainless-Async": f"async:{get_async_library()}",
543
+ **self._custom_headers,
544
+ }
545
+
546
+ def copy(
547
+ self,
548
+ *,
549
+ api_key: str | None = None,
550
+ environment: Literal["production", "environment_1"] | None = None,
551
+ base_url: str | httpx.URL | None = None,
552
+ timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
553
+ http_client: httpx.AsyncClient | None = None,
554
+ max_retries: int | NotGiven = NOT_GIVEN,
555
+ default_headers: Mapping[str, str] | None = None,
556
+ set_default_headers: Mapping[str, str] | None = None,
557
+ default_query: Mapping[str, object] | None = None,
558
+ set_default_query: Mapping[str, object] | None = None,
559
+ _extra_kwargs: Mapping[str, Any] = {},
560
+ ) -> Self:
561
+ """
562
+ Create a new client instance re-using the same options given to the current client with optional overriding.
563
+ """
564
+ if default_headers is not None and set_default_headers is not None:
565
+ raise ValueError("The `default_headers` and `set_default_headers` arguments are mutually exclusive")
566
+
567
+ if default_query is not None and set_default_query is not None:
568
+ raise ValueError("The `default_query` and `set_default_query` arguments are mutually exclusive")
569
+
570
+ headers = self._custom_headers
571
+ if default_headers is not None:
572
+ headers = {**headers, **default_headers}
573
+ elif set_default_headers is not None:
574
+ headers = set_default_headers
575
+
576
+ params = self._custom_query
577
+ if default_query is not None:
578
+ params = {**params, **default_query}
579
+ elif set_default_query is not None:
580
+ params = set_default_query
581
+
582
+ http_client = http_client or self._client
583
+ return self.__class__(
584
+ api_key=api_key or self.api_key,
585
+ base_url=base_url or self.base_url,
586
+ environment=environment or self._environment,
587
+ timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
588
+ http_client=http_client,
589
+ max_retries=max_retries if is_given(max_retries) else self.max_retries,
590
+ default_headers=headers,
591
+ default_query=params,
592
+ **_extra_kwargs,
593
+ )
594
+
595
+ # Alias for `copy` for nicer inline usage, e.g.
596
+ # client.with_options(timeout=10).foo.create(...)
597
+ with_options = copy
598
+
599
+ async def embed(
600
+ self,
601
+ *,
602
+ input: client_embed_params.Input,
603
+ model: str,
604
+ dimensions: Optional[int] | NotGiven = NOT_GIVEN,
605
+ encoding_format: Union[
606
+ Literal["float", "float16", "base64", "binary", "ubinary", "int8", "uint8"],
607
+ List[Literal["float", "float16", "base64", "binary", "ubinary", "int8", "uint8"]],
608
+ ]
609
+ | NotGiven = NOT_GIVEN,
610
+ normalized: bool | NotGiven = NOT_GIVEN,
611
+ prompt: Optional[str] | NotGiven = NOT_GIVEN,
612
+ # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
613
+ # The extra values given here take precedence over values defined on the client or passed to this method.
614
+ extra_headers: Headers | None = None,
615
+ extra_query: Query | None = None,
616
+ extra_body: Body | None = None,
617
+ timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
618
+ ) -> EmbedResponse:
619
+ """
620
+ Create embeddings for text or images using the specified model, encoding format,
621
+ and normalization.
622
+
623
+ Args: params: The parameters for creating embeddings.
624
+
625
+ Returns: EmbeddingCreateResponse: The response containing the embeddings.
626
+
627
+ Args:
628
+ input: The input to create embeddings for.
629
+
630
+ model: The model to use for creating embeddings.
631
+
632
+ dimensions: The number of dimensions to use for the embeddings.
633
+
634
+ encoding_format: The encoding format of the embeddings.
635
+
636
+ normalized: Whether to normalize the embeddings.
637
+
638
+ prompt: The prompt to use for the embedding creation.
639
+
640
+ extra_headers: Send extra headers
641
+
642
+ extra_query: Add additional query parameters to the request
643
+
644
+ extra_body: Add additional JSON properties to the request
645
+
646
+ timeout: Override the client-level default timeout for this request, in seconds
647
+ """
648
+ return await self.post(
649
+ "/v1/embeddings",
650
+ body=await async_maybe_transform(
651
+ {
652
+ "input": input,
653
+ "model": model,
654
+ "dimensions": dimensions,
655
+ "encoding_format": encoding_format,
656
+ "normalized": normalized,
657
+ "prompt": prompt,
658
+ },
659
+ client_embed_params.ClientEmbedParams,
660
+ ),
661
+ options=make_request_options(
662
+ extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
663
+ ),
664
+ cast_to=EmbedResponse,
665
+ )
666
+
667
+ async def info(
668
+ self,
669
+ *,
670
+ # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
671
+ # The extra values given here take precedence over values defined on the client or passed to this method.
672
+ extra_headers: Headers | None = None,
673
+ extra_query: Query | None = None,
674
+ extra_body: Body | None = None,
675
+ timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
676
+ ) -> InfoResponse:
677
+ """
678
+ Returns service information, including name and version.
679
+
680
+ Returns: InfoResponse: A response containing the service name and version.
681
+ """
682
+ return await self.get(
683
+ "/",
684
+ options=make_request_options(
685
+ extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
686
+ ),
687
+ cast_to=InfoResponse,
688
+ )
689
+
690
+ async def rerank(
691
+ self,
692
+ *,
693
+ input: object,
694
+ query: str,
695
+ model: str | NotGiven = NOT_GIVEN,
696
+ rank_fields: Optional[List[str]] | NotGiven = NOT_GIVEN,
697
+ return_input: bool | NotGiven = NOT_GIVEN,
698
+ top_k: int | NotGiven = NOT_GIVEN,
699
+ # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
700
+ # The extra values given here take precedence over values defined on the client or passed to this method.
701
+ extra_headers: Headers | None = None,
702
+ extra_query: Query | None = None,
703
+ extra_body: Body | None = None,
704
+ timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
705
+ ) -> RerankResponse:
706
+ """
707
+ Rerank different kind of documents for a given query.
708
+
709
+ Args: params: RerankingCreateParams: The parameters for reranking.
710
+
711
+ Returns: RerankingCreateResponse: The reranked documents for the input query.
712
+
713
+ Args:
714
+ query: The query to rerank the documents.
715
+
716
+ model: The model to use for reranking documents.
717
+
718
+ rank_fields: The fields of the documents to rank.
719
+
720
+ return_input: Whether to return the documents.
721
+
722
+ top_k: The number of documents to return.
723
+
724
+ extra_headers: Send extra headers
725
+
726
+ extra_query: Add additional query parameters to the request
727
+
728
+ extra_body: Add additional JSON properties to the request
729
+
730
+ timeout: Override the client-level default timeout for this request, in seconds
731
+ """
732
+ return await self.post(
733
+ "/v1/reranking",
734
+ body=await async_maybe_transform(
735
+ {
736
+ "input": input,
737
+ "query": query,
738
+ "model": model,
739
+ "rank_fields": rank_fields,
740
+ "return_input": return_input,
741
+ "top_k": top_k,
742
+ },
743
+ client_rerank_params.ClientRerankParams,
744
+ ),
745
+ options=make_request_options(
746
+ extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
747
+ ),
748
+ cast_to=RerankResponse,
749
+ )
750
+
751
+ @override
752
+ def _make_status_error(
753
+ self,
754
+ err_msg: str,
755
+ *,
756
+ body: object,
757
+ response: httpx.Response,
758
+ ) -> APIStatusError:
759
+ if response.status_code == 400:
760
+ return _exceptions.BadRequestError(err_msg, response=response, body=body)
761
+
762
+ if response.status_code == 401:
763
+ return _exceptions.AuthenticationError(err_msg, response=response, body=body)
764
+
765
+ if response.status_code == 403:
766
+ return _exceptions.PermissionDeniedError(err_msg, response=response, body=body)
767
+
768
+ if response.status_code == 404:
769
+ return _exceptions.NotFoundError(err_msg, response=response, body=body)
770
+
771
+ if response.status_code == 409:
772
+ return _exceptions.ConflictError(err_msg, response=response, body=body)
773
+
774
+ if response.status_code == 422:
775
+ return _exceptions.UnprocessableEntityError(err_msg, response=response, body=body)
776
+
777
+ if response.status_code == 429:
778
+ return _exceptions.RateLimitError(err_msg, response=response, body=body)
779
+
780
+ if response.status_code >= 500:
781
+ return _exceptions.InternalServerError(err_msg, response=response, body=body)
782
+ return APIStatusError(err_msg, response=response, body=body)
783
+
784
+
785
+ class MixedbreadWithRawResponse:
786
+ def __init__(self, client: Mixedbread) -> None:
787
+ self.document_ai = document_ai.DocumentAIResourceWithRawResponse(client.document_ai)
788
+ self.embeddings = embeddings.EmbeddingsResourceWithRawResponse(client.embeddings)
789
+ self.rerankings = rerankings.RerankingsResourceWithRawResponse(client.rerankings)
790
+ self.files = files.FilesResourceWithRawResponse(client.files)
791
+ self.vector_stores = vector_stores.VectorStoresResourceWithRawResponse(client.vector_stores)
792
+ self.chat = chat.ChatResourceWithRawResponse(client.chat)
793
+
794
+ self.embed = to_raw_response_wrapper(
795
+ client.embed,
796
+ )
797
+ self.info = to_raw_response_wrapper(
798
+ client.info,
799
+ )
800
+ self.rerank = to_raw_response_wrapper(
801
+ client.rerank,
802
+ )
803
+
804
+
805
+ class AsyncMixedbreadWithRawResponse:
806
+ def __init__(self, client: AsyncMixedbread) -> None:
807
+ self.document_ai = document_ai.AsyncDocumentAIResourceWithRawResponse(client.document_ai)
808
+ self.embeddings = embeddings.AsyncEmbeddingsResourceWithRawResponse(client.embeddings)
809
+ self.rerankings = rerankings.AsyncRerankingsResourceWithRawResponse(client.rerankings)
810
+ self.files = files.AsyncFilesResourceWithRawResponse(client.files)
811
+ self.vector_stores = vector_stores.AsyncVectorStoresResourceWithRawResponse(client.vector_stores)
812
+ self.chat = chat.AsyncChatResourceWithRawResponse(client.chat)
813
+
814
+ self.embed = async_to_raw_response_wrapper(
815
+ client.embed,
816
+ )
817
+ self.info = async_to_raw_response_wrapper(
818
+ client.info,
819
+ )
820
+ self.rerank = async_to_raw_response_wrapper(
821
+ client.rerank,
822
+ )
823
+
824
+
825
+ class MixedbreadWithStreamedResponse:
826
+ def __init__(self, client: Mixedbread) -> None:
827
+ self.document_ai = document_ai.DocumentAIResourceWithStreamingResponse(client.document_ai)
828
+ self.embeddings = embeddings.EmbeddingsResourceWithStreamingResponse(client.embeddings)
829
+ self.rerankings = rerankings.RerankingsResourceWithStreamingResponse(client.rerankings)
830
+ self.files = files.FilesResourceWithStreamingResponse(client.files)
831
+ self.vector_stores = vector_stores.VectorStoresResourceWithStreamingResponse(client.vector_stores)
832
+ self.chat = chat.ChatResourceWithStreamingResponse(client.chat)
833
+
834
+ self.embed = to_streamed_response_wrapper(
835
+ client.embed,
836
+ )
837
+ self.info = to_streamed_response_wrapper(
838
+ client.info,
839
+ )
840
+ self.rerank = to_streamed_response_wrapper(
841
+ client.rerank,
842
+ )
843
+
844
+
845
+ class AsyncMixedbreadWithStreamedResponse:
846
+ def __init__(self, client: AsyncMixedbread) -> None:
847
+ self.document_ai = document_ai.AsyncDocumentAIResourceWithStreamingResponse(client.document_ai)
848
+ self.embeddings = embeddings.AsyncEmbeddingsResourceWithStreamingResponse(client.embeddings)
849
+ self.rerankings = rerankings.AsyncRerankingsResourceWithStreamingResponse(client.rerankings)
850
+ self.files = files.AsyncFilesResourceWithStreamingResponse(client.files)
851
+ self.vector_stores = vector_stores.AsyncVectorStoresResourceWithStreamingResponse(client.vector_stores)
852
+ self.chat = chat.AsyncChatResourceWithStreamingResponse(client.chat)
853
+
854
+ self.embed = async_to_streamed_response_wrapper(
855
+ client.embed,
856
+ )
857
+ self.info = async_to_streamed_response_wrapper(
858
+ client.info,
859
+ )
860
+ self.rerank = async_to_streamed_response_wrapper(
861
+ client.rerank,
862
+ )
863
+
864
+
865
+ Client = Mixedbread
866
+
867
+ AsyncClient = AsyncMixedbread