litserve 0.2.5.dev0__tar.gz → 0.2.6.dev0__tar.gz

{litserve-0.2.5.dev0/src/litserve.egg-info → litserve-0.2.6.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: litserve
-Version: 0.2.5.dev0
+Version: 0.2.6.dev0
 Summary: Lightweight AI server.
 Home-page: https://github.com/Lightning-AI/litserve
 Download-URL: https://github.com/Lightning-AI/litserve
@@ -28,11 +28,11 @@ Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: fastapi>=0.100
-Requires-Dist: httpx
 Requires-Dist: uvicorn[standard]>=0.29.0
 Provides-Extra: test
 Requires-Dist: asgi-lifespan; extra == "test"
 Requires-Dist: coverage[toml]>=7.5.3; extra == "test"
+Requires-Dist: httpx>=0.27.0; extra == "test"
 Requires-Dist: lightning>2.0.0; extra == "test"
 Requires-Dist: mypy==1.11.2; extra == "test"
 Requires-Dist: numpy<2.0; extra == "test"
@@ -41,6 +41,7 @@ Requires-Dist: pillow; extra == "test"
 Requires-Dist: psutil; extra == "test"
 Requires-Dist: pytest-asyncio; extra == "test"
 Requires-Dist: pytest-cov; extra == "test"
+Requires-Dist: pytest-retry>=1.6.3; extra == "test"
 Requires-Dist: pytest>=8.0; extra == "test"
 Requires-Dist: python-multipart; extra == "test"
 Requires-Dist: requests; extra == "test"
@@ -208,7 +209,7 @@ Use LitServe to deploy any model or AI service: (Compound AI, Gen AI, classic ML
                 <a target="_blank" href="https://lightning.ai/bhimrajyadav/studios/deploy-background-removal-api-with-litserve">Background Removal</a>, <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-a-controlled-image-generation-api-controlnet">Control Stable Diffusion (ControlNet)</a>
 <strong>Speech:</strong>         <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-a-voice-clone-api-coqui-xtts-v2-model">Text-speech (XTTS V2)</a>, <a target="_blank" href="https://lightning.ai/bhimrajyadav/studios/deploy-a-speech-generation-api-using-parler-tts-powered-by-litserve">Parler-TTS</a>
 <strong>Classical ML:</strong>   <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-random-forest-with-litserve">Random forest</a>, <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-xgboost-with-litserve">XGBoost</a>
-<strong>Miscellaneous:</strong>  <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-an-media-conversion-api-with-ffmpeg">Media conversion API (ffmpeg)</a>, <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-both-pytorch-and-tensorflow-in-a-single-api">PyTorch + TensorFlow in one API</a>
+<strong>Miscellaneous:</strong>  <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-an-media-conversion-api-with-ffmpeg">Media conversion API (ffmpeg)</a>, <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-both-pytorch-and-tensorflow-in-a-single-api">PyTorch + TensorFlow in one API</a>, <a target="_blank" href="https://lightning.ai/lightning-ai/studios/openai-fault-tolerant-proxy-server">LLM proxy server</a>
 </pre>
 </pre>
 

{litserve-0.2.5.dev0 → litserve-0.2.6.dev0}/README.md

@@ -155,7 +155,7 @@ Use LitServe to deploy any model or AI service: (Compound AI, Gen AI, classic ML
                 <a target="_blank" href="https://lightning.ai/bhimrajyadav/studios/deploy-background-removal-api-with-litserve">Background Removal</a>, <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-a-controlled-image-generation-api-controlnet">Control Stable Diffusion (ControlNet)</a>
 <strong>Speech:</strong>         <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-a-voice-clone-api-coqui-xtts-v2-model">Text-speech (XTTS V2)</a>, <a target="_blank" href="https://lightning.ai/bhimrajyadav/studios/deploy-a-speech-generation-api-using-parler-tts-powered-by-litserve">Parler-TTS</a>
 <strong>Classical ML:</strong>   <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-random-forest-with-litserve">Random forest</a>, <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-xgboost-with-litserve">XGBoost</a>
-<strong>Miscellaneous:</strong>  <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-an-media-conversion-api-with-ffmpeg">Media conversion API (ffmpeg)</a>, <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-both-pytorch-and-tensorflow-in-a-single-api">PyTorch + TensorFlow in one API</a>
+<strong>Miscellaneous:</strong>  <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-an-media-conversion-api-with-ffmpeg">Media conversion API (ffmpeg)</a>, <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-both-pytorch-and-tensorflow-in-a-single-api">PyTorch + TensorFlow in one API</a>, <a target="_blank" href="https://lightning.ai/lightning-ai/studios/openai-fault-tolerant-proxy-server">LLM proxy server</a>
 </pre>
 </pre>
 

{litserve-0.2.5.dev0 → litserve-0.2.6.dev0}/requirements.txt

@@ -1,3 +1,2 @@
 fastapi >=0.100
-httpx
 uvicorn[standard] >=0.29.0

{litserve-0.2.5.dev0 → litserve-0.2.6.dev0}/src/litserve/__about__.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-__version__ = "0.2.5.dev0"
+__version__ = "0.2.6.dev0"
 __author__ = "Lightning-AI et al."
 __author_email__ = "community@lightning.ai"
 __license__ = "Apache-2.0"

{litserve-0.2.5.dev0 → litserve-0.2.6.dev0}/src/litserve/__init__.py

@@ -17,6 +17,16 @@ from litserve.api import LitAPI
 from litserve.callbacks import Callback
 from litserve.loggers import Logger
 from litserve.server import LitServer, Request, Response
-from litserve.specs.openai import OpenAISpec
+from litserve.specs import OpenAIEmbeddingSpec, OpenAISpec
 
-__all__ = ["LitAPI", "LitServer", "Request", "Response", "OpenAISpec", "test_examples", "Callback", "Logger"]
+__all__ = [
+    "LitAPI",
+    "LitServer",
+    "Request",
+    "Response",
+    "OpenAISpec",
+    "OpenAIEmbeddingSpec",
+    "test_examples",
+    "Callback",
+    "Logger",
+]

{litserve-0.2.5.dev0 → litserve-0.2.6.dev0}/src/litserve/connector.py

@@ -39,6 +39,21 @@ class _Connector:
         else:
             self._devices = devices
 
+        self.check_devices_and_accelerators()
+
+    def check_devices_and_accelerators(self):
+        """Check if the devices are in a valid fomra and raise an error if they are not."""
+        if self._accelerator in ["cuda", "mps"]:
+            if not isinstance(self._devices, int) and not (
+                isinstance(self._devices, list) and all(isinstance(device, int) for device in self._devices)
+            ):
+                raise ValueError(
+                    "devices must be an integer or a list of integers when using 'cuda' or 'mps', "
+                    f"instead got {self._devices}"
+                )
+        elif self._accelerator != "cpu":
+            raise ValueError(f"accelerator must be one of (cuda, mps, cpu), instead got {self._accelerator}")
+
     @property
     def accelerator(self):
         return self._accelerator

{litserve-0.2.5.dev0 → litserve-0.2.6.dev0}/src/litserve/docker_builder.py

@@ -67,7 +67,7 @@ CMD ["python", "/app/{server_filename}"]
 """
 
 CUDA_DOCKER_TEMPLATE = """# Change CUDA and cuDNN version here
-FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu20.04
+FROM nvidia/cuda:12.4.1-base-ubuntu22.04
 ARG PYTHON_VERSION=3.12
 
 ENV DEBIAN_FRONTEND=noninteractive

{litserve-0.2.5.dev0 → litserve-0.2.6.dev0}/src/litserve/loops.py

@@ -17,6 +17,7 @@ import logging
 import multiprocessing as mp
 import sys
 import time
+from abc import ABC
 from queue import Empty, Queue
 from typing import Dict, List, Optional, Tuple, Union
 
@@ -30,7 +31,6 @@ from litserve.utils import LitAPIStatus, PickleableHTTPException, WorkerSetupSta
 
 mp.allow_connection_pickling()
 
-
 try:
     import uvloop
 
@@ -389,6 +389,190 @@ def run_batched_streaming_loop(
             response_queues[response_queue_id].put((uid, (e, LitAPIStatus.ERROR)))
 
 
+class _BaseLoop(ABC):
+    """Loop runs an inference engine that executes a specific set of hooks, implemented in the LitAPI, in a predefined
+    order.
+
+    For a default loop, LitAPI must implement the following hooks:
+      - decode_request
+      - batch
+      - predict
+      - unbatch
+      - encode_response
+
+    To implement a custom loop, subclass this class and implement the `run` method. The `run` method should execute the
+    hooks in the desired order.
+
+    `__call__` method is the entry point for the worker process. It calls the `run` method in a loop until the worker is
+    terminated.
+
+    Example:
+
+    ```python
+    class TestLoop(_BaseLoop):
+        def run(
+            self,
+            lit_api: LitAPI,
+            lit_spec: Optional[LitSpec],
+            device: str,
+            worker_id: int,
+            request_queue: Queue,
+            response_queues: List[Queue],
+            max_batch_size: int,
+            batch_timeout: float,
+            stream: bool,
+            workers_setup_status: Dict[int, str],
+            callback_runner: CallbackRunner,
+        ):
+            item = request_queue.get()
+            if item is None:
+                return
+
+            response_queue_id, uid, timestamp, x_enc = item
+            # Expects LitAPI to implement the load_cache method
+            lit_api.load_cache(x_enc)
+            x = lit_api.decode_request(x_enc)
+            response = lit_api.predict(x)
+            response_enc = lit_api.encode_response(response)
+            response_queues[response_queue_id].put((uid, (response_enc, LitAPIStatus.OK)))
+    ```
+
+    """
+
+    def __call__(
+        self,
+        lit_api: LitAPI,
+        lit_spec: Optional[LitSpec],
+        device: str,
+        worker_id: int,
+        request_queue: Queue,
+        response_queues: List[Queue],
+        max_batch_size: int,
+        batch_timeout: float,
+        stream: bool,
+        workers_setup_status: Dict[int, str],
+        callback_runner: CallbackRunner,
+    ):
+        while True:
+            self.run(
+                lit_api,
+                lit_spec,
+                device,
+                worker_id,
+                request_queue,
+                response_queues,
+                max_batch_size,
+                batch_timeout,
+                stream,
+                workers_setup_status,
+                callback_runner,
+            )
+
+    def run(
+        self,
+        lit_api: LitAPI,
+        lit_spec: Optional[LitSpec],
+        device: str,
+        worker_id: int,
+        request_queue: Queue,
+        response_queues: List[Queue],
+        max_batch_size: int,
+        batch_timeout: float,
+        stream: bool,
+        workers_setup_status: Dict[int, str],
+        callback_runner: CallbackRunner,
+    ):
+        raise NotImplementedError
+
+
+class SingleLoop(_BaseLoop):
+    def __call__(
+        self,
+        lit_api: LitAPI,
+        lit_spec: Optional[LitSpec],
+        device: str,
+        worker_id: int,
+        request_queue: Queue,
+        response_queues: List[Queue],
+        max_batch_size: int,
+        batch_timeout: float,
+        stream: bool,
+        workers_setup_status: Dict[int, str],
+        callback_runner: CallbackRunner,
+    ):
+        run_single_loop(lit_api, lit_spec, request_queue, response_queues, callback_runner)
+
+
+class BatchedLoop(_BaseLoop):
+    def __call__(
+        self,
+        lit_api: LitAPI,
+        lit_spec: Optional[LitSpec],
+        device: str,
+        worker_id: int,
+        request_queue: Queue,
+        response_queues: List[Queue],
+        max_batch_size: int,
+        batch_timeout: float,
+        stream: bool,
+        workers_setup_status: Dict[int, str],
+        callback_runner: CallbackRunner,
+    ):
+        run_batched_loop(
+            lit_api,
+            lit_spec,
+            request_queue,
+            response_queues,
+            max_batch_size,
+            batch_timeout,
+            callback_runner,
+        )
+
+
+class StreamingLoop(_BaseLoop):
+    def __call__(
+        self,
+        lit_api: LitAPI,
+        lit_spec: Optional[LitSpec],
+        device: str,
+        worker_id: int,
+        request_queue: Queue,
+        response_queues: List[Queue],
+        max_batch_size: int,
+        batch_timeout: float,
+        stream: bool,
+        workers_setup_status: Dict[int, str],
+        callback_runner: CallbackRunner,
+    ):
+        run_streaming_loop(lit_api, lit_spec, request_queue, response_queues, callback_runner)
+
+
+class BatchedStreamingLoop(_BaseLoop):
+    def __call__(
+        self,
+        lit_api: LitAPI,
+        lit_spec: Optional[LitSpec],
+        device: str,
+        worker_id: int,
+        request_queue: Queue,
+        response_queues: List[Queue],
+        max_batch_size: int,
+        batch_timeout: float,
+        stream: bool,
+        workers_setup_status: Dict[int, str],
+        callback_runner: CallbackRunner,
+    ):
+        run_batched_streaming_loop(
+            lit_api,
+            lit_spec,
+            request_queue,
+            response_queues,
+            max_batch_size,
+            batch_timeout,
+            callback_runner,
+        )
+
+
 def inference_worker(
     lit_api: LitAPI,
     lit_spec: Optional[LitSpec],
@@ -401,6 +585,7 @@ def inference_worker(
     stream: bool,
     workers_setup_status: Dict[int, str],
     callback_runner: CallbackRunner,
+    loop: Union[str, _BaseLoop],
 ):
     callback_runner.trigger_event(EventTypes.BEFORE_SETUP, lit_api=lit_api)
     try:
@@ -412,25 +597,35 @@ def inference_worker(
     lit_api.device = device
     callback_runner.trigger_event(EventTypes.AFTER_SETUP, lit_api=lit_api)
 
-    logger.info(f"Setup complete for worker {worker_id}.")
+    print(f"Setup complete for worker {worker_id}.")
 
     if workers_setup_status:
         workers_setup_status[worker_id] = WorkerSetupStatus.READY
 
     if lit_spec:
         logging.info(f"LitServe will use {lit_spec.__class__.__name__} spec")
-    if stream:
-        if max_batch_size > 1:
-            run_batched_streaming_loop(
-                lit_api, lit_spec, request_queue, response_queues, max_batch_size, batch_timeout, callback_runner
-            )
-        else:
-            run_streaming_loop(lit_api, lit_spec, request_queue, response_queues, callback_runner)
-        return
 
-    if max_batch_size > 1:
-        run_batched_loop(
-            lit_api, lit_spec, request_queue, response_queues, max_batch_size, batch_timeout, callback_runner
+    if loop == "auto":
+        loop = (
+            BatchedStreamingLoop()
+            if stream and max_batch_size > 1
+            else StreamingLoop()
+            if stream
+            else BatchedLoop()
+            if max_batch_size > 1
+            else SingleLoop()
         )
-    else:
-        run_single_loop(lit_api, lit_spec, request_queue, response_queues, callback_runner)
+
+    loop(
+        lit_api,
+        lit_spec,
+        device,
+        worker_id,
+        request_queue,
+        response_queues,
+        max_batch_size,
+        batch_timeout,
+        stream,
+        workers_setup_status,
+        callback_runner,
+    )

litserve-0.2.6.dev0/src/litserve/schema/image.py

@@ -0,0 +1,75 @@
+import base64
+from io import BytesIO
+from typing import TYPE_CHECKING, Any, Optional
+
+from pydantic import BaseModel, field_serializer, model_validator
+
+if TYPE_CHECKING:
+    from PIL import Image
+
+
+class ImageInput(BaseModel):
+    image_data: Optional[str] = None
+
+    @model_validator(mode="after")
+    def validate_base64(self) -> "ImageInput":
+        """Ensure the string is a valid Base64."""
+        model_dump = self.model_dump()
+        for key, value in model_dump.items():
+            if value:
+                try:
+                    base64.b64decode(value)
+                except base64.binascii.Error:
+                    raise ValueError("Invalid Base64 string.")
+        return self
+
+    def get_image(self, key: Optional[str] = None) -> "Image.Image":
+        """Decode the Base64 string and return a PIL Image object."""
+        if key is None:
+            key = "image_data"
+        image_data = self.model_dump().get(key)
+        if not image_data:
+            raise ValueError(f"Missing image data for key '{key}'")
+        try:
+            from PIL import Image, UnidentifiedImageError
+        except ImportError:
+            raise ImportError("Pillow is required to use the ImageInput schema. Install it with `pip install Pillow`.")
+        try:
+            decoded_data = base64.b64decode(image_data)
+            return Image.open(BytesIO(decoded_data))
+        except UnidentifiedImageError as e:
+            raise ValueError(f"Error loading image from decoded data: {e}")
+
+
+class ImageOutput(BaseModel):
+    image: Any
+
+    @field_serializer("image")
+    def serialize_image(self, image: Any, _info):
+        """
+        Serialize a PIL Image into a base64 string.
+        Args:
+            image (Any): The image object to serialize.
+            _info: Metadata passed during serialization (not used here).
+
+        Returns:
+            str: Base64-encoded image string.
+        """
+        try:
+            from PIL import Image
+        except ImportError:
+            raise ImportError("Pillow is required to use the ImageOutput schema. Install it with `pip install Pillow`.")
+
+        if not isinstance(image, Image.Image):
+            raise TypeError(f"Expected a PIL Image, got {type(image)}")
+
+        # Save the image to a BytesIO buffer
+        buffer = BytesIO()
+        image.save(buffer, format="PNG")  # Default format is PNG
+        buffer.seek(0)
+
+        # Encode the buffer content to base64
+        base64_bytes = base64.b64encode(buffer.read())
+
+        # Decode to string for JSON serialization
+        return base64_bytes.decode("utf-8")

{litserve-0.2.5.dev0 → litserve-0.2.6.dev0}/src/litserve/server.py

@@ -14,6 +14,7 @@
 import asyncio
 import copy
 import inspect
+import json
 import logging
 import multiprocessing as mp
 import os
@@ -30,7 +31,7 @@ from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 import uvicorn
 from fastapi import Depends, FastAPI, HTTPException, Request, Response
-from fastapi.responses import StreamingResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi.security import APIKeyHeader
 from starlette.formparsers import MultiPartParser
 from starlette.middleware.gzip import GZipMiddleware
@@ -39,7 +40,7 @@ from litserve import LitAPI
 from litserve.callbacks.base import Callback, CallbackRunner, EventTypes
 from litserve.connector import _Connector
 from litserve.loggers import Logger, _LoggerConnector
-from litserve.loops import inference_worker
+from litserve.loops import _BaseLoop, inference_worker
 from litserve.middlewares import MaxSizeMiddleware, RequestCountMiddleware
 from litserve.python_client import client_template
 from litserve.specs import OpenAISpec
@@ -105,14 +106,42 @@ class LitServer:
         max_batch_size: int = 1,
         batch_timeout: float = 0.0,
         api_path: str = "/predict",
+        healthcheck_path: str = "/health",
+        info_path: str = "/info",
+        model_metadata: Optional[dict] = None,
         stream: bool = False,
         spec: Optional[LitSpec] = None,
         max_payload_size=None,
         track_requests: bool = False,
+        loop: Optional[Union[str, _BaseLoop]] = "auto",
         callbacks: Optional[Union[List[Callback], Callback]] = None,
         middlewares: Optional[list[Union[Callable, tuple[Callable, dict]]]] = None,
         loggers: Optional[Union[Logger, List[Logger]]] = None,
     ):
+        """Initialize a LitServer instance.
+
+        Args:
+            lit_api: The LitAPI instance to use for handling requests.
+            accelerator: The type of hardware accelerator to use (e.g., 'auto', 'cpu', 'cuda', 'mps').
+            devices: The number of devices to use (e.g., 'auto', 1, 2).
+            workers_per_device: The number of workers to use per device.
+            timeout: The timeout for requests in seconds.
+            max_batch_size: The maximum batch size.
+            batch_timeout: The timeout for batching requests in seconds.
+            api_path: The path for the prediction API endpoint.
+            healthcheck_path: The path for the health check endpoint.
+            info_path: The path for the server and model metadata info endpoint.
+            model_metadata: Metadata about the model, it will be shown via the `info_path` endpoint.
+            stream: Whether to enable streaming responses.
+            spec: The specification for the API such as OpenAISpec or OpenAIEmbeddingSpec.
+            max_payload_size: The maximum payload size for requests.
+            track_requests: Whether to track the number of active requests.
+            loop: The inference engine runs with this loop in the worker process.
+            callbacks: Callbacks to use for the server.
+            middlewares: ASGI middleware for the server.
+            loggers: Loggers to use for the server.
+
+        """
         if batch_timeout > timeout and timeout not in (False, -1):
             raise ValueError("batch_timeout must be less than timeout")
         if max_batch_size <= 0:
@@ -120,6 +149,12 @@ class LitServer:
         if isinstance(spec, OpenAISpec):
             stream = True
 
+        if loop is None:
+            loop = "auto"
+
+        if isinstance(loop, str) and loop != "auto":
+            raise ValueError("loop must be an instance of _BaseLoop or 'auto'")
+
         if middlewares is None:
             middlewares = []
         if not isinstance(middlewares, list):
@@ -137,6 +172,22 @@ class LitServer:
                 "Please provide a valid api path like '/predict', '/classify', or '/v1/predict'"
             )
 
+        if not healthcheck_path.startswith("/"):
+            raise ValueError(
+                "healthcheck_path must start with '/'. "
+                "Please provide a valid api path like '/health', '/healthcheck', or '/v1/health'"
+            )
+
+        if not info_path.startswith("/"):
+            raise ValueError(
+                "info_path must start with '/'. Please provide a valid api path like '/info', '/details', or '/v1/info'"
+            )
+
+        try:
+            json.dumps(model_metadata)
+        except (TypeError, ValueError):
+            raise ValueError("model_metadata must be JSON serializable.")
+
         # Check if the batch and unbatch methods are overridden in the lit_api instance
         batch_overridden = lit_api.batch.__code__ is not LitAPI.batch.__code__
         unbatch_overridden = lit_api.unbatch.__code__ is not LitAPI.unbatch.__code__
@@ -147,10 +198,14 @@ class LitServer:
                 "but the max_batch_size parameter was not set."
             )
 
+        self._loop = loop
         self.api_path = api_path
+        self.healthcheck_path = healthcheck_path
+        self.info_path = info_path
         self.track_requests = track_requests
+        self.timeout = timeout
         lit_api.stream = stream
-        lit_api.request_timeout = timeout
+        lit_api.request_timeout = self.timeout
         lit_api._sanitize(max_batch_size, spec=spec)
         self.app = FastAPI(lifespan=self.lifespan)
         self.app.response_queue_id = None
@@ -172,6 +227,7 @@ class LitServer:
         self.batch_timeout = batch_timeout
         self.stream = stream
         self.max_payload_size = max_payload_size
+        self.model_metadata = model_metadata
         self._connector = _Connector(accelerator=accelerator, devices=devices)
         self._callback_runner = CallbackRunner(callbacks)
 
@@ -242,6 +298,7 @@ class LitServer:
                     self.stream,
                     self.workers_setup_status,
                     self._callback_runner,
+                    self._loop,
                 ),
             )
             process.start()
@@ -316,7 +373,7 @@ class LitServer:
         async def index(request: Request) -> Response:
             return Response(content="litserve running")
 
-        @self.app.get("/health", dependencies=[Depends(self.setup_auth())])
+        @self.app.get(self.healthcheck_path, dependencies=[Depends(self.setup_auth())])
         async def health(request: Request) -> Response:
             nonlocal workers_ready
             if not workers_ready:
@@ -327,6 +384,24 @@ class LitServer:
 
             return Response(content="not ready", status_code=503)
 
+        @self.app.get(self.info_path, dependencies=[Depends(self.setup_auth())])
+        async def info(request: Request) -> Response:
+            return JSONResponse(
+                content={
+                    "model": self.model_metadata,
+                    "server": {
+                        "devices": self.devices,
+                        "workers_per_device": self.workers_per_device,
+                        "timeout": self.timeout,
+                        "max_batch_size": self.max_batch_size,
+                        "batch_timeout": self.batch_timeout,
+                        "stream": self.stream,
+                        "max_payload_size": self.max_payload_size,
+                        "track_requests": self.track_requests,
+                    },
+                }
+            )
+
         async def predict(request: self.request_type) -> self.response_type:
             self._callback_runner.trigger_event(
                 EventTypes.ON_REQUEST,

litserve-0.2.6.dev0/src/litserve/specs/__init__.py

@@ -0,0 +1,4 @@
+from litserve.specs.openai import OpenAISpec
+from litserve.specs.openai_embedding import OpenAIEmbeddingSpec
+
+__all__ = ["OpenAISpec", "OpenAIEmbeddingSpec"]

{litserve-0.2.5.dev0 → litserve-0.2.6.dev0}/src/litserve/specs/openai.py

@@ -156,7 +156,8 @@ class ChatCompletionRequest(BaseModel):
     temperature: Optional[float] = 0.7
     top_p: Optional[float] = 1.0
     n: Optional[int] = 1
-    max_tokens: Optional[int] = None
+    max_tokens: Optional[int] = None  # Kept for backward compatibility
+    max_completion_tokens: Optional[int] = None
     stop: Optional[Union[str, List[str]]] = None
     stream: Optional[bool] = False
     presence_penalty: Optional[float] = 0.0
@@ -365,7 +366,7 @@ class OpenAISpec(LitSpec):
         if request.stream:
             return StreamingResponse(
                 self.streaming_completion(request, responses),
-                media_type="application/x-ndjson",
+                media_type="text/event-stream",
                 background=background_tasks,
             )
 
@@ -394,9 +395,9 @@ class OpenAISpec(LitSpec):
 
             # Only use the last item from encode_response
             usage_info = sum(usage_infos)
-            chunk = ChatCompletionChunk(model=model, choices=choices, usage=None).json()
+            chunk = ChatCompletionChunk(model=model, choices=choices, usage=None)
             logger.debug(chunk)
-            yield f"data: {chunk}\n\n"
+            yield f"data: {chunk.model_dump_json()}\n\n"
 
         choices = [
             ChatCompletionStreamingChoice(
@@ -410,8 +411,8 @@ class OpenAISpec(LitSpec):
             model=model,
             choices=choices,
             usage=usage_info,
-        ).json()
-        yield f"data: {last_chunk}\n\n"
+        )
+        yield f"data: {last_chunk.model_dump_json()}\n\n"
         yield "data: [DONE]\n\n"
 
     async def non_streaming_completion(self, request: ChatCompletionRequest, generator_list: List[AsyncGenerator]):

litserve-0.2.6.dev0/src/litserve/specs/openai_embedding.py

@@ -0,0 +1,175 @@
+# Copyright The Lightning AI team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import asyncio
+import inspect
+import logging
+import time
+import uuid
+from typing import List, Literal, Optional, Union
+
+from fastapi import Request, Response, status
+from pydantic import BaseModel
+
+from litserve.specs.base import LitSpec
+from litserve.utils import LitAPIStatus
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingRequest(BaseModel):
+    input: Union[str, List[str]]
+    model: str
+    dimensions: Optional[int] = None
+    encoding_format: Literal["float"] = "float"
+
+    def ensure_list(self):
+        return self.input if isinstance(self.input, list) else [self.input]
+
+
+class Embedding(BaseModel):
+    index: int
+    embedding: List[float]
+    object: Literal["embedding"] = "embedding"
+
+
+class UsageInfo(BaseModel):
+    prompt_tokens: int = 0
+    total_tokens: int = 0
+
+
+class EmbeddingResponse(BaseModel):
+    data: List[Embedding]
+    model: str
+    object: Literal["list"] = "list"
+    usage: UsageInfo
+
+
+EMBEDDING_API_EXAMPLE = """
+Please follow the example below for guidance on how to use the OpenAI Embedding spec:
+
+```python
+import numpy as np
+from typing import List
+from litserve import LitAPI, OpenAIEmbeddingSpec
+
+class TestAPI(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()
+
+    def encode_response(self, output) -> dict:
+        return {"embeddings": output}
+
+if __name__ == "__main__":
+    import litserve as ls
+    server = ls.LitServer(TestAPI(), 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.options_embeddings, ["GET"])
+
+    def setup(self, server: "LitServer"):  # noqa: F821
+        from litserve import LitAPI
+
+        super().setup(server)
+
+        lit_api = self._server.lit_api
+        if inspect.isgeneratorfunction(lit_api.predict):
+            raise ValueError(
+                "You are using yield in your predict method, which is used for streaming.",
+                "OpenAIEmbeddingSpec doesn't support streaming because producing embeddings ",
+                "is not a sequential operation.",
+                "Please consider replacing yield with return in predict.\n",
+                EMBEDDING_API_EXAMPLE,
+            )
+
+        is_encode_response_original = lit_api.encode_response.__code__ is LitAPI.encode_response.__code__
+        if not is_encode_response_original and inspect.isgeneratorfunction(lit_api.encode_response):
+            raise ValueError(
+                "You are using yield in your encode_response method, which is used for streaming.",
+                "OpenAIEmbeddingSpec doesn't support streaming because producing embeddings ",
+                "is not a sequential operation.",
+                "Please consider replacing yield with return in encode_response.\n",
+                EMBEDDING_API_EXAMPLE,
+            )
+
+        print("OpenAI Embedding Spec is ready.")
+
+    def decode_request(self, request: EmbeddingRequest, context_kwargs: Optional[dict] = None) -> List[str]:
+        return request.ensure_list()
+
+    def encode_response(self, output: List[List[float]], context_kwargs: Optional[dict] = None) -> dict:
+        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):
+            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"
+                "Please ensure that your response is a dictionary with the following keys:\n"
+                "- 'embeddings' (required)\n"
+                "- 'prompt_tokens' (optional)\n"
+                "- 'total_tokens' (optional)\n"
+                f"{EMBEDDING_API_EXAMPLE}",
+            )
+        if "embeddings" not in response:
+            raise ValueError(
+                "The response does not contain the key 'embeddings'."
+                "The key 'embeddings' is required to ensure proper compatibility with the OpenAIEmbeddingSpec.\n"
+                "Please ensure that your response contains the key 'embeddings'.\n"
+                f"{EMBEDDING_API_EXAMPLE}"
+            )
+
+    async def embeddings(self, request: EmbeddingRequest):
+        response_queue_id = self.response_queue_id
+        logger.debug("Received embedding request: %s", request)
+        uid = uuid.uuid4()
+        event = asyncio.Event()
+        self._server.response_buffer[uid] = event
+
+        self._server.request_queue.put_nowait((response_queue_id, uid, time.monotonic(), request.model_copy()))
+        await event.wait()
+
+        response, status = self._server.response_buffer.pop(uid)
+
+        if status == LitAPIStatus.ERROR:
+            raise response
+
+        logger.debug(response)
+
+        self._validate_response(response)
+
+        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)
+
+    async def options_embeddings(self, request: Request):
+        return Response(status_code=status.HTTP_200_OK)

litserve-0.2.6.dev0/src/litserve/test_examples/openai_embedding_spec_example.py

@@ -0,0 +1,49 @@
+from typing import List
+
+import numpy as np
+
+from litserve.api import LitAPI
+
+
+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()
+
+    def encode_response(self, output) -> dict:
+        return {"embeddings": output}
+
+
+class TestEmbedBatchedAPI(TestEmbedAPI):
+    def predict(self, batch) -> List[List[List[float]]]:
+        return [np.random.rand(len(x), 768).tolist() for x in batch]
+
+
+class TestEmbedAPIWithUsage(TestEmbedAPI):
+    def encode_response(self, output) -> dict:
+        return {"embeddings": output, "prompt_tokens": 10, "total_tokens": 10}
+
+
+class TestEmbedAPIWithYieldPredict(TestEmbedAPI):
+    def predict(self, x):
+        yield from np.random.rand(768).tolist()
+
+
+class TestEmbedAPIWithYieldEncodeResponse(TestEmbedAPI):
+    def encode_response(self, output):
+        yield {"embeddings": output}
+
+
+class TestEmbedAPIWithNonDictOutput(TestEmbedAPI):
+    def encode_response(self, output):
+        return output
+
+
+class TestEmbedAPIWithMissingEmbeddings(TestEmbedAPI):
+    def encode_response(self, output):
+        return {"output": output}

{litserve-0.2.5.dev0 → litserve-0.2.6.dev0/src/litserve.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: litserve
-Version: 0.2.5.dev0
+Version: 0.2.6.dev0
 Summary: Lightweight AI server.
 Home-page: https://github.com/Lightning-AI/litserve
 Download-URL: https://github.com/Lightning-AI/litserve
@@ -28,11 +28,11 @@ Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: fastapi>=0.100
-Requires-Dist: httpx
 Requires-Dist: uvicorn[standard]>=0.29.0
 Provides-Extra: test
 Requires-Dist: asgi-lifespan; extra == "test"
 Requires-Dist: coverage[toml]>=7.5.3; extra == "test"
+Requires-Dist: httpx>=0.27.0; extra == "test"
 Requires-Dist: lightning>2.0.0; extra == "test"
 Requires-Dist: mypy==1.11.2; extra == "test"
 Requires-Dist: numpy<2.0; extra == "test"
@@ -41,6 +41,7 @@ Requires-Dist: pillow; extra == "test"
 Requires-Dist: psutil; extra == "test"
 Requires-Dist: pytest-asyncio; extra == "test"
 Requires-Dist: pytest-cov; extra == "test"
+Requires-Dist: pytest-retry>=1.6.3; extra == "test"
 Requires-Dist: pytest>=8.0; extra == "test"
 Requires-Dist: python-multipart; extra == "test"
 Requires-Dist: requests; extra == "test"
@@ -208,7 +209,7 @@ Use LitServe to deploy any model or AI service: (Compound AI, Gen AI, classic ML
                 <a target="_blank" href="https://lightning.ai/bhimrajyadav/studios/deploy-background-removal-api-with-litserve">Background Removal</a>, <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-a-controlled-image-generation-api-controlnet">Control Stable Diffusion (ControlNet)</a>
 <strong>Speech:</strong>         <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-a-voice-clone-api-coqui-xtts-v2-model">Text-speech (XTTS V2)</a>, <a target="_blank" href="https://lightning.ai/bhimrajyadav/studios/deploy-a-speech-generation-api-using-parler-tts-powered-by-litserve">Parler-TTS</a>
 <strong>Classical ML:</strong>   <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-random-forest-with-litserve">Random forest</a>, <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-xgboost-with-litserve">XGBoost</a>
-<strong>Miscellaneous:</strong>  <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-an-media-conversion-api-with-ffmpeg">Media conversion API (ffmpeg)</a>, <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-both-pytorch-and-tensorflow-in-a-single-api">PyTorch + TensorFlow in one API</a>
+<strong>Miscellaneous:</strong>  <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-an-media-conversion-api-with-ffmpeg">Media conversion API (ffmpeg)</a>, <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-both-pytorch-and-tensorflow-in-a-single-api">PyTorch + TensorFlow in one API</a>, <a target="_blank" href="https://lightning.ai/lightning-ai/studios/openai-fault-tolerant-proxy-server">LLM proxy server</a>
 </pre>
 </pre>
 

{litserve-0.2.5.dev0 → litserve-0.2.6.dev0}/src/litserve.egg-info/SOURCES.txt

@@ -26,9 +26,13 @@ src/litserve/callbacks/__init__.py
 src/litserve/callbacks/base.py
 src/litserve/callbacks/defaults/__init__.py
 src/litserve/callbacks/defaults/metric_callback.py
+src/litserve/schema/__init__.py
+src/litserve/schema/image.py
 src/litserve/specs/__init__.py
 src/litserve/specs/base.py
 src/litserve/specs/openai.py
+src/litserve/specs/openai_embedding.py
 src/litserve/test_examples/__init__.py
+src/litserve/test_examples/openai_embedding_spec_example.py
 src/litserve/test_examples/openai_spec_example.py
 src/litserve/test_examples/simple_example.py

{litserve-0.2.5.dev0 → litserve-0.2.6.dev0}/src/litserve.egg-info/requires.txt

@@ -1,5 +1,4 @@
 fastapi>=0.100
-httpx
 uvicorn[standard]>=0.29.0
 
 [perf]
@@ -10,6 +9,7 @@ uvloop
 [test]
 asgi-lifespan
 coverage[toml]>=7.5.3
+httpx>=0.27.0
 lightning>2.0.0
 mypy==1.11.2
 numpy<2.0
@@ -18,6 +18,7 @@ pillow
 psutil
 pytest-asyncio
 pytest-cov
+pytest-retry>=1.6.3
 pytest>=8.0
 python-multipart
 requests

litserve-0.2.5.dev0/src/litserve/specs/__init__.py

@@ -1,3 +0,0 @@
-from litserve.specs.openai import OpenAISpec
-
-__all__ = ["OpenAISpec"]