Add an optional local embedding server

.github/workflows/python-tests.yml

@@ -57,7 +57,7 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
-        run: uv sync --locked --no-group dev --no-group docs
+        run: uv sync --locked --no-group dev --no-group docs --extra server
 
       - name: Run pytest
         env:

.github/workflows/type-checkers.yml

@@ -25,7 +25,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Install dependencies
-        run: uv sync --locked --no-default-groups --group types
+        run: uv sync --locked --no-default-groups --group types --extra server
 
       - name: mypy
         run: |
@@ -36,4 +36,4 @@ jobs:
 
       - name: pyright
         run: |
-          uv run --no-sync pyright tests/type_stub.py
+          uv run --no-sync pyright tests/type_stub.py

.github/workflows/weekly-tests.yml

@@ -35,7 +35,7 @@ jobs:
         with:
           python-version: '3.14'
       - name: Install dependencies
-        run: uv sync --locked --no-group dev --no-group docs
+        run: uv sync --locked --no-group dev --no-group docs --extra server
 
       - name: Run pytest
         # explicit bash runs with pipefail, so tee keeps pytest's exit code

README.md

@@ -24,6 +24,9 @@ pip install fastembed
 pip install fastembed-gpu
 ```
 
+To run a local HTTP embedding service, install `fastembed[server]` and run
+`fastembed serve`. See the [embedding server guide](docs/embedding_server.md).
+
 ## 📖 Quickstart
 
 ```python

docs/embedding_server.md

@@ -0,0 +1,60 @@
+# Local embedding server
+
+Install the optional server dependencies and start a text embedding model.
+
+```bash
+pip install 'fastembed[server]'
+fastembed serve --model-name BAAI/bge-small-en-v1.5
+```
+
+`python -m fastembed serve` starts the same server. The default address is
+`127.0.0.1` and the default port is `8080`. The model is loaded once before
+the server starts. Omitting `--model-name` uses the library's default model.
+
+Use `--model-path /path/to/model-directory` to load existing model files without
+downloading them. The directory must contain the model and tokenizer files
+expected by the selected supported model. A standalone ONNX file is insufficient.
+The model name still selects the architecture and preprocessing.
+
+`--device cpu` is the default. `--device cuda` requires an ONNX Runtime installation
+with the CUDA execution provider and the corresponding CUDA libraries.
+`--cache-dir`, `--threads` and `--batch-size` control the model cache, ONNX Runtime
+threads and inference batch size. `fastembed serve --help` lists all options.
+
+## Requests
+
+```bash
+curl http://127.0.0.1:8080/v1/embeddings \
+  -H 'Content-Type: application/json' \
+  -d '{"model":"BAAI/bge-small-en-v1.5","input":["A local embedding","Another sentence"]}'
+```
+
+The response contains an ordered `data` list with an embedding and index for each
+input, the loaded model name and token usage from the model's tokenizer after
+truncation, including special tokens.
+`GET /v1/models` lists the loaded model and `GET /health` reports server readiness.
+
+The endpoint accepts a nonempty string or a list of up to 2048 nonempty strings.
+Token ID inputs are unsupported. Text is tokenized and truncated using the model's
+normal settings. `encoding_format` accepts `float` or `base64`. `dimensions` may
+only specify the model's native embedding size. An omitted `model` uses the loaded
+model, and requesting a different model returns an error without loading it.
+Requests run through `embed`, with the same preprocessing as the Python library.
+
+An OpenAI client can use the local endpoint with a placeholder key.
+
+```python
+from openai import OpenAI
+
+client = OpenAI(base_url="http://127.0.0.1:8080/v1", api_key="local")
+result = client.embeddings.create(
+    model="BAAI/bge-small-en-v1.5", input=["A local embedding"]
+)
+print(result.data[0].embedding)
+```
+
+The server does not implement authentication. Keep the default loopback binding
+for local use. To share it, configure access control with a reverse proxy before
+binding to an external interface with `--host`. Inference uses a worker thread
+and is serialized per model to avoid concurrent tokenizer mutation. Disconnecting
+a client does not cancel queued or running inference.

fastembed/__main__.py

@@ -0,0 +1,4 @@
+from fastembed.cli import main
+
+if __name__ == "__main__":
+    main()

fastembed/cli.py

