fix OpenAI embedding spec for batching (#500)

* refactor: enhance OpenAIEmbeddingSpec with improved response handling and new utility methods

- Introduced `get_num_items` method in `EmbeddingRequest` for better input management.
- Refactored embedding response handling into `_handle_embedding_response` for clarity and error checking.
- Updated endpoint registration to use `embeddings_endpoint` for consistency.
- Improved logging and response validation in the embeddings endpoint.

* update

* update

* refactor: enhance OpenAIEmbeddingSpec and TestAPI for improved input handling and response validation

- Updated `TestAPI` to support batching in the `predict` method, allowing for multiple inputs.
- Added `EMBEDDING_API_EXAMPLE_BATCHING` to provide guidance on using the OpenAI Embedding spec with batching.
- Improved response validation in `OpenAIEmbeddingSpec` to handle mismatches between requested and returned embeddings.
- Removed unused methods and streamlined the code for better clarity and maintainability.

* fix: improve error handling in BatchedLoop and add OpenAI embedding test

- Enhanced error logging in `BatchedLoop` to provide clearer messages when output length mismatches expected input count.
- Updated HTTPException to include a detailed message for better debugging.
- Introduced `openai_embedding_with_batching.py` for end-to-end testing of the OpenAI embedding API with batching support.
- Added assertions in `test_e2e_openai_embedding_with_batching` to validate model and embedding dimensions.

* fix

* fix

* fix: update OpenAI embedding prediction to handle single and multiple inputs

- Modified `predict` method in `TestEmbedAPI` to support both single and batch inputs by adjusting the random embedding generation.
- Added comprehensive tests for OpenAI embedding spec, covering single input, multiple inputs, usage validation, and error handling for missing or incorrect responses.
- Ensured that the tests validate the expected structure and content of the API responses.

* test: enhance OpenAI embedding tests for batching and error handling

- Added tests for client-side batching to validate error responses when dynamic batching is used.
- Improved assertions for status codes and response content in existing tests for better clarity.
- Utilized `copy.deepcopy` to ensure request data integrity during concurrent tests.

* update

* fix ci

* fix test

Author: aniketmaurya

src/litserve/api.py

@@ -176,6 +176,7 @@ def pre_setup(self, spec: Optional[LitSpec]):
 
         if spec:
             self._spec = spec
+            spec._max_batch_size = self.max_batch_size
             spec.pre_setup(self)
 
     def set_logger_queue(self, queue: Queue):

src/litserve/loops/simple_loops.py

@@ -326,11 +326,14 @@ def run_batched_loop(
                 outputs = lit_api.unbatch(y)
 
                 if len(outputs) != num_inputs:
+                    actual = len(outputs)
                     logger.error(
-                        f"LitAPI.predict/unbatch returned {len(outputs)} outputs, but expected {num_inputs}. "
-                        "Please check the predict/unbatch method of the LitAPI implementation."
+                        f"LitAPI.predict/unbatch returned {actual} outputs, but expected {num_inputs}. "
+                        "This suggests a possible issue in the predict or unbatch implementation.\n"
+                        "Hint: Ensure that LitAPI.predict returns a list with one prediction per input — "
+                        "the length of the returned list should match the number of inputs."
                     )
-                    raise HTTPException(500, "Batch size mismatch")
+                    raise HTTPException(500, detail="Batch size mismatch")
 
                 callback_runner.trigger_event(EventTypes.BEFORE_ENCODE_RESPONSE.value, lit_api=lit_api)
                 y_enc_list = []

src/litserve/specs/__init__.py

@@ -1,4 +1,12 @@
-from litserve.specs.openai import OpenAISpec
-from litserve.specs.openai_embedding import OpenAIEmbeddingSpec
+from litserve.specs.openai import ChatCompletionChunk, ChatCompletionRequest, ChatCompletionResponse, OpenAISpec
+from litserve.specs.openai_embedding import EmbeddingRequest, EmbeddingResponse, OpenAIEmbeddingSpec
 
-__all__ = ["OpenAISpec", "OpenAIEmbeddingSpec"]
+__all__ = [
+    "OpenAISpec",
+    "OpenAIEmbeddingSpec",
+    "EmbeddingRequest",
+    "EmbeddingResponse",
+    "ChatCompletionRequest",
+    "ChatCompletionResponse",
+    "ChatCompletionChunk",
+]

src/litserve/specs/base.py

@@ -25,6 +25,7 @@ def __init__(self):
         self._endpoints = []
 
         self._server: LitServer = None
+        self._max_batch_size = 1
 
     @property
     def stream(self):

src/litserve/specs/openai_embedding.py

@@ -14,9 +14,10 @@
 import asyncio
 import inspect
 import logging
+import sys
 import time
 import uuid
-from typing import TYPE_CHECKING, List, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, List, Literal, Optional, Union
 
 from fastapi import HTTPException, Request, Response, status
 from fastapi import status as status_code
@@ -28,6 +29,9 @@
 logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
+    import numpy as np
+    import torch
+
     from litserve import LitServer
 
 
@@ -38,6 +42,14 @@ class EmbeddingRequest(BaseModel):
     encoding_format: Literal["float", "base64"] = "float"
     user: Optional[str] = None
 
+    def get_num_items(self) -> int:
+        """Return the number of sentences or tokens in the input."""
+        if isinstance(self.input, list):
+            if isinstance(self.input[0], list):
+                return len(self.input[0])
+            return len(self.input)
+        return 1
+
     def ensure_list(self):
         return self.input if isinstance(self.input, list) else [self.input]
 
@@ -66,34 +78,54 @@ class EmbeddingResponse(BaseModel):
 ```python
 import numpy as np
 from typing import List
-from litserve import LitAPI, OpenAIEmbeddingSpec
+from litserve.specs import OpenAIEmbeddingSpec, EmbeddingRequest
+import litserve as ls
 
-class TestAPI(LitAPI):
+class TestAPI(ls.LitAPI):
     def setup(self, device):
         self.model = None
 
-    def decode_request(self, request) -> List[str]:
-        return request.ensure_list()
+    def predict(self, inputs) -> List[List[float]]:
+        # inputs is a string
+        return np.random.rand(1, 768).tolist()
 
-    def predict(self, x) -> List[List[float]]:
-        return np.random.rand(len(x), 768).tolist()
-
-    def encode_response(self, output) -> dict:
-        return {"embeddings": output}
 
 if __name__ == "__main__":
-    import litserve as ls
     server = ls.LitServer(TestAPI(), spec=OpenAIEmbeddingSpec())
     server.run()
 ```
 """
 
+EMBEDDING_API_EXAMPLE_BATCHING = """
+Please follow the example below for guidance on how to use the OpenAI Embedding spec with batching:
+
+```python
+import numpy as np
+from typing import List
+from litserve.specs import OpenAIEmbeddingSpec, EmbeddingRequest
+import litserve as ls
+
+class TestAPI(ls.LitAPI):
+    def setup(self, device):
+        self.model = None
+
+    def predict(self, inputs) -> List[List[float]]:
+        # inputs is a list of texts (List[str])
+        return np.random.rand(len(inputs), 768)
+
+if __name__ == "__main__":
+    api = TestAPI(max_batch_size=2, batch_timeout=0.4)
+    server = ls.LitServer(api, spec=OpenAIEmbeddingSpec())
+    server.run()
+```
+"""
+
 
 class OpenAIEmbeddingSpec(LitSpec):
     def __init__(self):
         super().__init__()
         # register the endpoint
-        self.add_endpoint("/v1/embeddings", self.embeddings, ["POST"])
+        self.add_endpoint("/v1/embeddings", self.embeddings_endpoint, ["POST"])
         self.add_endpoint("/v1/embeddings", self.options_embeddings, ["GET"])
 
     def setup(self, server: "LitServer"):
@@ -124,17 +156,21 @@ def setup(self, server: "LitServer"):
         print("OpenAI Embedding Spec is ready.")
 
     def decode_request(self, request: EmbeddingRequest, context_kwargs: Optional[dict] = None) -> List[str]:
-        return request.ensure_list()
+        return request.input
 
-    def encode_response(self, output: List[List[float]], context_kwargs: Optional[dict] = None) -> dict:
+    def encode_response(
+        self, output: List[List[float]], context_kwargs: Optional[dict] = None
+    ) -> Union[dict, EmbeddingResponse]:
         usage = {
             "prompt_tokens": context_kwargs.get("prompt_tokens", 0) if context_kwargs else 0,
             "total_tokens": context_kwargs.get("total_tokens", 0) if context_kwargs else 0,
         }
         return {"embeddings": output} | usage
 
-    def _validate_response(self, response: dict) -> None:
-        if not isinstance(response, dict):
+    def _validate_response(self, response: Union[dict, List[Embedding], Any]) -> None:
+        if isinstance(response, list) and all(isinstance(item, Embedding) for item in response):
+            return
+        if not isinstance(response, (dict, EmbeddingResponse)):
             raise ValueError(
                 f"Expected response to be a dictionary, but got type {type(response)}.",
                 "The response should be a dictionary to ensure proper compatibility with the OpenAIEmbeddingSpec.\n\n"
@@ -152,8 +188,60 @@ def _validate_response(self, response: dict) -> None:
                 f"{EMBEDDING_API_EXAMPLE}"
             )
 
-    async def embeddings(self, request: EmbeddingRequest) -> EmbeddingResponse:
+    def _handle_embedding_response(
+        self, embeddings: Union[List, "np.ndarray", "torch.Tensor", "List[List[float]]"], num_items: int = 1
+    ) -> List[Embedding]:
+        ndim = None
+        if "torch" in sys.modules:
+            import torch
+
+            if isinstance(embeddings, torch.Tensor):
+                ndim = embeddings.ndim
+        if "numpy" in sys.modules:
+            import numpy as np
+
+            if isinstance(embeddings, np.ndarray):
+                ndim = embeddings.ndim
+
+        # expand_dims for torch.Tensor or np.ndarray
+        if ndim == 1:
+            embeddings = embeddings[None, :]
+
+        if ndim is not None:
+            embeddings = embeddings.tolist()
+
+        # expand dims for list of floats
+        if isinstance(embeddings, (list, tuple)) and isinstance(embeddings[0], (int, float)):
+            embeddings = [embeddings]
+
+        # check if we have total num_items number of embeddings vectors
+        num_response_items = len(embeddings)
+        if num_response_items != num_items:
+            logger.debug("mismatch between number of requested and returned embeddings: %s", embeddings)
+            raise ValueError(
+                f"Mismatch between requested and returned embeddings: "
+                f"expected {num_items}, but got {num_response_items}. "
+                f"This may indicate a bug in the LitAPI embedding implementation."
+            )
+
+        result = []
+        for i, embedding in enumerate(embeddings):
+            result.append(Embedding(index=i, embedding=embedding))
+
+        return result
+
+    async def embeddings_endpoint(self, request: EmbeddingRequest) -> EmbeddingResponse:
         response_queue_id = self.response_queue_id
+        num_items = request.get_num_items()
+        if num_items > 1 and self._max_batch_size > 1:
+            raise HTTPException(
+                status_code=400,
+                detail=(
+                    "The OpenAIEmbedding spec does not support dynamic batching when client-side batching is used. "
+                    "To resolve this, either set `max_batch_size=1` or send a single input from the client."
+                ),
+            )
+
         logger.debug("Received embedding request: %s", request)
         uid = uuid.uuid4()
         event = asyncio.Event()
@@ -174,9 +262,9 @@ async def embeddings(self, request: EmbeddingRequest) -> EmbeddingResponse:
         logger.debug(response)
 
         self._validate_response(response)
+        data: List[Embedding] = self._handle_embedding_response(response["embeddings"], num_items)
 
         usage = UsageInfo(**response)
-        data = [Embedding(index=i, embedding=embedding) for i, embedding in enumerate(response["embeddings"])]
 
         return EmbeddingResponse(data=data, model=request.model, usage=usage)
 

src/litserve/test_examples/openai_embedding_spec_example.py

@@ -9,11 +9,9 @@ class TestEmbedAPI(LitAPI):
     def setup(self, device):
         self.model = None
 
-    def decode_request(self, request) -> List[str]:
-        return request.ensure_list()
-
     def predict(self, x) -> List[List[float]]:
-        return np.random.rand(len(x), 768).tolist()
+        n = len(x) if isinstance(x, list) else 1
+        return np.random.rand(n, 768).tolist()
 
     def encode_response(self, output) -> dict:
         return {"embeddings": output}

tests/e2e/openai_embedding_with_batching.py

@@ -0,0 +1,20 @@
+import numpy as np
+
+import litserve as ls
+
+
+class EmbeddingsAPI(ls.LitAPI):
+    def setup(self, device):
+        def model(x):
+            return np.random.rand(len(x), 768)
+
+        self.model = model
+
+    def predict(self, inputs):
+        return self.model(inputs)
+
+
+if __name__ == "__main__":
+    api = EmbeddingsAPI(max_batch_size=10, batch_timeout=2)
+    server = ls.LitServer(api, spec=ls.OpenAIEmbeddingSpec())
+    server.run(port=8000)

tests/e2e/test_e2e.py

@@ -15,6 +15,7 @@
 import os
 import subprocess
 import time
+from concurrent.futures import ThreadPoolExecutor
 from functools import wraps
 
 import psutil
@@ -390,3 +391,26 @@ def test_e2e_default_async_streaming():
             outputs.append(json.loads(line.decode("utf-8"))["output"])
 
     assert outputs == list(range(10)), "server didn't return expected output"
+
+
+@e2e_from_file("tests/e2e/openai_embedding_with_batching.py")
+def test_e2e_openai_embedding_with_batching():
+    model = "text-embedding-3-large"
+    client = OpenAI(
+        base_url="http://127.0.0.1:8000/v1",
+        api_key="lit",  # required, but unused
+    )
+    futures = []
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        futures.append(executor.submit(client.embeddings.create, model=model, input=["This is the first request"]))
+        futures.append(executor.submit(client.embeddings.create, model=model, input=["This is the second request"]))
+        futures.append(executor.submit(client.embeddings.create, model=model, input=["This is the first request"]))
+        futures.append(executor.submit(client.embeddings.create, model=model, input=["This is the second request"]))
+
+    responses = [future.result() for future in futures]
+    for response in responses:
+        assert response.model == model, f"Expected model to be {model} but got {response.model}"
+        assert len(response.data[0].embedding) == 768, (
+            f"Expected 768 dimensions but got {len(response.data[0].embedding)}"
+        )
+    assert len(responses) == 4, f"Expected 4 responses but got {len(responses)}"