| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import TypeAlias, cast |
| 4 | |
| 5 | from fastapi.responses import JSONResponse, Response, StreamingResponse |
| 6 | from typing_extensions import assert_never |
| 7 | |
| 8 | from vllm.logger import init_logger |
| 9 | from vllm.outputs import PoolingRequestOutput |
| 10 | from vllm.utils.serial_utils import EmbedDType, Endianness |
| 11 | |
| 12 | from ..base.io_processor import PoolingIOProcessor |
| 13 | from ..base.serving import PoolingServing |
| 14 | from ..typing import PoolingServeContext |
| 15 | from ..utils import ( |
| 16 | BytesEncodingFormat, |
| 17 | JsonEncodingFormat, |
| 18 | build_pooling_bytes_streaming_response, |
| 19 | encode_pooling_output_float, |
| 20 | encode_pooling_output_float_or_ndarray, |
| 21 | get_json_response_cls, |
| 22 | get_pooling_output_encoder, |
| 23 | get_pooling_usage, |
| 24 | ) |
| 25 | from .io_processor import EmbedIOProcessor |
| 26 | from .protocol import ( |
| 27 | CohereBilledUnits, |
| 28 | CohereEmbedRequest, |
| 29 | CohereEmbedResponse, |
| 30 | CohereMeta, |
| 31 | EmbeddingRequest, |
| 32 | EmbeddingResponse, |
| 33 | EmbeddingResponseData, |
| 34 | build_typed_embeddings, |
| 35 | ) |
| 36 | |
| 37 | logger = init_logger(__name__) |
| 38 | |
| 39 | |
| 40 | EmbeddingServeContext: TypeAlias = PoolingServeContext[EmbeddingRequest] |
| 41 | |
| 42 | |
| 43 | class ServingEmbedding(PoolingServing): |
| 44 | """Embedding API supporting both OpenAI and Cohere formats.""" |
| 45 | |
| 46 | request_id_prefix = "embd" |
| 47 | io_processor: EmbedIOProcessor |
| 48 | |
| 49 | def __init__(self, *args, **kwargs): |
| 50 | super().__init__(*args, **kwargs) |
| 51 | |
| 52 | self.json_response_cls = get_json_response_cls() |
| 53 | |
| 54 | def init_io_processor(self, *args, **kwargs) -> EmbedIOProcessor: |
| 55 | return EmbedIOProcessor(*args, **kwargs) |
| 56 | |
| 57 | async def _preprocessing( |
| 58 | self, io_processor: PoolingIOProcessor, ctx: PoolingServeContext |
| 59 | ): |
| 60 | await super()._preprocessing(io_processor, ctx) |
| 61 | self.io_processor.maybe_pre_process_chunked(ctx) |
| 62 | |
| 63 | def _build_response( |
| 64 | self, |
| 65 | ctx: PoolingServeContext, |
| 66 | ) -> Response: |
| 67 | if isinstance(ctx.request, CohereEmbedRequest): |
| 68 | return self._build_cohere_response_from_ctx(ctx) |
| 69 | return self._build_openai_response(ctx) |
| 70 | |
| 71 | def _build_openai_response( |
| 72 | self, |
| 73 | ctx: EmbeddingServeContext, |
| 74 | ) -> JSONResponse | StreamingResponse: |
| 75 | encoding_format = ctx.request.encoding_format |
| 76 | embed_dtype = ctx.request.embed_dtype |
| 77 | endianness = ctx.request.endianness |
| 78 | |
| 79 | if encoding_format == "float" or encoding_format == "base64": |
| 80 | return self._openai_json_response( |
| 81 | ctx.final_res_batch, |
| 82 | ctx.request_id, |
| 83 | ctx.created_time, |
| 84 | ctx.model_name, |
| 85 | encoding_format, |
| 86 | embed_dtype, |
| 87 | endianness, |
| 88 | ) |
| 89 | |
| 90 | if encoding_format == "bytes" or encoding_format == "bytes_only": |
| 91 | return self._openai_bytes_response( |
| 92 | ctx.final_res_batch, |
| 93 | ctx.request_id, |
| 94 | ctx.created_time, |
| 95 | ctx.model_name, |
| 96 | encoding_format, |
| 97 | embed_dtype, |
| 98 | endianness, |
| 99 | ) |
| 100 | |
| 101 | assert_never(encoding_format) |
| 102 | |
| 103 | def _openai_json_response( |
| 104 | self, |
| 105 | final_res_batch: list[PoolingRequestOutput], |
| 106 | request_id: str, |
| 107 | created_time: int, |
| 108 | model_name: str, |
| 109 | encoding_format: JsonEncodingFormat, |
| 110 | embed_dtype: EmbedDType, |
| 111 | endianness: Endianness, |
| 112 | ) -> JSONResponse: |
| 113 | use_ndarray_response = ( |
| 114 | encoding_format == "float" |
| 115 | and self.json_response_cls.__name__ == "ORJSONResponse" |
| 116 | ) |
| 117 | if use_ndarray_response: |
| 118 | ndarray_items: list[dict[str, object]] = [] |
| 119 | |
| 120 | for idx, final_res in enumerate(final_res_batch): |
| 121 | item_dict = EmbeddingResponseData( |
| 122 | index=idx, |
| 123 | embedding=[], |
| 124 | ).model_dump() |
| 125 | item_dict["embedding"] = encode_pooling_output_float_or_ndarray( |
| 126 | final_res |
| 127 | ) |
| 128 | ndarray_items.append(item_dict) |
| 129 | ndarray_response = EmbeddingResponse( |
| 130 | id=request_id, |
| 131 | created=created_time, |
| 132 | model=model_name, |
| 133 | data=[], # type: ignore[arg-type] |
| 134 | usage=get_pooling_usage(final_res_batch), |
| 135 | ).model_dump() |
| 136 | ndarray_response["data"] = ndarray_items |
| 137 | |
| 138 | return self.json_response_cls(content=ndarray_response) |
| 139 | |
| 140 | encode_fn = get_pooling_output_encoder( |
| 141 | encoding_format=encoding_format, |
| 142 | embed_dtype=embed_dtype, |
| 143 | endianness=endianness, |
| 144 | ) |
| 145 | |
| 146 | items: list[EmbeddingResponseData] = [] |
| 147 | |
| 148 | for idx, final_res in enumerate(final_res_batch): |
| 149 | item = EmbeddingResponseData( |
| 150 | index=idx, |
| 151 | embedding=cast(list[float] | str, encode_fn(final_res)), |
| 152 | ) |
| 153 | |
| 154 | items.append(item) |
| 155 | |
| 156 | response = EmbeddingResponse( |
| 157 | id=request_id, |
| 158 | created=created_time, |
| 159 | model=model_name, |
| 160 | data=items, |
| 161 | usage=get_pooling_usage(final_res_batch), |
| 162 | ) |
| 163 | return self.json_response_cls(content=response.model_dump()) |
| 164 | |
| 165 | def _openai_bytes_response( |
| 166 | self, |
| 167 | final_res_batch: list[PoolingRequestOutput], |
| 168 | request_id: str, |
| 169 | created_time: int, |
| 170 | model_name: str, |
| 171 | encoding_format: BytesEncodingFormat, |
| 172 | embed_dtype: EmbedDType, |
| 173 | endianness: Endianness, |
| 174 | ) -> StreamingResponse: |
| 175 | return build_pooling_bytes_streaming_response( |
| 176 | pooling_outputs=final_res_batch, |
| 177 | request_id=request_id, |
| 178 | created_time=created_time, |
| 179 | model_name=model_name, |
| 180 | encoding_format=encoding_format, |
| 181 | embed_dtype=embed_dtype, |
| 182 | endianness=endianness, |
| 183 | ) |
| 184 | |
| 185 | def _build_cohere_response_from_ctx( |
| 186 | self, |
| 187 | ctx: PoolingServeContext, |
| 188 | ) -> JSONResponse: |
| 189 | request = ctx.request |
| 190 | assert isinstance(request, CohereEmbedRequest) |
| 191 | |
| 192 | all_floats = [ |
| 193 | cast(list[float], encode_pooling_output_float(out)) |
| 194 | for out in ctx.final_res_batch |
| 195 | ] |
| 196 | total_tokens = get_pooling_usage(ctx.final_res_batch).prompt_tokens |
| 197 | |
| 198 | has_image_input = request.images is not None or any( |
| 199 | content.type == "image_url" |
| 200 | for input_item in request.inputs or [] |
| 201 | for content in input_item.content |
| 202 | ) |
| 203 | image_tokens = total_tokens if has_image_input else 0 |
| 204 | texts_echo = request.texts |
| 205 | |
| 206 | embedding_types = request.embedding_types or ["float"] |
| 207 | embeddings_obj = build_typed_embeddings(all_floats, embedding_types) |
| 208 | |
| 209 | input_tokens = total_tokens - image_tokens |
| 210 | response = CohereEmbedResponse( |
| 211 | id=ctx.request_id, |
| 212 | embeddings=embeddings_obj, |
| 213 | texts=texts_echo, |
| 214 | meta=CohereMeta( |
| 215 | billed_units=CohereBilledUnits( |
| 216 | input_tokens=input_tokens, |
| 217 | image_tokens=image_tokens, |
| 218 | ), |
| 219 | ), |
| 220 | ) |
| 221 | return self.json_response_cls(content=response.model_dump(exclude_none=True)) |