@@ -0,0 +1,62 @@
+import argparse
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+
+def _positive_int(value: str) -> int:
+    number = int(value)
+    if number <= 0:
+        raise argparse.ArgumentTypeError("must be a positive integer")
+    return number
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(prog="fastembed")
+    commands = parser.add_subparsers(dest="command", required=True)
+    serve = commands.add_parser("serve", help="Serve local text embeddings over HTTP")
+    serve.add_argument(
+        "--model-name", help="Supported model name, using the library default if omitted"
+    )
+    serve.add_argument(
+        "--model-path", type=Path, help="Local directory containing the model files"
+    )
+    serve.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
+    serve.add_argument("--host", default="127.0.0.1")
+    serve.add_argument("--port", type=_positive_int, default=8080)
+    serve.add_argument("--threads", type=_positive_int, help="ONNX Runtime threads")
+    serve.add_argument("--batch-size", type=_positive_int, default=256)
+    serve.add_argument("--cache-dir", help="Model cache directory")
+    args = parser.parse_args(argv)
+    if args.port > 65535:
+        parser.error("port must be between 1 and 65535")
+    if args.model_path is not None and not args.model_path.is_dir():
+        parser.error("model-path must be an existing directory containing the model files")
+
+    try:
+        import uvicorn
+
+        from fastembed.server import create_app
+    except ImportError as exc:
+        parser.exit(
+            1, f"Server dependencies are unavailable. Install 'fastembed[server]'.\n{exc}\n"
+        )
+
+    from fastembed import TextEmbedding
+    from fastembed.common.types import Device
+
+    options: dict[str, Any] = {
+        "cuda": Device(args.device),
+        "cache_dir": args.cache_dir,
+        "threads": args.threads,
+    }
+    if args.model_name is not None:
+        options["model_name"] = args.model_name
+    if args.model_path is not None:
+        options["specific_model_path"] = str(args.model_path)
+    try:
+        model = TextEmbedding(**options)
+        app = create_app(model, batch_size=args.batch_size)
+    except Exception as exc:
+        parser.exit(1, f"Unable to load the embedding model.\n{exc}\n")
+    uvicorn.run(app, host=args.host, port=args.port)

fastembed/server.py

