Merge pull request #3 from mixedbread-ai/feature/obj-reranking

Add object reranking

mixedbread_ai/__init__.py

@@ -10,12 +10,12 @@
     EmbeddingsResponseEncodingFormat,
     EncodingFormat,
     ForbiddenErrorBody,
-    Input,
     InternalError,
     InvalidMatryoshkaDimensionsError,
     InvalidMatryoshkaModelError,
     ModelNotFoundError,
     MultiModalInput,
+    MultiModalRerankingInput,
     MultipleEncodingsEmbedding,
     MultipleEncodingsEmbeddingItem,
     MxbaiApiError,
@@ -56,14 +56,14 @@
     "EncodingFormat",
     "ForbiddenError",
     "ForbiddenErrorBody",
-    "Input",
     "InternalError",
     "InternalServerError",
     "InvalidMatryoshkaDimensionsError",
     "InvalidMatryoshkaModelError",
     "MixedbreadAIEnvironment",
     "ModelNotFoundError",
     "MultiModalInput",
+    "MultiModalRerankingInput",
     "MultipleEncodingsEmbedding",
     "MultipleEncodingsEmbeddingItem",
     "MxbaiApiError",

mixedbread_ai/client.py

@@ -23,9 +23,9 @@
 from .types.embeddings_request_encoding_format import EmbeddingsRequestEncodingFormat
 from .types.embeddings_response import EmbeddingsResponse
 from .types.forbidden_error_body import ForbiddenErrorBody
-from .types.input import Input
 from .types.internal_error import InternalError
 from .types.multi_modal_input import MultiModalInput
+from .types.multi_modal_reranking_input import MultiModalRerankingInput
 from .types.not_found_error_body import NotFoundErrorBody
 from .types.query import Query
 from .types.reranking_response import RerankingResponse
@@ -192,20 +192,23 @@ def reranking(
         self,
         *,
         model: str,
-        input: Input,
         query: Query,
+        input: MultiModalRerankingInput,
+        rank_fields: typing.Optional[typing.Sequence[str]] = OMIT,
         top_k: typing.Optional[int] = OMIT,
         return_input: typing.Optional[bool] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> RerankingResponse:
         """
         Parameters:
-            - model: str. The model to use for creating embeddings
-
-            - input: Input. The input documents to rerank
+            - model: str. The model to use for reranking documents
 
             - query: Query. The query to rerank the documents
 
+            - input: MultiModalRerankingInput.
+
+            - rank_fields: typing.Optional[typing.Sequence[str]].
+
             - top_k: typing.Optional[int]. The number of documents to return
 
             - return_input: typing.Optional[bool]. Whether to return the documents
@@ -216,9 +219,11 @@ def reranking(
         from mixedbread-ai.client import MixedbreadAI
 
         client = MixedbreadAI(api_key="YOUR_API_KEY", )
-        client.reranking(model="model", input=["input"], query=TextDocument(text="text", ), top_k=10, return_input=False, )
+        client.reranking(model="model", query=TextDocument(text="text", ), input=["input"], top_k=10, return_input=False, )
         """
-        _request: typing.Dict[str, typing.Any] = {"model": model, "input": input, "query": query}
+        _request: typing.Dict[str, typing.Any] = {"model": model, "query": query, "input": input}
+        if rank_fields is not OMIT:
+            _request["rank_fields"] = rank_fields
         if top_k is not OMIT:
             _request["top_k"] = top_k
         if return_input is not OMIT:
@@ -423,20 +428,23 @@ async def reranking(
         self,
         *,
         model: str,
-        input: Input,
         query: Query,
+        input: MultiModalRerankingInput,
+        rank_fields: typing.Optional[typing.Sequence[str]] = OMIT,
         top_k: typing.Optional[int] = OMIT,
         return_input: typing.Optional[bool] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> RerankingResponse:
         """
         Parameters:
-            - model: str. The model to use for creating embeddings
-
-            - input: Input. The input documents to rerank
+            - model: str. The model to use for reranking documents
 
             - query: Query. The query to rerank the documents
 
+            - input: MultiModalRerankingInput.
+
+            - rank_fields: typing.Optional[typing.Sequence[str]].
+
             - top_k: typing.Optional[int]. The number of documents to return
 
             - return_input: typing.Optional[bool]. Whether to return the documents
@@ -447,9 +455,11 @@ async def reranking(
         from mixedbread-ai.client import AsyncMixedbreadAI
 
         client = AsyncMixedbreadAI(api_key="YOUR_API_KEY", )
-        await client.reranking(model="model", input=["input"], query=TextDocument(text="text", ), top_k=10, return_input=False, )
+        await client.reranking(model="model", query=TextDocument(text="text", ), input=["input"], top_k=10, return_input=False, )
         """
-        _request: typing.Dict[str, typing.Any] = {"model": model, "input": input, "query": query}
+        _request: typing.Dict[str, typing.Any] = {"model": model, "query": query, "input": input}
+        if rank_fields is not OMIT:
+            _request["rank_fields"] = rank_fields
         if top_k is not OMIT:
             _request["top_k"] = top_k
         if return_input is not OMIT:

mixedbread_ai/types/__init__.py

@@ -9,12 +9,12 @@
 from .embeddings_response_encoding_format import EmbeddingsResponseEncodingFormat
 from .encoding_format import EncodingFormat
 from .forbidden_error_body import ForbiddenErrorBody
-from .input import Input
 from .internal_error import InternalError
 from .invalid_matryoshka_dimensions_error import InvalidMatryoshkaDimensionsError
 from .invalid_matryoshka_model_error import InvalidMatryoshkaModelError
 from .model_not_found_error import ModelNotFoundError
 from .multi_modal_input import MultiModalInput
+from .multi_modal_reranking_input import MultiModalRerankingInput
 from .multiple_encodings_embedding import MultipleEncodingsEmbedding
 from .multiple_encodings_embedding_item import MultipleEncodingsEmbeddingItem
 from .mxbai_api_error import MxbaiApiError
@@ -42,12 +42,12 @@
     "EmbeddingsResponseEncodingFormat",
     "EncodingFormat",
     "ForbiddenErrorBody",
-    "Input",
     "InternalError",
     "InvalidMatryoshkaDimensionsError",
     "InvalidMatryoshkaModelError",
     "ModelNotFoundError",
     "MultiModalInput",
+    "MultiModalRerankingInput",
     "MultipleEncodingsEmbedding",
     "MultipleEncodingsEmbeddingItem",
     "MxbaiApiError",

mixedbread_ai/types/multi_modal_reranking_input.py

@@ -0,0 +1,5 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+MultiModalRerankingInput = typing.Union[typing.List[str], typing.List[typing.Dict[str, typing.Any]]]