@@ -0,0 +1,105 @@
+"""Optional HTTP interface for a single local text embedding model."""
+
+import base64
+from threading import Lock
+from typing import Annotated, Any, Literal
+
+import numpy as np
+from fastapi import FastAPI, Request
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse
+from loguru import logger
+from pydantic import BaseModel, ConfigDict, Field, StrictStr
+
+from fastembed.text.text_embedding_base import TextEmbeddingBase
+
+TextInput = Annotated[StrictStr, Field(min_length=1)]
+
+
+class EmbeddingsRequest(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    input: TextInput | Annotated[list[TextInput], Field(min_length=1, max_length=2048)]
+    model: StrictStr | None = None
+    encoding_format: Literal["float", "base64"] = "float"
+    dimensions: Annotated[int, Field(strict=True, gt=0)] | None = None
+    user: StrictStr | None = None
+
+
+def _error(
+    message: str,
+    status: int = 400,
+    param: str | None = None,
+    error_type: str = "invalid_request_error",
+) -> JSONResponse:
+    return JSONResponse(
+        status_code=status,
+        content={"error": {"message": message, "type": error_type, "param": param, "code": None}},
+    )
+
+
+def create_app(model: TextEmbeddingBase, *, batch_size: int = 256) -> FastAPI:
+    """Serve an initialized model. Inference is serialized outside the ASGI event loop."""
+    if batch_size <= 0:
+        raise ValueError("batch_size must be positive")
+    app = FastAPI(title="FastEmbed")
+    lock = Lock()
+    embedding_size = model.embedding_size
+
+    @app.exception_handler(RequestValidationError)
+    async def validation_error(request: Request, exc: RequestValidationError) -> JSONResponse:
+        # Validation details can contain the original text. Return only the field name.
+        location = exc.errors()[0]["loc"]
+        param = location[1] if len(location) > 1 and isinstance(location[1], str) else None
+        return _error("Invalid embedding request.", param=param)
+
+    @app.get("/health")
+    async def health() -> dict[str, str]:
+        return {"status": "ok"}
+
+    @app.get("/v1/models")
+    async def models() -> dict[str, Any]:
+        return {
+            "object": "list",
+            "data": [
+                {"id": model.model_name, "object": "model", "created": 0, "owned_by": "fastembed"}
+            ],
+        }
+
+    @app.post("/v1/embeddings", response_model=None)
+    def embeddings(body: EmbeddingsRequest) -> dict[str, Any] | JSONResponse:
+        if body.model is not None and body.model != model.model_name:
+            return _error("The requested model is not loaded.", status=404, param="model")
+        if body.dimensions is not None and body.dimensions != embedding_size:
+            return _error(
+                "Only the model's native embedding dimensions are supported.", param="dimensions"
+            )
+        inputs = [body.input] if isinstance(body.input, str) else body.input
+        try:
+            with lock:
+                vectors = list(model.embed(inputs, batch_size=batch_size))
+                tokens = model.token_count(inputs)
+            if len(vectors) != len(inputs):
+                raise ValueError("The model returned an unexpected number of embeddings")
+            data = []
+            for index, vector in enumerate(vectors):
+                values = np.asarray(vector, dtype="<f4")
+                if values.shape != (embedding_size,) or not np.isfinite(values).all():
+                    raise ValueError("The model returned an invalid embedding")
+                encoded = (
+                    base64.b64encode(values.tobytes()).decode("ascii")
+                    if body.encoding_format == "base64"
+                    else values.tolist()
+                )
+                data.append({"object": "embedding", "index": index, "embedding": encoded})
+            return {
+                "object": "list",
+                "data": data,
+                "model": model.model_name,
+                "usage": {"prompt_tokens": tokens, "total_tokens": tokens},
+            }
+        except Exception:
+            logger.exception("Embedding generation failed")
+            return _error("Embedding generation failed.", status=500, error_type="server_error")
+
+    return app

pyproject.toml

@@ -42,10 +42,21 @@ dependencies = [
 Homepage = "https://github.com/qdrant/fastembed"
 Repository = "https://github.com/qdrant/fastembed"
 
+[project.scripts]
+fastembed = "fastembed.cli:main"
+
+[project.optional-dependencies]
+server = [
+    "fastapi>=0.115,<1",
+    "pydantic>=2,<3",
+    "uvicorn>=0.30,<1",
+]
+
 [dependency-groups]
 test = [
     "pytest>=7.4.2,<10.0.0",
     "ruff>=0.3.1,<1.0",
+    "httpx>=0.27,<1",
 ]
 dev = [
     "notebook>=7.0.2",

tests/test_cli.py

@@ -0,0 +1,92 @@
+from pathlib import Path
+from unittest.mock import Mock
+
+import pytest
+
+from fastembed.cli import main
+from fastembed.common.types import Device
+
+
+@pytest.fixture(autouse=True)
+def server_dependencies() -> None:
+    pytest.importorskip("fastapi")
+    pytest.importorskip("uvicorn")
+
+
+def test_serve_loads_one_model(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+    model = Mock()
+    factory = Mock(return_value=model)
+    create_app = Mock()
+    run = Mock()
+    monkeypatch.setattr("fastembed.TextEmbedding", factory)
+    monkeypatch.setattr("fastembed.server.create_app", create_app)
+    monkeypatch.setattr("uvicorn.run", run)
+    main(
+        [
+            "serve",
+            "--model-name",
+            "test-model",
+            "--model-path",
+            str(tmp_path),
+            "--device",
+            "cuda",
+            "--host",
+            "0.0.0.0",
+            "--port",
+            "9000",
+            "--threads",
+            "2",
+            "--batch-size",
+            "16",
+            "--cache-dir",
+            str(tmp_path / "cache"),
+        ]
+    )
+    factory.assert_called_once_with(
+        model_name="test-model",
+        specific_model_path=str(tmp_path),
+        cuda=Device.CUDA,
+        cache_dir=str(tmp_path / "cache"),
+        threads=2,
+    )
+    create_app.assert_called_once_with(model, batch_size=16)
+    run.assert_called_once_with(create_app.return_value, host="0.0.0.0", port=9000)
+
+
+def test_serve_defaults(monkeypatch: pytest.MonkeyPatch) -> None:
+    factory, create_app, run = Mock(), Mock(), Mock()
+    monkeypatch.setattr("fastembed.TextEmbedding", factory)
+    monkeypatch.setattr("fastembed.server.create_app", create_app)
+    monkeypatch.setattr("uvicorn.run", run)
+    main(["serve"])
+    factory.assert_called_once_with(cuda=Device.CPU, cache_dir=None, threads=None)
+    run.assert_called_once_with(create_app.return_value, host="127.0.0.1", port=8080)
+
+
+@pytest.mark.parametrize(
+    "args",
+    [
+        [],
+        ["serve", "--port", "0"],
+        ["serve", "--port", "65536"],
+        ["serve", "--batch-size", "-1"],
+        ["serve", "--threads", "0"],
+        ["serve", "--device", "invalid"],
+        ["serve", "--model-path", "missing-directory"],
+    ],
+)
+def test_invalid_arguments(args: list[str]) -> None:
+    with pytest.raises(SystemExit) as exc:
+        main(args)
+    assert exc.value.code == 2
+
+
+def test_model_load_error(monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture) -> None:
+    monkeypatch.setattr("fastembed.TextEmbedding", Mock(side_effect=ValueError("invalid model")))
+    run = Mock()
+    monkeypatch.setattr("uvicorn.run", run)
+    with pytest.raises(SystemExit) as exc:
+        main(["serve", "--model-name", "not-supported"])
+    assert exc.value.code == 1
+    assert "Unable to load the embedding model" in capsys.readouterr().err
+    run.assert_not_called()