mirror of
https://github.com/mudler/LocalAI.git
synced 2026-08-04 20:33:05 -04:00
Compare commits
1 Commits
feat/dllm-
...
fix/vllm-c
| Author | SHA1 | Date | |
|---|---|---|---|
|
a4c9698daf |
@@ -122,7 +122,7 @@ The per-backend prefix match only sees files under a backend's own directory, so
|
||||
|
||||
| Changed path | Rebuilds |
|
||||
|---|---|
|
||||
| `backend/backend.proto` | nothing if the edit is additive-only, otherwise everything (see below) |
|
||||
| `backend/backend.proto` | everything (all languages compile or copy it) |
|
||||
| `backend/Dockerfile.<x>` | the Linux entries whose `dockerfile:` names it |
|
||||
| `backend/python/common/` | Python, Linux + Darwin |
|
||||
| `scripts/build/package-gpu-libs.sh` | Python, Linux only |
|
||||
@@ -132,17 +132,6 @@ The per-backend prefix match only sees files under a backend's own directory, so
|
||||
|
||||
Deliberately excluded: `backend/index.yaml` (gallery metadata, never enters an image), `.github/backend-matrix.yml` (adding a backend would rebuild all of them), `backend/Dockerfile.base-grpc-builder` (owned by `base-images.yml`), and the root `Makefile` (touched in ~11% of commits, and its backend-relevant edits arrive alongside the backend directory anyway). `make test-ci-scripts` pins all of this.
|
||||
|
||||
#### `backend/backend.proto` is content-filtered, not path-filtered
|
||||
|
||||
Every language consumes the proto, so a path rule for it can only ever say "rebuild all 473 images". It changes in ~1.3% of commits, and that was enough to make it the single largest CI cost driver in the repo: on 2026-07-29 four runs totalling 935 queued jobs traced to nothing but a proto edit, one of which (#11158) was a six-line diff adding `bool cache_prompt = 8;`.
|
||||
|
||||
An additive proto edit cannot change how a backend that never references the new symbol behaves, so `filterMatrix()` suppresses the rule for one. `changed-backends.js` fetches `backend/backend.proto` at the base revision (same contents-API pattern as `.github/backend-matrix.yml`) and hands both texts to `protoChangeIsAdditive()`, which compares them structurally rather than textually:
|
||||
|
||||
- **Additive, rebuilds nothing**: a new field with an unused number, a new message, a new enum value, a new RPC. Comment, whitespace and ordering changes also land here.
|
||||
- **Breaking, rebuilds everything**: a removed, renumbered, retyped or renamed field, a dropped RPC, a changed `option` or `package`. So does an unresolvable base revision, matching the run-all posture used for a truncated diff.
|
||||
|
||||
Checked against every proto commit in the preceding six months, all nine resolvable ones classify as additive. Note the tradeoff this accepts: generated stubs do change for an additive edit, so image bytes would differ on a rebuild even though behavior does not. That is the same standard already applied when the filter declines to rebuild on unrelated `pkg/` changes, and the weekly cron remains the backstop.
|
||||
|
||||
The Sunday 06:00 UTC cron on `backend.yml` exists specifically because path filtering can leave Python backends frozen on stale wheels. `DEPS_REFRESH` (below) only fires when the build actually runs, so an untouched Python backend would never re-resolve its unpinned deps. The weekly cron is the safety net.
|
||||
|
||||
## The `DEPS_REFRESH` cache-buster (Python backends)
|
||||
|
||||
@@ -1,138 +0,0 @@
|
||||
# Working on the dllm Backend
|
||||
|
||||
`mudler/dllm.cpp` is a standalone C++/ggml engine for DiffusionGemma
|
||||
block-diffusion models. LocalAI wraps it with a **pure-Go** backend at
|
||||
`backend/go/dllm/` that dlopens `libdllm.so` via purego (ebitengine/purego) -
|
||||
NOT cgo, and NOT a C++ grpc-server fork. The Go side owns chat templating
|
||||
(gemma4 renderer) and output parsing (gemma4 streaming parser) and implements
|
||||
the rich gRPC interface (`PredictRich`/`PredictStreamRich`, ChatDelta replies).
|
||||
|
||||
> NOTE: github.com/mudler/dllm.cpp is still **private** (publishing is
|
||||
> planned). Until then the Makefile's anonymous clone fails; use the local-dev
|
||||
> symlink shortcut documented at the top of `backend/go/dllm/Makefile`
|
||||
> (symlink an out-of-tree `build/libdllm.so` into the backend dir and skip the
|
||||
> clone), or a git credential helper with repo access.
|
||||
|
||||
## Pin
|
||||
|
||||
`backend/go/dllm/Makefile` pins `DLLM_VERSION?=<sha>` at the top
|
||||
(whisper / parakeet-cpp / ds4 convention). The bump-deps bot
|
||||
(`.github/workflows/bump_deps.yaml`) tracks `mudler/dllm.cpp` `main` and
|
||||
rewrites that variable. After a manual bump: `make -C backend/go/dllm purge &&
|
||||
make -C backend/go/dllm` (the clone is keyed on the directory existing, not
|
||||
the sha).
|
||||
|
||||
## C-ABI and the serialization contract
|
||||
|
||||
The binding covers the 9-symbol flat C-ABI from dllm.cpp's
|
||||
`include/dllm_capi.h` (ABI v1; `main.go` hard-fails on a version mismatch):
|
||||
`abi_version, load, free, last_error, free_string, tokenize_json, generate,
|
||||
generate_stream, cancel`. Contract points the Go wiring encodes (`capi.go`
|
||||
header comment has the full list):
|
||||
|
||||
- **One ctx = one concurrent generate/tokenize.** A per-model worker
|
||||
goroutine (`Dllm.jobs` in `dllm.go`) owns ALL C calls, making the
|
||||
serialization structural instead of lock discipline.
|
||||
- **`dllm_capi_cancel` is the ONE exception**: it only flips an atomic and may
|
||||
be called from any goroutine mid-generate, so `Dllm.Cancel` bypasses the
|
||||
worker queue. The flag resets at the start of each generate, so a watchdog
|
||||
racing a new generate must re-issue cancel.
|
||||
- **`last_error` is a borrowed pointer** and must only be read AFTER the
|
||||
failing call returned (never while a generate is in flight on the same ctx).
|
||||
- **Free vs in-flight requests**: requests hold `genMu.RLock` for their full
|
||||
duration; `Free` takes the write lock, so it only runs when nothing is in
|
||||
flight, then drains and closes the worker. Post-Free requests get a clean
|
||||
"model not loaded" error.
|
||||
- `tokenize_json`/`generate` return malloc'd `char*` (bound as `uintptr`,
|
||||
copied, then `dllm_capi_free_string`d); opts/params JSON must be a FLAT
|
||||
object of scalars (`buildOptsJSON` rejects anything else).
|
||||
|
||||
## Wire shape
|
||||
|
||||
| RPC | Implementation |
|
||||
|---|---|
|
||||
| LoadModel | `dllm_capi_load` (params: `n_gpu_layers`, `n_threads`, `ctx_len`); `Options[]` parsed into per-request gen opts (`eb_*`, `blocks`, `kv_cache`) by `parseModelGenOpts` |
|
||||
| PredictRich | render (if templated) → `dllm_capi_generate` → parse → ONE Reply with aggregated ChatDeltas + legacy `Message` bytes |
|
||||
| PredictStreamRich | `dllm_capi_generate_stream`; per committed diffusion block → UTF-8 holdback → parser.Feed → one Reply per non-empty delta batch (channel closed by the CALLER, per `pkg/grpc/interface.go`) |
|
||||
| Predict / PredictStream | Legacy paths, delegate to the rich pair (legacy stream INVERTS channel ownership: the impl closes) |
|
||||
| TokenizeString | `dllm_capi_tokenize_json` (C side prepends BOS per `vocab.add_bos`) |
|
||||
| Cancel | `dllm_capi_cancel`, exposed as the `grpc.Cancellable` capability (`pkg/grpc/interface.go`): the gRPC server arms it via `context.AfterFunc` on the Predict/PredictStream context, so client disconnects/timeouts abort the in-flight generate - llama.cpp `IsCancelled()` parity for Go backends |
|
||||
|
||||
`n_threads` and `ctx_len` are accepted-but-ignored by the engine at the
|
||||
current pin (the context bound comes from GGUF `n_ctx_train`); they are sent
|
||||
for forward compatibility.
|
||||
|
||||
## Renderer / parser (the templated chat path)
|
||||
|
||||
With `use_tokenizer_template` + raw Messages, the backend owns templating and
|
||||
parsing (the ds4 precedent, but in Go):
|
||||
|
||||
- `gemma4_renderer.go` - `RenderGemma4(msgs, toolsJSON, enableThinking,
|
||||
addGenerationPrompt)`. The file embeds the FULL `tokenizer.chat_template`
|
||||
jinja (17466 bytes, md5 `8c34cf93c7a7815b3fdb300a009c4c17`) extracted
|
||||
verbatim from `diffusiongemma-26B-A4B-it-BF16.gguf` via gguf-py - e.g.
|
||||
`python scripts/dump_gguf.py model.gguf | grep -A400 chat_template` in the
|
||||
dllm.cpp checkout - as a numbered comment block; every Go rule cites its
|
||||
"tpl L<n>" line. Re-verify the md5 before blaming the renderer for a
|
||||
mismatch with a new GGUF. **BOS exception**: the template emits
|
||||
`{{- bos_token -}}` but the renderer deliberately does NOT - dllm.cpp's
|
||||
`run_generate` tokenizes with `prepend_bos = vocab.add_bos` (true for
|
||||
gemma4), so a literal `<bos>` would double it.
|
||||
- `gemma4_parser.go` - streaming state machine turning raw model text
|
||||
(fragments can split anywhere, including mid-marker) into ChatDeltas:
|
||||
thought channels → `reasoning_content`, `<|tool_call>call:name{...}` →
|
||||
ToolCallDelta, `<turn|>` → done. Marker grammar cross-checked against vLLM
|
||||
PR #45163's gemma4 tool/reasoning parsers. Malformed payloads are re-emitted
|
||||
raw as content, never dropped.
|
||||
- Thinking is **opt-in** for this family (`Metadata["enable_thinking"]`,
|
||||
default OFF - the inverse of ds4): the template gates every thinking branch
|
||||
on `enable_thinking`, and the no-thinking render pre-closes an empty thought
|
||||
channel, so the parser always starts in content state.
|
||||
- **UTF-8 boundary holdback** (`splitValidUTF8` in `dllm.go`): per-block
|
||||
detokenization can split a multi-byte character across block boundaries, and
|
||||
grpc-go refuses to marshal invalid UTF-8 in proto3 strings. An incomplete
|
||||
trailing sequence (at most 3 bytes) is carried into the next block; genuinely
|
||||
undecodable bytes become U+FFFD.
|
||||
|
||||
Without `use_tokenizer_template`, the prompt passes through verbatim and the
|
||||
output is NOT gemma4-parsed (plain content, like any non-autoparsing backend).
|
||||
|
||||
## Tests
|
||||
|
||||
| Layer | Gate | What |
|
||||
|---|---|---|
|
||||
| `backend/go/dllm/*_test.go` (renderer/parser/wiring) | none - run in plain `go test ./backend/go/dllm/...` | Ginkgo specs over a fake `generator` seam; canonical renderer fixtures from transformers' `test_modeling_diffusion_gemma.py`, parser tables from the vLLM gemma4 parsers |
|
||||
| `backend/go/dllm/dllm_test.go` C-ABI smoke | `DLLM_TEST_LIBRARY` + `DLLM_TEST_TINY_MODEL` (dllm.cpp's `tests/fixtures/tiny_with_vocab.gguf`); Skips when unset | Drives the real `libdllm.so`: ABI check, load, tokenize `[2,18]`, deterministic generate, cancel (incl. mid-stream `Dllm.Cancel` aborting a deliberately slow `eb_max_steps:256` run in ~10ms) |
|
||||
| `tests/e2e-backends/dllm_test.go` | `BACKEND_TEST_DLLM=1` + `BACKEND_BINARY` (packaged run.sh) + `BACKEND_TEST_MODEL_FILE` (tiny fixture) | Templated chat round trip (Messages + UseTokenizerTemplate) over the real gRPC binary, non-streaming + streaming; plus client-context cancellation mid-stream (proves the `Cancellable` server plumbing end to end) |
|
||||
| Real-model e2e | `BACKEND_TEST_DLLM_REAL_MODEL_FILE` (26B BF16, ~50 GB) + `BACKEND_TEST_DLLM_REAL_GPU_LAYERS` | CUDA-13-class hardware only |
|
||||
|
||||
Tool-call e2e is deliberately absent from the tiny-model spec: the fixture has
|
||||
random weights and cannot be coaxed into emitting tool markup; the unit tables
|
||||
carry that coverage.
|
||||
|
||||
## Build matrix
|
||||
|
||||
`cpu-dllm` (amd64 + arm64), `cuda13-dllm` (amd64), and
|
||||
`cuda13-nvidia-l4t-arm64-dllm` (arm64 CUDA: Jetson / DGX Spark GB10), via
|
||||
`.github/backend-matrix.yml`. No darwin/Metal. CUDA builds forward
|
||||
`-DDLLM_CUDA=ON` (dllm.cpp gates ggml's CUDA behind its own flag - a bare
|
||||
`-DGGML_CUDA=ON` is overridden by the cache FORCE). `libdllm.so` is
|
||||
self-contained (ggml statically absorbed, PIC), so `package.sh` only ships
|
||||
the binary, `run.sh` and that one .so (the parakeet-cpp-style stub layout;
|
||||
no ldd walk yet).
|
||||
|
||||
## Known limitations
|
||||
|
||||
- **Cancel granularity**: the C-ABI cancel flag is per-ctx and resets on
|
||||
every generate entry, so a Cancel racing a NEW generate can be lost, and
|
||||
with requests queued on the worker it aborts whichever generate is
|
||||
currently running (acceptable: the server de-registers the hook on normal
|
||||
completion, one process serves one model).
|
||||
- **Throughput**: ~0.15 tok/s on the 26B at default settings (GB10) - every
|
||||
denoise step recomputes the full prompt+canvas. The upstream prefix-KV
|
||||
cache (dllm.cpp P3) is the fix; `kv_cache:on` errors until it lands
|
||||
(`auto`/`off` are accepted no-ops).
|
||||
- **Repo privacy**: see the note at the top - CI clone of dllm.cpp needs the
|
||||
repo published (or credentials) before the backend images can build.
|
||||
- Engine spec/validation references: dllm.cpp `docs/validation.md` and
|
||||
LocalAI `docs/superpowers/specs/2026-06-10-dllm-cpp-design.md`.
|
||||
@@ -28,10 +28,6 @@ if [ -z "${BUILD_TYPE:-}" ]; then
|
||||
# variants with it (the host never *selects* SME unless it has it, but every variant must
|
||||
# still compile).
|
||||
if [ "${TARGETARCH}" = "arm64" ]; then
|
||||
# The prebuilt base inherits default ports.ubuntu.com sources; honor the
|
||||
# APT_*_MIRROR build args here like the from-source path does, so this
|
||||
# apt step survives a mirror outage.
|
||||
sh /LocalAI/.docker/apt-mirror.sh || true
|
||||
apt-get update -qq && apt-get install -y -qq gcc-14 g++-14
|
||||
export CC=gcc-14 CXX=g++-14
|
||||
fi
|
||||
|
||||
200
.github/backend-matrix.yml
vendored
200
.github/backend-matrix.yml
vendored
@@ -66,34 +66,6 @@ include:
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
platform-tag: 'amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-kokoro'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'true'
|
||||
backend: "kokoro"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/arm64'
|
||||
platform-tag: 'arm64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-kokoro'
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'true'
|
||||
backend: "kokoro"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
@@ -756,19 +728,6 @@ include:
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "8"
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-cuda-12-trellis2cpp'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "trellis2cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "8"
|
||||
@@ -1729,19 +1688,6 @@ include:
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-cuda-13-trellis2cpp'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "trellis2cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
@@ -1755,19 +1701,6 @@ include:
|
||||
backend: "stablediffusion-ggml"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/arm64'
|
||||
skip-drivers: 'false'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-nvidia-l4t-cuda-13-arm64-trellis2cpp'
|
||||
base-image: "ubuntu:24.04"
|
||||
ubuntu-version: '2404'
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
backend: "trellis2cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
@@ -1911,19 +1844,6 @@ include:
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-cuda-13-dllm'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "dllm"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
@@ -1937,19 +1857,6 @@ include:
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/arm64'
|
||||
skip-drivers: 'false'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-nvidia-l4t-cuda-13-arm64-dllm'
|
||||
base-image: "ubuntu:24.04"
|
||||
ubuntu-version: '2404'
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
backend: "dllm"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
@@ -3332,35 +3239,6 @@ include:
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
# trellis2cpp
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
platform-tag: 'amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-trellis2cpp'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "trellis2cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/arm64'
|
||||
platform-tag: 'arm64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-trellis2cpp'
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "trellis2cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
# sam3-cpp
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
@@ -3686,34 +3564,6 @@ include:
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'vulkan'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
platform-tag: 'amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-vulkan-trellis2cpp'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "trellis2cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'vulkan'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/arm64'
|
||||
platform-tag: 'arm64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-vulkan-trellis2cpp'
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "trellis2cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "0"
|
||||
@@ -3727,19 +3577,6 @@ include:
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2204'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/arm64'
|
||||
skip-drivers: 'false'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-nvidia-l4t-arm64-trellis2cpp'
|
||||
base-image: "nvcr.io/nvidia/l4t-jetpack:r36.4.0"
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
backend: "trellis2cpp"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2204'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "0"
|
||||
@@ -5571,35 +5408,6 @@ include:
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
# valkey-store
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
platform-tag: 'amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-valkey-store'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "valkey-store"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/arm64'
|
||||
platform-tag: 'arm64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-valkey-store'
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "valkey-store"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
# rfdetr
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
@@ -6141,10 +5949,6 @@ includeDarwin:
|
||||
tag-suffix: "-metal-darwin-arm64-stablediffusion-ggml"
|
||||
build-type: "metal"
|
||||
lang: "go"
|
||||
- backend: "trellis2cpp"
|
||||
tag-suffix: "-metal-darwin-arm64-trellis2cpp"
|
||||
build-type: "metal"
|
||||
lang: "go"
|
||||
- backend: "whisper"
|
||||
tag-suffix: "-metal-darwin-arm64-whisper"
|
||||
build-type: "metal"
|
||||
@@ -6322,10 +6126,6 @@ includeDarwin:
|
||||
tag-suffix: "-metal-darwin-arm64-cloud-proxy"
|
||||
build-type: "metal"
|
||||
lang: "go"
|
||||
- backend: "valkey-store"
|
||||
tag-suffix: "-metal-darwin-arm64-valkey-store"
|
||||
build-type: "metal"
|
||||
lang: "go"
|
||||
- backend: "llama-cpp-quantization"
|
||||
tag-suffix: "-metal-darwin-arm64-llama-cpp-quantization"
|
||||
build-type: "mps"
|
||||
|
||||
8
.github/workflows/bump_deps.yaml
vendored
8
.github/workflows/bump_deps.yaml
vendored
@@ -50,10 +50,6 @@ jobs:
|
||||
variable: "PARAKEET_VERSION"
|
||||
branch: "master"
|
||||
file: "backend/go/parakeet-cpp/Makefile"
|
||||
- repository: "mudler/dllm.cpp"
|
||||
variable: "DLLM_VERSION"
|
||||
branch: "main"
|
||||
file: "backend/go/dllm/Makefile"
|
||||
- repository: "mudler/vllm.cpp"
|
||||
variable: "VLLM_CPP_VERSION"
|
||||
branch: "main"
|
||||
@@ -82,10 +78,6 @@ jobs:
|
||||
variable: "STABLEDIFFUSION_GGML_VERSION"
|
||||
branch: "master"
|
||||
file: "backend/go/stablediffusion-ggml/Makefile"
|
||||
- repository: "localai-org/trellis2cpp"
|
||||
variable: "TRELLIS2CPP_VERSION"
|
||||
branch: "pbr-textures"
|
||||
file: "backend/go/trellis2cpp/Makefile"
|
||||
- repository: "mudler/go-piper"
|
||||
variable: "PIPER_VERSION"
|
||||
branch: "master"
|
||||
|
||||
36
.github/workflows/test-extra.yml
vendored
36
.github/workflows/test-extra.yml
vendored
@@ -38,7 +38,6 @@ jobs:
|
||||
acestep-cpp: ${{ steps.detect.outputs.acestep-cpp }}
|
||||
qwen3-tts-cpp: ${{ steps.detect.outputs.qwen3-tts-cpp }}
|
||||
magpie-tts-cpp: ${{ steps.detect.outputs.magpie-tts-cpp }}
|
||||
trellis2cpp: ${{ steps.detect.outputs.trellis2cpp }}
|
||||
rfdetr-cpp: ${{ steps.detect.outputs.rfdetr-cpp }}
|
||||
locate-anything-cpp: ${{ steps.detect.outputs.locate-anything-cpp }}
|
||||
vibevoice-cpp: ${{ steps.detect.outputs.vibevoice-cpp }}
|
||||
@@ -936,41 +935,6 @@ jobs:
|
||||
- name: Test rfdetr-cpp
|
||||
run: |
|
||||
make --jobs=5 --output-sync=target -C backend/go/rfdetr-cpp test
|
||||
# Weight-free packaged-backend smoke for trellis2cpp. Starting run.sh loads
|
||||
# libtrellis2 + ggml, resolves the complete C ABI (including remeshing), and
|
||||
# answers gRPC Health without downloading or loading the multi-GB model set.
|
||||
tests-trellis2cpp:
|
||||
needs: detect-changes
|
||||
if: needs.detect-changes.outputs.trellis2cpp == 'true' || needs.detect-changes.outputs.run-all == 'true'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 90
|
||||
steps:
|
||||
- name: Clone
|
||||
uses: actions/checkout@v7
|
||||
with:
|
||||
submodules: true
|
||||
- name: Dependencies
|
||||
run: |
|
||||
sudo apt-get update
|
||||
sudo apt-get install -y build-essential cmake curl unzip
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@v5
|
||||
- name: Display Go version
|
||||
run: go version
|
||||
- name: Proto Dependencies
|
||||
run: |
|
||||
curl -L -s https://github.com/protocolbuffers/protobuf/releases/download/v26.1/protoc-26.1-linux-x86_64.zip -o protoc.zip && \
|
||||
unzip -j -d /usr/local/bin protoc.zip bin/protoc && \
|
||||
rm protoc.zip
|
||||
go install google.golang.org/protobuf/cmd/protoc-gen-go@v1.34.2
|
||||
go install google.golang.org/grpc/cmd/protoc-gen-go-grpc@1958fcbe2ca8bd93af633f11e97d44e567e945af
|
||||
PATH="$PATH:$HOME/go/bin" make protogen-go
|
||||
- name: Build trellis2cpp
|
||||
run: |
|
||||
make --jobs=5 --output-sync=target -C backend/go/trellis2cpp
|
||||
- name: Test trellis2cpp
|
||||
run: |
|
||||
make --jobs=5 --output-sync=target -C backend/go/trellis2cpp test
|
||||
# Per-backend e2e for locate-anything-cpp: builds the .so + Go binary and
|
||||
# runs `make -C backend/go/locate-anything-cpp test`. test.sh fetches the
|
||||
# locate-anything-q8_0 GGUF (~6.3 GB, NVIDIA LocateAnything-3B) from the
|
||||
|
||||
1
.gitignore
vendored
1
.gitignore
vendored
@@ -30,7 +30,6 @@ LocalAI
|
||||
# Go backend packages whose main lives under backend/go/.
|
||||
/cloud-proxy
|
||||
/local-store
|
||||
/valkey-store
|
||||
# prevent above rules from omitting the helm chart
|
||||
!charts/*
|
||||
# prevent above rules from omitting the api/localai folder
|
||||
|
||||
@@ -26,7 +26,6 @@ LocalAI follows the Linux kernel project's [guidelines for AI coding assistants]
|
||||
| [.agents/vllm-backend.md](.agents/vllm-backend.md) | Working on the vLLM / vLLM-omni backends — native parsers, ChatDelta, CPU build, libnuma packaging, backend hooks |
|
||||
| [.agents/sglang-backend.md](.agents/sglang-backend.md) | Working on the SGLang backend — `engine_args` validation against ServerArgs, speculative-decoding (EAGLE/EAGLE3/DFLASH/MTP) recipes, parser handling |
|
||||
| [.agents/ds4-backend.md](.agents/ds4-backend.md) | Working on the ds4 backend - DSML state machine, thinking modes, KV cache, Metal+CUDA matrix |
|
||||
| [.agents/dllm-backend.md](.agents/dllm-backend.md) | Working on the dllm backend (DiffusionGemma block-diffusion) - purego C-ABI binding, per-ctx serialization contract, gemma4 renderer/parser, gated test layers |
|
||||
| [.agents/testing-mcp-apps.md](.agents/testing-mcp-apps.md) | Testing MCP Apps (interactive tool UIs) in the React UI |
|
||||
| [.agents/api-endpoints-and-auth.md](.agents/api-endpoints-and-auth.md) | Adding API endpoints, auth middleware, feature permissions, user access control |
|
||||
| [.agents/debugging-backends.md](.agents/debugging-backends.md) | Debugging runtime backend failures, dependency conflicts, rebuilding backends |
|
||||
|
||||
31
Makefile
31
Makefile
@@ -1,5 +1,5 @@
|
||||
# Disable parallel execution for backend builds
|
||||
.NOTPARALLEL: backends/diffusers backends/llama-cpp backends/turboquant backends/bonsai backends/outetts backends/piper backends/stablediffusion-ggml backends/trellis2cpp backends/trellis2cpp-darwin backends/whisper backends/crispasr backends/parakeet-cpp backends/dllm backends/moss-transcribe-cpp backends/faster-whisper backends/silero-vad backends/local-store backends/valkey-store backends/cloud-proxy backends/huggingface backends/rfdetr backends/rfdetr-cpp backends/insightface backends/speaker-recognition backends/kitten-tts backends/kokoro backends/chatterbox backends/llama-cpp-darwin backends/neutts build-darwin-python-backend build-darwin-go-backend backends/mlx backends/diffuser-darwin backends/mlx-vlm backends/mlx-audio backends/mlx-distributed backends/stablediffusion-ggml-darwin backends/vllm backends/vllm-omni backends/longcat-video backends/sglang backends/moonshine backends/pocket-tts backends/qwen-tts backends/faster-qwen3-tts backends/qwen-asr backends/nemo backends/voxcpm backends/whisperx backends/ace-step backends/acestep-cpp backends/fish-speech backends/voxtral backends/opus backends/trl backends/llama-cpp-quantization backends/kokoros backends/sam3-cpp backends/qwen3-tts-cpp backends/moss-tts-cpp backends/magpie-tts-cpp backends/vllm-cpp backends/omnivoice-cpp backends/vibevoice-cpp backends/localvqe backends/tinygrad backends/sherpa-onnx backends/ds4 backends/ds4-darwin backends/liquid-audio backends/supertonic backends/depth-anything-cpp backends/privacy-filter backends/privacy-filter-darwin
|
||||
.NOTPARALLEL: backends/diffusers backends/llama-cpp backends/turboquant backends/bonsai backends/outetts backends/piper backends/stablediffusion-ggml backends/whisper backends/crispasr backends/parakeet-cpp backends/moss-transcribe-cpp backends/faster-whisper backends/silero-vad backends/local-store backends/cloud-proxy backends/huggingface backends/rfdetr backends/rfdetr-cpp backends/insightface backends/speaker-recognition backends/kitten-tts backends/kokoro backends/chatterbox backends/llama-cpp-darwin backends/neutts build-darwin-python-backend build-darwin-go-backend backends/mlx backends/diffuser-darwin backends/mlx-vlm backends/mlx-audio backends/mlx-distributed backends/stablediffusion-ggml-darwin backends/vllm backends/vllm-omni backends/longcat-video backends/sglang backends/moonshine backends/pocket-tts backends/qwen-tts backends/faster-qwen3-tts backends/qwen-asr backends/nemo backends/voxcpm backends/whisperx backends/ace-step backends/acestep-cpp backends/fish-speech backends/voxtral backends/opus backends/trl backends/llama-cpp-quantization backends/kokoros backends/sam3-cpp backends/qwen3-tts-cpp backends/moss-tts-cpp backends/magpie-tts-cpp backends/vllm-cpp backends/omnivoice-cpp backends/vibevoice-cpp backends/localvqe backends/tinygrad backends/sherpa-onnx backends/ds4 backends/ds4-darwin backends/liquid-audio backends/supertonic backends/depth-anything-cpp backends/privacy-filter backends/privacy-filter-darwin
|
||||
|
||||
GOCMD=go
|
||||
GOTEST=$(GOCMD) test
|
||||
@@ -69,7 +69,7 @@ else
|
||||
GORELEASER=$(shell which goreleaser)
|
||||
endif
|
||||
|
||||
TEST_PATHS?=./api/... ./pkg/... ./core/... ./backend/go/cloud-proxy/... ./backend/go/local-store/... ./backend/go/valkey-store/...
|
||||
TEST_PATHS?=./api/... ./pkg/... ./core/... ./backend/go/cloud-proxy/... ./backend/go/local-store/...
|
||||
|
||||
## Coverage output and the committed baseline that CI compares against.
|
||||
## The gate is strict: total coverage must never decrease (no tolerance).
|
||||
@@ -386,15 +386,6 @@ test-stores: backends/local-store
|
||||
BACKENDS_PATH=$(abspath ./)/backends \
|
||||
$(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --flake-attempts $(TEST_FLAKES) -v -r tests/integration
|
||||
|
||||
## Valkey-backed vector-store integration. Requires a running Valkey Search
|
||||
## server (valkey/valkey-bundle:9.1.0) reachable at $$VALKEY_ADDR — the suite
|
||||
## skips itself when VALKEY_ADDR is unset. Builds the backend on demand and
|
||||
## points the model loader at it via BACKENDS_PATH. Label-filtered to the
|
||||
## valkey specs so it does not also run the in-memory local-store suite.
|
||||
test-valkey-store: backends/valkey-store
|
||||
BACKENDS_PATH=$(abspath ./)/backends \
|
||||
$(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --flake-attempts $(TEST_FLAKES) --label-filter='valkey' -v -r tests/integration
|
||||
|
||||
test-opus:
|
||||
@echo 'Running opus backend tests'
|
||||
$(MAKE) -C backend/go/opus libopusshim.so
|
||||
@@ -603,8 +594,6 @@ prepare-test-extra: protogen-python
|
||||
$(MAKE) -C backend/rust/kokoros kokoros-grpc
|
||||
$(MAKE) -C backend/go/rfdetr-cpp
|
||||
$(MAKE) -C backend/go/locate-anything-cpp
|
||||
$(MAKE) -C backend/go/trellis2cpp
|
||||
$(MAKE) -C backend/go/valkey-store
|
||||
|
||||
test-extra: prepare-test-extra
|
||||
$(MAKE) -C backend/python/transformers test
|
||||
@@ -637,8 +626,6 @@ test-extra: prepare-test-extra
|
||||
$(MAKE) -C backend/go/depth-anything-cpp test
|
||||
$(MAKE) -C backend/go/supertonic test
|
||||
$(MAKE) -C backend/go/vllm-cpp test
|
||||
$(MAKE) -C backend/go/trellis2cpp test
|
||||
$(MAKE) -C backend/go/valkey-store test
|
||||
|
||||
##
|
||||
## End-to-end gRPC tests that exercise a built backend container image.
|
||||
@@ -1231,10 +1218,6 @@ backends/stablediffusion-ggml-darwin:
|
||||
BACKEND=stablediffusion-ggml BUILD_TYPE=metal $(MAKE) build-darwin-go-backend
|
||||
./local-ai backends install "ocifile://$(abspath ./backend-images/stablediffusion-ggml.tar)"
|
||||
|
||||
backends/trellis2cpp-darwin:
|
||||
BACKEND=trellis2cpp BUILD_TYPE=metal $(MAKE) build-darwin-go-backend
|
||||
./local-ai backends install "ocifile://$(abspath ./backend-images/trellis2cpp.tar)"
|
||||
|
||||
backend-images:
|
||||
mkdir -p backend-images
|
||||
|
||||
@@ -1262,18 +1245,13 @@ BACKEND_PRIVACY_FILTER = privacy-filter|privacy-filter|.|false|false
|
||||
# Golang backends
|
||||
BACKEND_PIPER = piper|golang|.|false|true
|
||||
BACKEND_LOCAL_STORE = local-store|golang|.|false|true
|
||||
BACKEND_VALKEY_STORE = valkey-store|golang|.|false|true
|
||||
BACKEND_CLOUD_PROXY = cloud-proxy|golang|.|false|true
|
||||
BACKEND_HUGGINGFACE = huggingface|golang|.|false|true
|
||||
BACKEND_SILERO_VAD = silero-vad|golang|.|false|true
|
||||
BACKEND_STABLEDIFFUSION_GGML = stablediffusion-ggml|golang|.|--progress=plain|true
|
||||
BACKEND_TRELLIS2CPP = trellis2cpp|golang|.|--progress=plain|true
|
||||
BACKEND_WHISPER = whisper|golang|.|false|true
|
||||
BACKEND_CRISPASR = crispasr|golang|.|false|true
|
||||
BACKEND_PARAKEET_CPP = parakeet-cpp|golang|.|false|true
|
||||
# dllm is mudler/dllm.cpp, the DiffusionGemma block-diffusion engine,
|
||||
# wrapped by the purego backend at backend/go/dllm.
|
||||
BACKEND_DLLM = dllm|golang|.|false|true
|
||||
BACKEND_MOSS_TRANSCRIBE_CPP = moss-transcribe-cpp|golang|.|false|true
|
||||
BACKEND_DEPTH_ANYTHING_CPP = depth-anything-cpp|golang|.|false|true
|
||||
BACKEND_VOXTRAL = voxtral|golang|.|false|true
|
||||
@@ -1366,16 +1344,13 @@ $(eval $(call generate-docker-build-target,$(BACKEND_DS4)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_PRIVACY_FILTER)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_PIPER)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_LOCAL_STORE)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_VALKEY_STORE)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_CLOUD_PROXY)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_HUGGINGFACE)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_SILERO_VAD)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_STABLEDIFFUSION_GGML)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_TRELLIS2CPP)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_WHISPER)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_CRISPASR)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_PARAKEET_CPP)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_DLLM)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_MOSS_TRANSCRIBE_CPP)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_DEPTH_ANYTHING_CPP)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_VOXTRAL)))
|
||||
@@ -1433,7 +1408,7 @@ $(eval $(call generate-docker-build-target,$(BACKEND_SUPERTONIC)))
|
||||
docker-save-%: backend-images
|
||||
docker save local-ai-backend:$* -o backend-images/$*.tar
|
||||
|
||||
docker-build-backends: docker-build-llama-cpp docker-build-ik-llama-cpp docker-build-turboquant docker-build-bonsai docker-build-ds4 docker-build-rerankers docker-build-vllm docker-build-vllm-omni docker-build-longcat-video docker-build-sglang docker-build-transformers docker-build-outetts docker-build-diffusers docker-build-kokoro docker-build-faster-whisper docker-build-crispasr docker-build-coqui docker-build-chatterbox docker-build-vibevoice docker-build-liquid-audio docker-build-moonshine docker-build-pocket-tts docker-build-qwen-tts docker-build-fish-speech docker-build-faster-qwen3-tts docker-build-qwen-asr docker-build-nemo docker-build-voxcpm docker-build-whisperx docker-build-ace-step docker-build-acestep-cpp docker-build-voxtral docker-build-mlx-distributed docker-build-trl docker-build-llama-cpp-quantization docker-build-tinygrad docker-build-kokoros docker-build-sam3-cpp docker-build-rfdetr-cpp docker-build-qwen3-tts-cpp docker-build-moss-tts-cpp docker-build-magpie-tts-cpp docker-build-vllm-cpp docker-build-omnivoice-cpp docker-build-vibevoice-cpp docker-build-localvqe docker-build-insightface docker-build-speaker-recognition docker-build-sherpa-onnx docker-build-cloud-proxy docker-build-supertonic docker-build-depth-anything-cpp docker-build-moss-transcribe-cpp docker-build-privacy-filter docker-build-trellis2cpp docker-build-valkey-store
|
||||
docker-build-backends: docker-build-llama-cpp docker-build-ik-llama-cpp docker-build-turboquant docker-build-bonsai docker-build-ds4 docker-build-rerankers docker-build-vllm docker-build-vllm-omni docker-build-longcat-video docker-build-sglang docker-build-transformers docker-build-outetts docker-build-diffusers docker-build-kokoro docker-build-faster-whisper docker-build-crispasr docker-build-coqui docker-build-chatterbox docker-build-vibevoice docker-build-liquid-audio docker-build-moonshine docker-build-pocket-tts docker-build-qwen-tts docker-build-fish-speech docker-build-faster-qwen3-tts docker-build-qwen-asr docker-build-nemo docker-build-voxcpm docker-build-whisperx docker-build-ace-step docker-build-acestep-cpp docker-build-voxtral docker-build-mlx-distributed docker-build-trl docker-build-llama-cpp-quantization docker-build-tinygrad docker-build-kokoros docker-build-sam3-cpp docker-build-rfdetr-cpp docker-build-qwen3-tts-cpp docker-build-moss-tts-cpp docker-build-magpie-tts-cpp docker-build-vllm-cpp docker-build-omnivoice-cpp docker-build-vibevoice-cpp docker-build-localvqe docker-build-insightface docker-build-speaker-recognition docker-build-sherpa-onnx docker-build-cloud-proxy docker-build-supertonic docker-build-depth-anything-cpp docker-build-moss-transcribe-cpp docker-build-privacy-filter
|
||||
|
||||
########################################################
|
||||
### Mock Backend for E2E Tests
|
||||
|
||||
@@ -245,7 +245,6 @@ Most backends wrap a best-in-class upstream engine. A handful of them are native
|
||||
| [depth-anything.cpp](https://github.com/mudler/depth-anything.cpp) | Depth Anything 3 monocular metric depth + camera pose estimation |
|
||||
| [face-detect.cpp](https://github.com/mudler/face-detect.cpp) | Face detection, recognition, demographics and anti-spoofing (SCRFD/ArcFace, YuNet/SFace), replacing the Python insightface backend |
|
||||
| [free-splatter.cpp](https://github.com/localai-org/free-splatter.cpp) | Pose-free 3D reconstruction (FreeSplatter): turns a handful of plain photos into 3D Gaussians, no camera poses or GPU required |
|
||||
| [trellis2.cpp](https://github.com/localai-org/trellis2cpp) | C++/GGML port of Microsoft TRELLIS.2: single-image to textured 3D mesh (GLB with PBR materials) |
|
||||
| [privacy-filter.cpp](https://github.com/localai-org/privacy-filter.cpp) | Standalone GGML PII/NER token-classification engine powering LocalAI's PII redaction tier |
|
||||
| [LocalVQE](https://github.com/localai-org/LocalVQE) | Joint acoustic echo cancellation, noise suppression, and dereverberation |
|
||||
| [local-store](https://github.com/mudler/LocalAI) | Local-first vector database for embeddings (shipped in-tree) |
|
||||
|
||||
@@ -111,10 +111,6 @@ RUN make -BC /LocalAI/backend/cpp/llama-cpp package
|
||||
# ============================================================================
|
||||
FROM ${BUILDER_BASE_IMAGE} AS builder-prebuilt
|
||||
|
||||
ARG APT_MIRROR
|
||||
ENV APT_MIRROR=${APT_MIRROR}
|
||||
ARG APT_PORTS_MIRROR
|
||||
ENV APT_PORTS_MIRROR=${APT_PORTS_MIRROR}
|
||||
ARG BUILD_TYPE
|
||||
ENV BUILD_TYPE=${BUILD_TYPE}
|
||||
ARG CUDA_DOCKER_ARCH
|
||||
|
||||
@@ -56,7 +56,6 @@ The backend system provides language-specific Dockerfiles that handle the build
|
||||
- **stablediffusion-ggml**: Stable Diffusion in Go with GGML Cpp backend
|
||||
- **piper**: Text-to-speech synthesis Golang with C bindings using rhaspy/piper
|
||||
- **local-store**: Vector storage backend
|
||||
- **valkey-store**: Durable vector storage backend backed by Valkey Search (FT.*)
|
||||
|
||||
#### C++ Backends (`cpp/`)
|
||||
- **llama-cpp**: Llama.cpp integration
|
||||
|
||||
@@ -16,7 +16,6 @@ service Backend {
|
||||
rpc Embedding(PredictOptions) returns (EmbeddingResult) {}
|
||||
rpc GenerateImage(GenerateImageRequest) returns (Result) {}
|
||||
rpc GenerateVideo(GenerateVideoRequest) returns (Result) {}
|
||||
rpc Generate3D(Generate3DRequest) returns (Result) {}
|
||||
rpc AudioTranscription(TranscriptRequest) returns (TranscriptResult) {}
|
||||
rpc AudioTranscriptionStream(TranscriptRequest) returns (stream TranscriptStreamResponse) {}
|
||||
// AudioTranscriptionLive is the bidirectional live-microphone ASR RPC. The
|
||||
@@ -182,13 +181,6 @@ message ScoreRequest {
|
||||
// PredictOptions.ModelIdentity for the full rationale. Empty means "no
|
||||
// identity supplied" and backends MUST skip the check.
|
||||
string ModelIdentity = 5;
|
||||
// Byte length of the prompt prefix that stays identical across
|
||||
// repeated scoring calls (e.g. a classifier's option-list system
|
||||
// prompt — everything before the per-turn probe text). Backends that
|
||||
// snapshot state (hybrid/recurrent models cannot rewind otherwise)
|
||||
// use it to place a reuse point exactly at the boundary, so the next
|
||||
// call re-processes only the tokens after it. 0 means unknown.
|
||||
int32 stable_prefix_len = 6;
|
||||
}
|
||||
|
||||
// CandidateScore is one row in the ScoreResponse, matching by index
|
||||
@@ -501,11 +493,6 @@ message ModelOptions {
|
||||
// Proxy carries the cloud-proxy backend's per-model configuration.
|
||||
// Empty for non-proxy backends.
|
||||
ProxyOptions Proxy = 74;
|
||||
|
||||
// EnableScore reserves backend resources for the Score RPC. It is derived
|
||||
// from the model's explicit `known_usecases: [score]` declaration so models
|
||||
// that never score retain their ordinary serving footprint.
|
||||
bool EnableScore = 75;
|
||||
}
|
||||
|
||||
// ProxyOptions configures the cloud-proxy backend. UpstreamURL and
|
||||
@@ -521,12 +508,6 @@ message ProxyOptions {
|
||||
string api_key_file = 5;
|
||||
string upstream_model = 6;
|
||||
int32 request_timeout_seconds = 7;
|
||||
// cache_prompt enables automatic Anthropic prompt-cache breakpoints
|
||||
// (cache_control: ephemeral) on the stable prefix — system, tools, and
|
||||
// the last message block — when translating to the Anthropic provider.
|
||||
// Cuts input cost on repeated/agentic calls (cache read = 0.1x). Only
|
||||
// meaningful for mode=translate + provider=anthropic; ignored otherwise.
|
||||
bool cache_prompt = 8;
|
||||
}
|
||||
|
||||
message Result {
|
||||
@@ -659,20 +640,6 @@ message GenerateVideoRequest {
|
||||
string ModelIdentity = 15;
|
||||
}
|
||||
|
||||
message Generate3DRequest {
|
||||
string src = 1; // Path to the staged conditioning image (3D generation is image-conditioned)
|
||||
string dst = 2; // Output path for the generated binary glTF (.glb) asset
|
||||
int32 seed = 3; // <=0 lets the backend pick a random seed
|
||||
int32 step = 4; // Flow sampling steps; <=0 uses the backend default
|
||||
float cfg_scale = 5; // Classifier-free guidance scale; <=0 uses the backend default
|
||||
int32 texture_steps = 6; // Texture flow sampling steps; <=0 uses the backend default
|
||||
string quality = 7; // Mesh pipeline: ""|"auto"|"coarse"|"512"|"1024"
|
||||
string background = 8; // Conditioning-image background handling: ""|"auto"|"keep"|"black"|"white"
|
||||
// Backend-specific per-request generation parameters. Values are strings
|
||||
// and are validated/coerced by the selected backend.
|
||||
map<string, string> params = 9;
|
||||
}
|
||||
|
||||
message TTSRequest {
|
||||
string text = 1;
|
||||
string model = 2;
|
||||
|
||||
@@ -41,7 +41,6 @@ define bonsai-build
|
||||
# and are applied by apply-patches.sh below.
|
||||
rm -rf $(CURRENT_MAKEFILE_DIR)/../bonsai-$(1)-build/patches
|
||||
$(MAKE) -C $(CURRENT_MAKEFILE_DIR)/../bonsai-$(1)-build purge
|
||||
bash $(LLAMA_CPP_DIR)/disable-score-task.sh $(CURRENT_MAKEFILE_DIR)/../bonsai-$(1)-build/grpc-server.cpp
|
||||
$(info $(GREEN)I bonsai build info:$(1)$(RESET))
|
||||
LLAMA_REPO=$(LLAMA_REPO) LLAMA_VERSION=$(BONSAI_VERSION) \
|
||||
$(MAKE) -C $(CURRENT_MAKEFILE_DIR)/../bonsai-$(1)-build llama.cpp
|
||||
@@ -78,7 +77,6 @@ bonsai-cpu-all:
|
||||
# and are applied by apply-patches.sh below.
|
||||
rm -rf $(CURRENT_MAKEFILE_DIR)/../bonsai-cpu-all-build/patches
|
||||
$(MAKE) -C $(CURRENT_MAKEFILE_DIR)/../bonsai-cpu-all-build purge
|
||||
bash $(LLAMA_CPP_DIR)/disable-score-task.sh $(CURRENT_MAKEFILE_DIR)/../bonsai-cpu-all-build/grpc-server.cpp
|
||||
$(info $(GREEN)I bonsai build info:cpu-all-variants$(RESET))
|
||||
LLAMA_REPO=$(LLAMA_REPO) LLAMA_VERSION=$(BONSAI_VERSION) \
|
||||
$(MAKE) -C $(CURRENT_MAKEFILE_DIR)/../bonsai-cpu-all-build llama.cpp
|
||||
|
||||
@@ -1,10 +1,10 @@
|
||||
# ds4 backend Makefile.
|
||||
#
|
||||
# Upstream pin lives below as DS4_VERSION?=54b36ed9ba42da31b24f2d1a5feb075c2475dbb1
|
||||
# Upstream pin lives below as DS4_VERSION?=0a7ad776b9068348e6cb09df8cafa9cadd285298
|
||||
# (.github/bump_deps.sh) can find and update it - matches the
|
||||
# llama-cpp / ik-llama-cpp / turboquant convention.
|
||||
|
||||
DS4_VERSION?=54b36ed9ba42da31b24f2d1a5feb075c2475dbb1
|
||||
DS4_VERSION?=0a7ad776b9068348e6cb09df8cafa9cadd285298
|
||||
DS4_REPO?=https://github.com/antirez/ds4
|
||||
|
||||
CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
|
||||
IK_LLAMA_VERSION?=b054a8b983827c01aec59d4dc273a27c492c51c4
|
||||
IK_LLAMA_VERSION?=0a4e10c7fb65d2dd5a4afb78339c7d373a8cdfaa
|
||||
LLAMA_REPO?=https://github.com/ikawrakow/ik_llama.cpp
|
||||
|
||||
CMAKE_ARGS?=
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
|
||||
LLAMA_VERSION?=1cbfd1988311775425d36c0ce066590f7d3049cf
|
||||
LLAMA_VERSION?=0d47ea7427463093e69128bf2c2f9cd06b3ee5b3
|
||||
LLAMA_REPO?=https://github.com/ggerganov/llama.cpp
|
||||
|
||||
CMAKE_ARGS?=
|
||||
|
||||
@@ -1,43 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Mark a copied gRPC server as targeting a llama.cpp fork that does not carry
|
||||
# LocalAI's slot-based Score patches. The RPC remains present in the shared
|
||||
# protobuf service, but responds with UNIMPLEMENTED instead of referencing
|
||||
# server task types and common_params fields absent from those forks.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
if [[ $# -ne 1 ]]; then
|
||||
echo "usage: $0 <grpc-server.cpp>" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
SRC=$1
|
||||
|
||||
if [[ ! -f "$SRC" ]]; then
|
||||
echo "grpc-server.cpp not found at $SRC" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
if grep -q '^#define LOCALAI_LLAMA_CPP_NO_SCORE_TASK' "$SRC"; then
|
||||
echo "==> $SRC already disables the LocalAI score task, skipping"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
awk '
|
||||
!done && /^#include/ {
|
||||
print "#define LOCALAI_LLAMA_CPP_NO_SCORE_TASK 1"
|
||||
print "// ^ injected by disable-score-task.sh for an unpatched llama.cpp fork"
|
||||
print ""
|
||||
done = 1
|
||||
}
|
||||
{ print }
|
||||
END {
|
||||
if (!done) {
|
||||
print "disable-score-task.sh: no #include anchor found" > "/dev/stderr"
|
||||
exit 1
|
||||
}
|
||||
}
|
||||
' "$SRC" > "$SRC.tmp"
|
||||
mv "$SRC.tmp" "$SRC"
|
||||
|
||||
echo "==> LocalAI score task disabled in $SRC"
|
||||
@@ -152,6 +152,40 @@ static std::string base64_encode_bytes(const unsigned char* data, size_t len) {
|
||||
|
||||
bool loaded_model; // TODO: add a mutex for this, but happens only once loading the model
|
||||
|
||||
// Score bypasses the slot loop (see the comment on Score below) so it
|
||||
// must not run concurrently with any slot-loop RPC. These counters
|
||||
// are a defence-in-depth tripwire — ModelConfig.Validate already
|
||||
// rejects llama-cpp configs that mix score with chat/completion/
|
||||
// embeddings, so a healthy deployment never trips them. seq_cst is
|
||||
// load-bearing for the increment-then-check pattern below.
|
||||
static std::atomic<int> slot_loop_inflight{0};
|
||||
static std::atomic<int> score_inflight{0};
|
||||
|
||||
// Increment-then-check, not check-then-increment: two simultaneous
|
||||
// racers both observe the other's increment and both abort cleanly.
|
||||
// Reversed, both could see zero and proceed.
|
||||
struct conflict_guard {
|
||||
std::atomic<int>& self;
|
||||
conflict_guard(const char* rpc, std::atomic<int>& self_, std::atomic<int>& other, const char* other_name)
|
||||
: self(self_) {
|
||||
self.fetch_add(1, std::memory_order_seq_cst);
|
||||
int o = other.load(std::memory_order_seq_cst);
|
||||
if (o > 0) {
|
||||
fprintf(stderr,
|
||||
"FATAL: %s called with %s=%d. The llama-cpp backend cannot "
|
||||
"service Score and slot-loop RPCs concurrently — Score "
|
||||
"bypasses the slot loop and races the llama_context. Bind "
|
||||
"Score-using features to a model dedicated to scoring "
|
||||
"(known_usecases: [score] with no chat/completion/embeddings).\n",
|
||||
rpc, other_name, o);
|
||||
std::abort();
|
||||
}
|
||||
}
|
||||
~conflict_guard() {
|
||||
self.fetch_sub(1, std::memory_order_seq_cst);
|
||||
}
|
||||
};
|
||||
|
||||
static std::function<void(int)> shutdown_handler;
|
||||
static std::atomic_flag is_terminating = ATOMIC_FLAG_INIT;
|
||||
|
||||
@@ -698,22 +732,6 @@ static void params_parse(server_context& /*ctx_server*/, const backend::ModelOpt
|
||||
// If conversion fails, keep default value (0)
|
||||
}
|
||||
}
|
||||
#ifndef LOCALAI_LLAMA_CPP_NO_SCORE_TASK
|
||||
} else if (!strcmp(optname, "n_rs_seq") || !strcmp(optname, "rs_seq")) {
|
||||
// Recurrent-state rollback snapshots per sequence. Hybrid models
|
||||
// (deltanet/conv layers) cannot rewind their state, so without
|
||||
// snapshots any prompt-cache reuse that needs a rewind — e.g. a
|
||||
// score task whose probe changed under a stable option-list
|
||||
// prefix — falls back to a full re-prefill. Costs recurrent-state
|
||||
// memory x (1 + N) per sequence; unsupported archs clamp to 0.
|
||||
if (optval != NULL) {
|
||||
try {
|
||||
params.n_rs_seq = std::stoi(optval_str);
|
||||
} catch (const std::exception& e) {
|
||||
// If conversion fails, keep default value (0)
|
||||
}
|
||||
}
|
||||
#endif
|
||||
} else if (!strcmp(optname, "slot_prompt_similarity") || !strcmp(optname, "sps")) {
|
||||
if (optval != NULL) {
|
||||
try {
|
||||
@@ -1374,17 +1392,6 @@ static void params_parse(server_context& /*ctx_server*/, const backend::ModelOpt
|
||||
}
|
||||
}
|
||||
|
||||
#ifndef LOCALAI_LLAMA_CPP_NO_SCORE_TASK
|
||||
// Score-task suffix forking: reserve seq ids (and recurrent-state cells)
|
||||
// beyond the slots so one scoring call decodes all candidate tails in a
|
||||
// single batch (SERVER_TASK_TYPE_SCORE, patches/). Requires the unified
|
||||
// KV cache — with per-sequence streams the extra ids would shrink every
|
||||
// sequence's context to n_ctx / n_seq_max. Decided after both option
|
||||
// passes so an explicit kv_unified:false wins and disables forking.
|
||||
params.score_enabled = request->enablescore();
|
||||
params.n_seq_score_forks = params.score_enabled && params.kv_unified ? SERVER_SCORE_FORK_SEQS : 0;
|
||||
#endif
|
||||
|
||||
// Terminate/pad the override vectors only after BOTH the named-option loop
|
||||
// and the generic passthrough (common_params_parse above) have pushed their
|
||||
// real entries, so back() is the null sentinel the model loader asserts on.
|
||||
@@ -1471,16 +1478,6 @@ public:
|
||||
common_params params;
|
||||
params_parse(ctx_server, request, params);
|
||||
|
||||
#ifndef LOCALAI_LLAMA_CPP_NO_SCORE_TASK
|
||||
if (params.score_enabled && !params.kv_unified) {
|
||||
const std::string error_msg =
|
||||
"Score requires the unified KV cache; remove kv_unified:false or remove score from known_usecases";
|
||||
result->set_message(error_msg);
|
||||
result->set_success(false);
|
||||
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, error_msg);
|
||||
}
|
||||
#endif
|
||||
|
||||
common_init();
|
||||
// Ensure debug logs are enabled after common_init() sets up logging
|
||||
common_log_set_verbosity_thold(params.verbosity);
|
||||
@@ -1683,6 +1680,7 @@ public:
|
||||
if (params_base.model.path.empty()) {
|
||||
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
|
||||
}
|
||||
conflict_guard guard("PredictStream", slot_loop_inflight, score_inflight, "score_inflight");
|
||||
json data = parse_options(true, request, params_base, ctx_server.get_llama_context());
|
||||
|
||||
|
||||
@@ -2251,6 +2249,7 @@ public:
|
||||
if (params_base.model.path.empty()) {
|
||||
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
|
||||
}
|
||||
conflict_guard guard("Predict", slot_loop_inflight, score_inflight, "score_inflight");
|
||||
json data = parse_options(true, request, params_base, ctx_server.get_llama_context());
|
||||
|
||||
data["stream"] = false;
|
||||
@@ -2784,6 +2783,7 @@ public:
|
||||
if (params_base.model.path.empty()) {
|
||||
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
|
||||
}
|
||||
conflict_guard guard("Embedding", slot_loop_inflight, score_inflight, "score_inflight");
|
||||
json body = parse_options(false, request, params_base, ctx_server.get_llama_context());
|
||||
|
||||
body["stream"] = false;
|
||||
@@ -2893,6 +2893,7 @@ public:
|
||||
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "\"documents\" must be a non-empty string array");
|
||||
}
|
||||
|
||||
conflict_guard guard("Rerank", slot_loop_inflight, score_inflight, "score_inflight");
|
||||
|
||||
// Create and queue the task
|
||||
auto rd = ctx_server.get_response_reader();
|
||||
@@ -2969,16 +2970,37 @@ public:
|
||||
// Score returns the model's joint log-probability of each candidate
|
||||
// continuation given a shared prompt.
|
||||
//
|
||||
// Scoring runs as a single SERVER_TASK_TYPE_SCORE task through the
|
||||
// slot loop (added by patches/ on top of upstream server-context), so
|
||||
// it is safe to interleave with generation on the same process and it
|
||||
// reuses any KV prefix the slot already holds across turns. The task
|
||||
// decodes the shared prefix (prompt + longest common candidate token
|
||||
// prefix) once on the slot's sequence; every candidate's unique tail
|
||||
// then rides its own forked sequence and all tails are decoded
|
||||
// together in one batch, so a warm scoring call costs roughly one
|
||||
// forward pass over the new prompt tokens plus one batched pass over
|
||||
// the candidate tails.
|
||||
// WHY bypass the slot/task queue: upstream server_context exposes
|
||||
// get_llama_context as "main thread only" and the slot loop's
|
||||
// update_slots() owns the context whenever a task is in flight.
|
||||
// No public synchronization primitive is available — so Score is
|
||||
// unsafe to call concurrently with active generation through this
|
||||
// backend. In practice routing-classifier calls happen before the
|
||||
// request is routed to a generation backend, so the model used
|
||||
// for Score is typically idle. Concurrent Score calls are
|
||||
// serialised by a local mutex; KV-cache state is isolated behind
|
||||
// a dedicated sequence ID cleared between candidates.
|
||||
//
|
||||
// A patch to server-context.cpp that adds SERVER_TASK_TYPE_SCORE
|
||||
// and routes scoring through the slot loop would be the correct
|
||||
// long-term fix; tracked as a follow-up.
|
||||
//
|
||||
// Perf TODO (measured: ~450 ms warm for 3 candidates on Arch-
|
||||
// Router-1.5B Q4_K_M + Intel SYCL): the current loop re-decodes
|
||||
// `prompt + candidate` from scratch for every candidate, throwing
|
||||
// away the prompt's KV cache between iterations. A smarter
|
||||
// version would:
|
||||
// 1. Decode just the prompt once into score_seq_id.
|
||||
// 2. Snapshot/cp that sequence (llama_memory_seq_cp) into a
|
||||
// per-candidate sequence id.
|
||||
// 3. For each candidate, decode only its tokens onto the copy
|
||||
// (continuing from the saved prompt state), read logits.
|
||||
// 4. llama_memory_seq_rm the copy.
|
||||
// Estimated speedup: 3-candidate calls 450 ms -> ~150-200 ms,
|
||||
// 6-candidate calls 630 ms -> ~220 ms. Single source-file change,
|
||||
// no proto / Go-side changes needed. Worth doing once routing is
|
||||
// wired into the middleware and Score is on the hot path of every
|
||||
// chat request.
|
||||
grpc::Status Score(ServerContext* context, const backend::ScoreRequest* request, backend::ScoreResponse* response) override {
|
||||
auto auth = checkAuth(context);
|
||||
if (!auth.ok()) return auth;
|
||||
@@ -2987,21 +3009,40 @@ public:
|
||||
if (params_base.model.path.empty()) {
|
||||
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
|
||||
}
|
||||
#ifdef LOCALAI_LLAMA_CPP_NO_SCORE_TASK
|
||||
(void) request;
|
||||
(void) response;
|
||||
return grpc::Status(grpc::StatusCode::UNIMPLEMENTED,
|
||||
"Score is unavailable in this llama.cpp fork backend");
|
||||
#else
|
||||
if (!params_base.score_enabled) {
|
||||
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION,
|
||||
"Score was not enabled when the model was loaded; add score to known_usecases");
|
||||
}
|
||||
if (request->candidates_size() == 0) {
|
||||
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "candidates must be non-empty");
|
||||
}
|
||||
|
||||
// Tripwire against the slot loop. Acquired before score_mutex
|
||||
// so it fires even when this Score is queued behind another.
|
||||
conflict_guard guard("Score", score_inflight, slot_loop_inflight, "slot_loop_inflight");
|
||||
|
||||
// Serialise concurrent Score calls. The slot loop is still
|
||||
// free to race with us — see the class comment above.
|
||||
static std::mutex score_mutex;
|
||||
std::lock_guard<std::mutex> score_lock(score_mutex);
|
||||
|
||||
llama_context * lctx = ctx_server.get_llama_context();
|
||||
if (lctx == nullptr) {
|
||||
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "llama context unavailable (sleeping?)");
|
||||
}
|
||||
const llama_vocab * vocab = ctx_server.impl->vocab;
|
||||
const int32_t n_vocab = llama_vocab_n_tokens(vocab);
|
||||
const int32_t n_ctx = llama_n_ctx(lctx);
|
||||
llama_memory_t mem = llama_get_memory(lctx);
|
||||
|
||||
// The KV-cache is sized to seq_to_stream.size() at load
|
||||
// (typically equal to n_slots, often 1). Sequence IDs must
|
||||
// be in [0, n_seq_max), so we can't pick a high-value
|
||||
// "private" ID — we have to share with the slot. We clear
|
||||
// the cache before AND after each candidate to keep
|
||||
// scoring isolated from whatever state the slot held, and
|
||||
// the static mutex above guarantees no other Score call is
|
||||
// racing in the meantime. The slot loop is still free to
|
||||
// race (see comment on this method) — Score must not run
|
||||
// concurrently with generation through this backend.
|
||||
const llama_seq_id score_seq_id = 0;
|
||||
llama_memory_seq_rm(mem, score_seq_id, -1, -1);
|
||||
|
||||
// Tokenize the shared prompt once with add_special=true so
|
||||
// BOS is prepended when the model requires it. parse_special
|
||||
@@ -3010,15 +3051,6 @@ public:
|
||||
std::vector<llama_token> prompt_tokens = common_tokenize(vocab, prompt, /*add_special=*/true, /*parse_special=*/true);
|
||||
const int32_t prompt_len = (int32_t) prompt_tokens.size();
|
||||
|
||||
// Per candidate: full prompt+candidate token list and the
|
||||
// divergence point, kept for piece rendering and empty-candidate
|
||||
// handling after the task comes back.
|
||||
std::vector<std::vector<llama_token>> cand_tokens(request->candidates_size());
|
||||
std::vector<int32_t> cand_divergence(request->candidates_size(), 0);
|
||||
|
||||
// candidates that actually have tokens to score
|
||||
std::vector<int32_t> included;
|
||||
|
||||
for (int ci = 0; ci < request->candidates_size(); ci++) {
|
||||
const std::string & candidate_text = request->candidates(ci);
|
||||
|
||||
@@ -3035,135 +3067,9 @@ public:
|
||||
break;
|
||||
}
|
||||
}
|
||||
divergence = std::min<int32_t>(divergence, (int32_t) full_tokens.size());
|
||||
|
||||
const int32_t cand_len = (int32_t) full_tokens.size() - divergence;
|
||||
if (cand_len > 0 && divergence < 1) {
|
||||
// Need at least one prior token (typically BOS) to
|
||||
// predict the first candidate token's logit. Tokeniser
|
||||
// models without BOS + an empty prompt fall in here.
|
||||
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
|
||||
"Score: prompt produced no leading tokens; need at least one (e.g. BOS) to predict candidate");
|
||||
}
|
||||
if (cand_len > SERVER_SCORE_MAX_CAND_TOKENS) {
|
||||
// The context reserves logits outputs for at most this many
|
||||
// candidate tokens per slot (server_n_outputs_max).
|
||||
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
|
||||
"Score: candidate " + std::to_string(ci) + " is " + std::to_string(cand_len) +
|
||||
" tokens; the maximum is " + std::to_string(SERVER_SCORE_MAX_CAND_TOKENS));
|
||||
}
|
||||
|
||||
cand_divergence[ci] = divergence;
|
||||
cand_tokens[ci] = std::move(full_tokens);
|
||||
|
||||
if (cand_len > 0) {
|
||||
included.push_back(ci);
|
||||
}
|
||||
}
|
||||
|
||||
auto rd = ctx_server.get_response_reader();
|
||||
bool posted_task = false;
|
||||
|
||||
// Shared prefix bounds, needed again when stitching the results:
|
||||
// n_shared is the longest common token prefix of the scored
|
||||
// candidates, n_score_prompt the earliest divergence from the
|
||||
// bare prompt (scored logprobs start there).
|
||||
int32_t n_shared = 0;
|
||||
int32_t n_score_prompt = 0;
|
||||
|
||||
if (!included.empty()) {
|
||||
const auto & first = cand_tokens[included[0]];
|
||||
|
||||
// the common prefix of a set is the shortest common prefix
|
||||
// against any fixed member
|
||||
n_shared = (int32_t) first.size();
|
||||
for (int32_t ci : included) {
|
||||
const auto & ft = cand_tokens[ci];
|
||||
const int32_t lim = std::min<int32_t>(n_shared, (int32_t) ft.size());
|
||||
int32_t match = 0;
|
||||
while (match < lim && ft[match] == first[match]) {
|
||||
match++;
|
||||
}
|
||||
n_shared = match;
|
||||
}
|
||||
|
||||
// below its divergence every candidate equals the prompt
|
||||
// tokens, so n_score_prompt <= n_shared always holds
|
||||
n_score_prompt = cand_divergence[included[0]];
|
||||
for (int32_t ci : included) {
|
||||
n_score_prompt = std::min(n_score_prompt, cand_divergence[ci]);
|
||||
}
|
||||
|
||||
// Map the caller's stable-prefix byte length onto a token
|
||||
// index: the last prompt token that ends at or before the
|
||||
// boundary. A checkpoint forced there survives every future
|
||||
// probe under the same option list, which is what keeps
|
||||
// repeat scoring cheap on models that cannot rewind state.
|
||||
int32_t n_stable_prompt = 0;
|
||||
if (request->stable_prefix_len() > 0) {
|
||||
size_t consumed = 0;
|
||||
for (int32_t ti = 0; ti < n_score_prompt; ti++) {
|
||||
const size_t piece_len = common_token_to_piece(vocab, prompt_tokens[ti]).size();
|
||||
// BOS and other zero-length specials consume no prompt bytes
|
||||
if (consumed + piece_len > (size_t) request->stable_prefix_len()) {
|
||||
break;
|
||||
}
|
||||
consumed += piece_len;
|
||||
n_stable_prompt = ti + 1;
|
||||
}
|
||||
}
|
||||
|
||||
server_task task(SERVER_TASK_TYPE_SCORE);
|
||||
task.id = rd.queue_tasks.get_new_id();
|
||||
task.index = 0;
|
||||
task.tokens = server_tokens(llama_tokens(first.begin(), first.begin() + n_shared), false);
|
||||
task.n_score_prompt = n_score_prompt;
|
||||
task.n_stable_prompt = n_stable_prompt;
|
||||
task.score_suffixes.reserve(included.size());
|
||||
for (int32_t ci : included) {
|
||||
task.score_suffixes.emplace_back(cand_tokens[ci].begin() + n_shared, cand_tokens[ci].end());
|
||||
}
|
||||
|
||||
std::vector<server_task> tasks;
|
||||
tasks.push_back(std::move(task));
|
||||
rd.post_tasks(std::move(tasks));
|
||||
posted_task = true;
|
||||
}
|
||||
|
||||
// Wait for the shared-prefix and per-candidate logprob vectors.
|
||||
// Context overflow and decode failures surface here as task errors.
|
||||
std::vector<float> shared_logprobs;
|
||||
std::vector<std::vector<float>> cand_logprobs;
|
||||
if (posted_task) {
|
||||
auto all_results = rd.wait_for_all([&context]() { return context->IsCancelled(); });
|
||||
if (all_results.is_terminated) {
|
||||
return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled by client");
|
||||
}
|
||||
if (all_results.error) {
|
||||
return grpc::Status(grpc::StatusCode::INTERNAL,
|
||||
all_results.error->to_json().value("message", "Error in receiving score results"));
|
||||
}
|
||||
if (all_results.results.size() != 1) {
|
||||
return grpc::Status(grpc::StatusCode::INTERNAL, "expected a single score result");
|
||||
}
|
||||
auto * score_res = dynamic_cast<server_task_result_score*>(all_results.results[0].get());
|
||||
if (score_res == nullptr) {
|
||||
return grpc::Status(grpc::StatusCode::INTERNAL, "unexpected result type for score task");
|
||||
}
|
||||
shared_logprobs = std::move(score_res->shared_logprobs);
|
||||
cand_logprobs = std::move(score_res->cand_logprobs);
|
||||
if (cand_logprobs.size() != included.size()) {
|
||||
return grpc::Status(grpc::StatusCode::INTERNAL, "score result candidate count mismatch");
|
||||
}
|
||||
}
|
||||
|
||||
size_t inc = 0; // index into included / cand_logprobs
|
||||
for (int ci = 0; ci < request->candidates_size(); ci++) {
|
||||
const int32_t divergence = cand_divergence[ci];
|
||||
const int32_t cand_len = (int32_t) cand_tokens[ci].size() - divergence;
|
||||
|
||||
backend::CandidateScore * cs = response->add_candidates();
|
||||
cs->set_num_tokens(cand_len > 0 ? cand_len : 0);
|
||||
cs->set_num_tokens(cand_len);
|
||||
if (cand_len <= 0) {
|
||||
cs->set_log_prob(0.0);
|
||||
if (request->length_normalize()) {
|
||||
@@ -3171,57 +3077,101 @@ public:
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// Stitch the candidate's scored logprobs back together: the
|
||||
// stretch inside the shared prefix (identical for every
|
||||
// candidate) followed by its forked suffix. Suffix entries
|
||||
// before the candidate's own divergence are prompt tokens
|
||||
// decoded only as context — not scored.
|
||||
std::vector<float> lp;
|
||||
lp.reserve(cand_len);
|
||||
for (int32_t t = divergence; t < n_shared; t++) {
|
||||
const int32_t idx = t - n_score_prompt;
|
||||
if (idx < 0 || idx >= (int32_t) shared_logprobs.size()) {
|
||||
return grpc::Status(grpc::StatusCode::INTERNAL,
|
||||
"Score: shared logprob index out of range for candidate " + std::to_string(ci));
|
||||
}
|
||||
lp.push_back(shared_logprobs[idx]);
|
||||
if (divergence < 1) {
|
||||
// Need at least one prior token (typically BOS) to
|
||||
// predict the first candidate token's logit. Tokeniser
|
||||
// models without BOS + an empty prompt fall in here.
|
||||
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
|
||||
"Score: prompt produced no leading tokens; need at least one (e.g. BOS) to predict candidate");
|
||||
}
|
||||
const auto & sfx_lp = cand_logprobs[inc++];
|
||||
for (int32_t j = std::max(0, divergence - n_shared); j < (int32_t) sfx_lp.size(); j++) {
|
||||
lp.push_back(sfx_lp[j]);
|
||||
if ((int32_t) full_tokens.size() > n_ctx) {
|
||||
return grpc::Status(grpc::StatusCode::OUT_OF_RANGE,
|
||||
"Score: prompt+candidate exceeds context size (got " +
|
||||
std::to_string(full_tokens.size()) + ", n_ctx=" + std::to_string(n_ctx) + ")");
|
||||
}
|
||||
|
||||
if ((int32_t) lp.size() != cand_len) {
|
||||
// Build a batch covering the entire prompt+candidate. We
|
||||
// need logits at (divergence-1) onward — those are the
|
||||
// predictions for each candidate token.
|
||||
llama_batch batch = llama_batch_init((int32_t) full_tokens.size(), 0, 1);
|
||||
for (int32_t i = 0; i < (int32_t) full_tokens.size(); i++) {
|
||||
batch.token[i] = full_tokens[i];
|
||||
batch.pos[i] = i;
|
||||
batch.n_seq_id[i] = 1;
|
||||
batch.seq_id[i][0] = score_seq_id;
|
||||
// logits[i] is "do we want the prediction *for the
|
||||
// next token*, computed from this position?"
|
||||
// We want predictions for candidate tokens at
|
||||
// positions divergence .. full_tokens.size()-1, which
|
||||
// come from logits at positions (divergence-1) ..
|
||||
// (full_tokens.size()-2).
|
||||
bool need_logit = (i >= divergence - 1) && (i < (int32_t) full_tokens.size() - 1);
|
||||
batch.logits[i] = need_logit ? 1 : 0;
|
||||
}
|
||||
batch.n_tokens = (int32_t) full_tokens.size();
|
||||
|
||||
// Decode the batch. If decode fails (e.g. KV slot
|
||||
// exhaustion), surface as INTERNAL — the caller will
|
||||
// typically fall back to a sampling-based classifier.
|
||||
int decode_err = llama_decode(lctx, batch);
|
||||
if (decode_err != 0) {
|
||||
llama_batch_free(batch);
|
||||
llama_memory_seq_rm(mem, score_seq_id, -1, -1);
|
||||
return grpc::Status(grpc::StatusCode::INTERNAL,
|
||||
"Score: result for candidate " + std::to_string(ci) + " is missing token logprobs");
|
||||
"llama_decode failed during Score: " + std::to_string(decode_err));
|
||||
}
|
||||
|
||||
// Sum log-probabilities of the actual candidate tokens.
|
||||
double total_log_prob = 0.0;
|
||||
for (int32_t k = 0; k < cand_len; k++) {
|
||||
const float token_log_prob = lp[k];
|
||||
if (std::isnan(token_log_prob)) {
|
||||
// The k-th candidate token sits at full_tokens index
|
||||
// (divergence + k). Its predicting logit is at batch
|
||||
// position (divergence + k - 1).
|
||||
int32_t logit_pos = divergence + k - 1;
|
||||
const float * logits = llama_get_logits_ith(lctx, logit_pos);
|
||||
if (logits == nullptr) {
|
||||
llama_batch_free(batch);
|
||||
llama_memory_seq_rm(mem, score_seq_id, -1, -1);
|
||||
return grpc::Status(grpc::StatusCode::INTERNAL,
|
||||
"Score: incomplete result for candidate " + std::to_string(ci) +
|
||||
" at token " + std::to_string(k));
|
||||
"llama_get_logits_ith returned null at position " + std::to_string(logit_pos));
|
||||
}
|
||||
total_log_prob += (double) token_log_prob;
|
||||
llama_token target_token = full_tokens[divergence + k];
|
||||
|
||||
// Compute log_softmax(logits)[target_token] with the
|
||||
// max-subtraction stability trick.
|
||||
float max_logit = logits[0];
|
||||
for (int32_t v = 1; v < n_vocab; v++) {
|
||||
if (logits[v] > max_logit) max_logit = logits[v];
|
||||
}
|
||||
double sum_exp = 0.0;
|
||||
for (int32_t v = 0; v < n_vocab; v++) {
|
||||
sum_exp += std::exp((double)(logits[v] - max_logit));
|
||||
}
|
||||
double token_log_prob = (double)(logits[target_token] - max_logit) - std::log(sum_exp);
|
||||
total_log_prob += token_log_prob;
|
||||
|
||||
if (request->include_token_logprobs()) {
|
||||
backend::TokenLogProb * tlp = cs->add_tokens();
|
||||
tlp->set_token(common_token_to_piece(vocab, cand_tokens[ci][divergence + k]));
|
||||
std::string piece = common_token_to_piece(lctx, target_token);
|
||||
tlp->set_token(piece);
|
||||
tlp->set_log_prob(token_log_prob);
|
||||
}
|
||||
}
|
||||
|
||||
cs->set_log_prob(total_log_prob);
|
||||
if (request->length_normalize()) {
|
||||
if (request->length_normalize() && cand_len > 0) {
|
||||
cs->set_length_normalized_log_prob(total_log_prob / (double) cand_len);
|
||||
}
|
||||
|
||||
llama_batch_free(batch);
|
||||
// Drop this candidate's KV-cache contribution so the next
|
||||
// candidate starts from a clean state. Without this, the
|
||||
// next decode would conflict at positions 0..N-1 for our
|
||||
// sequence ID.
|
||||
llama_memory_seq_rm(mem, score_seq_id, -1, -1);
|
||||
}
|
||||
|
||||
return grpc::Status::OK;
|
||||
#endif
|
||||
}
|
||||
|
||||
grpc::Status TokenizeString(ServerContext* context, const backend::PredictOptions* request, backend::TokenizationResponse* response) override {
|
||||
@@ -3232,6 +3182,7 @@ public:
|
||||
if (params_base.model.path.empty()) {
|
||||
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
|
||||
}
|
||||
conflict_guard guard("TokenizeString", slot_loop_inflight, score_inflight, "score_inflight");
|
||||
json body = parse_options(false, request, params_base, ctx_server.get_llama_context());
|
||||
body["stream"] = false;
|
||||
|
||||
@@ -3253,6 +3204,7 @@ public:
|
||||
|
||||
grpc::Status GetMetrics(ServerContext* /*context*/, const backend::MetricsRequest* /*request*/, backend::MetricsResponse* response) override {
|
||||
|
||||
conflict_guard guard("GetMetrics", slot_loop_inflight, score_inflight, "score_inflight");
|
||||
|
||||
// request slots data using task queue
|
||||
auto rd = ctx_server.get_response_reader();
|
||||
|
||||
@@ -1,599 +0,0 @@
|
||||
diff --git a/common/common.cpp b/common/common.cpp
|
||||
index 8f13217..fc584e1 100644
|
||||
--- a/common/common.cpp
|
||||
+++ b/common/common.cpp
|
||||
@@ -1591,8 +1591,10 @@ struct llama_context_params common_context_params_to_llama(const common_params &
|
||||
auto cparams = llama_context_default_params();
|
||||
|
||||
cparams.n_ctx = params.n_ctx;
|
||||
- cparams.n_seq_max = params.n_parallel;
|
||||
- cparams.n_rs_seq = params.speculative.need_n_rs_seq();
|
||||
+ // score-task forks need seq ids (and recurrent-state cells) of their
|
||||
+ // own beyond the parallel slots
|
||||
+ cparams.n_seq_max = params.n_parallel + params.n_seq_score_forks;
|
||||
+ cparams.n_rs_seq = std::max(params.speculative.need_n_rs_seq(), (uint32_t) std::max(0, params.n_rs_seq));
|
||||
cparams.n_outputs_max = std::max(params.n_outputs_max, 0);
|
||||
cparams.n_batch = params.n_batch;
|
||||
cparams.n_ubatch = params.n_ubatch;
|
||||
diff --git a/common/common.h b/common/common.h
|
||||
index bffc176..e313bd6 100644
|
||||
--- a/common/common.h
|
||||
+++ b/common/common.h
|
||||
@@ -455,6 +455,9 @@ struct common_params {
|
||||
int32_t n_keep = 0; // number of tokens to keep from initial prompt
|
||||
int32_t n_chunks = -1; // max number of chunks to process (-1 = unlimited)
|
||||
int32_t n_parallel = 1; // number of parallel sequences to decode
|
||||
+ int32_t n_seq_score_forks = 0; // extra seq ids beyond n_parallel, reserved for server score-task forks
|
||||
+ int32_t n_rs_seq = 0; // recurrent-state rollback snapshots per seq (hybrid models cannot rewind without them; lets score tasks reuse a cached prompt across probe changes)
|
||||
+ bool score_enabled = false; // reserve server resources for the Score task type
|
||||
int32_t n_sequences = 1; // number of sequences to decode
|
||||
int32_t n_outputs_max = 0; // max outputs in a batch (0 = n_batch)
|
||||
int32_t grp_attn_n = 1; // group-attention factor
|
||||
diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt
|
||||
index 780df32..1d2fe8f 100644
|
||||
--- a/tools/CMakeLists.txt
|
||||
+++ b/tools/CMakeLists.txt
|
||||
@@ -41,3 +41,4 @@ else()
|
||||
add_subdirectory(fit-params)
|
||||
add_subdirectory(results)
|
||||
endif()
|
||||
+add_subdirectory(grpc-server)
|
||||
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
|
||||
index 715477e..de5bed8 100644
|
||||
--- a/tools/server/server-context.cpp
|
||||
+++ b/tools/server/server-context.cpp
|
||||
@@ -49,7 +49,16 @@ static uint32_t server_n_outputs_max(const common_params & params) {
|
||||
|
||||
const uint32_t n_outputs_per_seq = 1 + common_speculative_n_max(¶ms.speculative);
|
||||
|
||||
- const uint64_t n_outputs = (uint64_t) params.n_parallel * n_outputs_per_seq;
|
||||
+ // score tasks (SERVER_TASK_TYPE_SCORE) output logits for every candidate
|
||||
+ // token, so reserve room for a bounded candidate tail per parallel slot
|
||||
+ if (!params.score_enabled) {
|
||||
+ return std::max<uint32_t>(1, std::min<uint64_t>(n_batch,
|
||||
+ (uint64_t) params.n_parallel * n_outputs_per_seq));
|
||||
+ }
|
||||
+
|
||||
+ const uint32_t n_outputs_score_seq = 1 + SERVER_SCORE_MAX_CAND_TOKENS;
|
||||
+
|
||||
+ const uint64_t n_outputs = (uint64_t) params.n_parallel * std::max(n_outputs_per_seq, n_outputs_score_seq);
|
||||
|
||||
return std::max<uint32_t>(1, std::min<uint64_t>(n_batch, n_outputs));
|
||||
}
|
||||
@@ -202,6 +211,26 @@ struct server_slot {
|
||||
|
||||
std::vector<completion_token_output> generated_token_probs;
|
||||
|
||||
+ // SERVER_TASK_TYPE_SCORE: shared-prefix token logprobs harvested
|
||||
+ // incrementally across batch views (NaN = not yet produced)
|
||||
+ std::vector<float> score_logprobs;
|
||||
+
|
||||
+ // SERVER_TASK_TYPE_SCORE: per-candidate suffix token logprobs; entry
|
||||
+ // [c][0] comes from the last shared token's logits during prompt
|
||||
+ // processing, the rest from the forked suffix decode
|
||||
+ std::vector<std::vector<float>> score_cand_logprobs;
|
||||
+
|
||||
+ // SERVER_TASK_TYPE_SCORE: the prompt completed but some candidate has
|
||||
+ // suffix tokens beyond the first, so a forked decode is still needed
|
||||
+ bool score_suffix_pending = false;
|
||||
+
|
||||
+ // SERVER_TASK_TYPE_SCORE: where the current task's tokens diverged from
|
||||
+ // the slot's previous cache. When the memory cannot rewind there and a
|
||||
+ // re-prefill follows, a checkpoint at this position lets the next
|
||||
+ // scoring call over the same stable prefix (e.g. a classifier's option
|
||||
+ // list) resume from it instead of re-processing the whole prompt.
|
||||
+ int32_t score_divergence = -1;
|
||||
+
|
||||
bool has_next_token = true;
|
||||
bool has_new_line = false;
|
||||
bool truncated = false;
|
||||
@@ -311,6 +340,10 @@ struct server_slot {
|
||||
}
|
||||
generated_tokens.clear();
|
||||
generated_token_probs.clear();
|
||||
+ score_logprobs.clear();
|
||||
+ score_cand_logprobs.clear();
|
||||
+ score_suffix_pending = false;
|
||||
+ score_divergence = -1;
|
||||
json_schema = json();
|
||||
|
||||
// clear speculative decoding stats
|
||||
@@ -2205,6 +2238,229 @@ private:
|
||||
queue_results.send(std::move(res));
|
||||
}
|
||||
|
||||
+ // log(sum(exp(logits))) with max-subtraction for stability — the
|
||||
+ // log_softmax denominator shared by every token read from one output
|
||||
+ static double score_log_denom(const float * logits, int32_t n_vocab) {
|
||||
+ float max_logit = logits[0];
|
||||
+ for (int32_t v = 1; v < n_vocab; ++v) {
|
||||
+ max_logit = std::max(max_logit, logits[v]);
|
||||
+ }
|
||||
+ double sum_exp = 0.0;
|
||||
+ for (int32_t v = 0; v < n_vocab; ++v) {
|
||||
+ sum_exp += std::exp((double)(logits[v] - max_logit));
|
||||
+ }
|
||||
+ return (double) max_logit + std::log(sum_exp);
|
||||
+ }
|
||||
+
|
||||
+ // Harvest logprobs for SCORE tasks from the current batch view: the
|
||||
+ // shared-prefix scored tokens, and — from the last shared token's
|
||||
+ // logits — the first suffix token of every candidate. The scored
|
||||
+ // region can straddle ubatch boundaries for long prompts, so this
|
||||
+ // accumulates view by view instead of reading everything when the
|
||||
+ // prompt completes.
|
||||
+ void collect_score_logprobs(server_slot & slot, const llama_batch & batch) {
|
||||
+ const int32_t n_prompt = slot.task->n_score_prompt;
|
||||
+ const int32_t n_total = slot.task->n_tokens();
|
||||
+ const auto & suffixes = slot.task->score_suffixes;
|
||||
+
|
||||
+ const size_t n_shared_scored = (size_t) std::max(0, n_total - n_prompt);
|
||||
+
|
||||
+ if (slot.score_logprobs.size() != n_shared_scored) {
|
||||
+ slot.score_logprobs.assign(n_shared_scored, NAN);
|
||||
+ }
|
||||
+ if (slot.score_cand_logprobs.size() != suffixes.size()) {
|
||||
+ slot.score_cand_logprobs.resize(suffixes.size());
|
||||
+ for (size_t c = 0; c < suffixes.size(); ++c) {
|
||||
+ slot.score_cand_logprobs[c].assign(suffixes[c].size(), NAN);
|
||||
+ }
|
||||
+ }
|
||||
+
|
||||
+ const int32_t n_vocab = llama_vocab_n_tokens(vocab);
|
||||
+
|
||||
+ for (int32_t i = 0; i < batch.n_tokens; ++i) {
|
||||
+ if (!batch.logits[i] || batch.seq_id[i][0] != slot.id) {
|
||||
+ continue;
|
||||
+ }
|
||||
+
|
||||
+ // the output at position p predicts the task token at index p + 1;
|
||||
+ // score tasks are text-only, so positions equal token indices
|
||||
+ const int32_t target = batch.pos[i] + 1;
|
||||
+ if (target < n_prompt || target > n_total) {
|
||||
+ continue;
|
||||
+ }
|
||||
+
|
||||
+ const float * logits = llama_get_logits_ith(slot.ctx_tgt, i);
|
||||
+ if (logits == nullptr) {
|
||||
+ SLT_ERR(slot, "failed to get logits for score target %d\n", target);
|
||||
+ continue;
|
||||
+ }
|
||||
+
|
||||
+ const double log_denom = score_log_denom(logits, n_vocab);
|
||||
+
|
||||
+ if (target < n_total) {
|
||||
+ const llama_token tok = slot.task->tokens[target];
|
||||
+ slot.score_logprobs[target - n_prompt] = (float) ((double) logits[tok] - log_denom);
|
||||
+ } else {
|
||||
+ // the last shared token predicts the first suffix token of
|
||||
+ // every candidate
|
||||
+ for (size_t c = 0; c < suffixes.size(); ++c) {
|
||||
+ if (!suffixes[c].empty()) {
|
||||
+ slot.score_cand_logprobs[c][0] = (float) ((double) logits[suffixes[c][0]] - log_denom);
|
||||
+ }
|
||||
+ }
|
||||
+ }
|
||||
+ }
|
||||
+ }
|
||||
+
|
||||
+ void send_score(server_slot & slot) {
|
||||
+ auto res = std::make_unique<server_task_result_score>();
|
||||
+ res->id = slot.task->id;
|
||||
+ res->index = slot.task->index;
|
||||
+ res->shared_logprobs = std::move(slot.score_logprobs);
|
||||
+ res->cand_logprobs = std::move(slot.score_cand_logprobs);
|
||||
+
|
||||
+ slot.score_logprobs.clear();
|
||||
+ slot.score_cand_logprobs.clear();
|
||||
+
|
||||
+ SLT_DBG(slot, "sending score result, n_shared = %zu, n_cand = %zu\n",
|
||||
+ res->shared_logprobs.size(), res->cand_logprobs.size());
|
||||
+
|
||||
+ queue_results.send(std::move(res));
|
||||
+ }
|
||||
+
|
||||
+ // Decode the candidate suffixes of a completed score prompt: fork one
|
||||
+ // sequence per candidate off the slot's shared prefix (metadata-only
|
||||
+ // for the unified KV cache, copy-on-write for recurrent state) and
|
||||
+ // decode all unique suffix tokens in as few llama_decode calls as the
|
||||
+ // fork/batch/output budgets allow, harvesting a logprob for every
|
||||
+ // suffix token that predicts a following one.
|
||||
+ bool decode_score_suffixes(server_slot & slot) {
|
||||
+ const auto & suffixes = slot.task->score_suffixes;
|
||||
+
|
||||
+ auto * mem = llama_get_memory(ctx_tgt);
|
||||
+
|
||||
+ // seq ids beyond the slots are reserved for score forks at context
|
||||
+ // creation (common_params::n_seq_score_forks)
|
||||
+ const int32_t seq_base = (int32_t) slots.size();
|
||||
+ const int32_t n_forks_max = std::min<int32_t>(SERVER_SCORE_FORK_SEQS, (int32_t) llama_n_seq_max(ctx_tgt) - seq_base);
|
||||
+
|
||||
+ if (n_forks_max < 1) {
|
||||
+ SLT_ERR(slot, "no fork sequences reserved for score suffixes (n_seq_max = %d, n_slots = %d)\n",
|
||||
+ (int32_t) llama_n_seq_max(ctx_tgt), seq_base);
|
||||
+ return false;
|
||||
+ }
|
||||
+
|
||||
+ const int32_t n_batch_max = llama_n_batch(ctx_tgt);
|
||||
+ const int32_t n_vocab = llama_vocab_n_tokens(vocab);
|
||||
+ const llama_pos pos0 = slot.prompt.tokens.pos_next();
|
||||
+
|
||||
+ std::vector<size_t> pending;
|
||||
+ for (size_t c = 0; c < suffixes.size(); ++c) {
|
||||
+ // single-token suffixes were fully scored from the last shared
|
||||
+ // token's logits during prompt processing
|
||||
+ if (suffixes[c].size() > 1) {
|
||||
+ if ((int32_t) suffixes[c].size() > n_batch_max) {
|
||||
+ SLT_ERR(slot, "score suffix of candidate %zu (%zu tokens) exceeds n_batch (%d)\n",
|
||||
+ c, suffixes[c].size(), n_batch_max);
|
||||
+ return false;
|
||||
+ }
|
||||
+ pending.push_back(c);
|
||||
+ }
|
||||
+ }
|
||||
+
|
||||
+ size_t next = 0;
|
||||
+ while (next < pending.size()) {
|
||||
+ // greedy-pack candidates into one decode within the fork,
|
||||
+ // batch and reserved-output budgets
|
||||
+ std::vector<size_t> chunk;
|
||||
+ int32_t n_tok = 0;
|
||||
+ int32_t n_out = 0;
|
||||
+ while (next < pending.size() && (int32_t) chunk.size() < n_forks_max) {
|
||||
+ const int32_t m = (int32_t) suffixes[pending[next]].size();
|
||||
+ if (!chunk.empty() && (n_tok + m > n_batch_max || n_out + m - 1 > SERVER_SCORE_MAX_CAND_TOKENS)) {
|
||||
+ break;
|
||||
+ }
|
||||
+ chunk.push_back(pending[next]);
|
||||
+ n_tok += m;
|
||||
+ n_out += m - 1;
|
||||
+ next++;
|
||||
+ }
|
||||
+
|
||||
+ llama_batch fb = llama_batch_init(n_tok, 0, 1);
|
||||
+
|
||||
+ for (size_t k = 0; k < chunk.size(); ++k) {
|
||||
+ const llama_seq_id seq = seq_base + (llama_seq_id) k;
|
||||
+ const auto & sfx = suffixes[chunk[k]];
|
||||
+
|
||||
+ llama_memory_seq_rm(mem, seq, -1, -1);
|
||||
+ llama_memory_seq_cp(mem, slot.id, seq, -1, -1);
|
||||
+
|
||||
+ for (size_t j = 0; j < sfx.size(); ++j) {
|
||||
+ common_batch_add(fb, sfx[j], pos0 + (llama_pos) j, { seq }, j + 1 < sfx.size());
|
||||
+ }
|
||||
+ }
|
||||
+
|
||||
+ const int ret = llama_decode(ctx_tgt, fb);
|
||||
+
|
||||
+ if (ret == 0) {
|
||||
+ int32_t i = 0;
|
||||
+ for (size_t k = 0; k < chunk.size(); ++k) {
|
||||
+ const auto & sfx = suffixes[chunk[k]];
|
||||
+ auto & out = slot.score_cand_logprobs[chunk[k]];
|
||||
+
|
||||
+ for (size_t j = 0; j < sfx.size(); ++j, ++i) {
|
||||
+ if (j + 1 >= sfx.size()) {
|
||||
+ continue; // last suffix token predicts nothing
|
||||
+ }
|
||||
+ const float * logits = llama_get_logits_ith(ctx_tgt, i);
|
||||
+ if (logits == nullptr) {
|
||||
+ SLT_ERR(slot, "failed to get logits for suffix token %zu of score candidate %zu\n", j, chunk[k]);
|
||||
+ continue;
|
||||
+ }
|
||||
+ const double log_denom = score_log_denom(logits, n_vocab);
|
||||
+ out[j + 1] = (float) ((double) logits[sfx[j + 1]] - log_denom);
|
||||
+ }
|
||||
+ }
|
||||
+ }
|
||||
+
|
||||
+ for (size_t k = 0; k < chunk.size(); ++k) {
|
||||
+ llama_memory_seq_rm(mem, seq_base + (llama_seq_id) k, -1, -1);
|
||||
+ }
|
||||
+
|
||||
+ llama_batch_free(fb);
|
||||
+
|
||||
+ if (ret != 0) {
|
||||
+ SLT_ERR(slot, "score suffix decode failed, ret = %d\n", ret);
|
||||
+ return false;
|
||||
+ }
|
||||
+ }
|
||||
+
|
||||
+ return true;
|
||||
+ }
|
||||
+
|
||||
+ // score slots whose prompt completed this iteration decode their
|
||||
+ // candidate suffixes here, after every batch view was consumed — a
|
||||
+ // mid-view llama_decode would clobber logits other slots still read
|
||||
+ void update_score_suffixes() {
|
||||
+ for (auto & slot : slots) {
|
||||
+ if (!slot.score_suffix_pending) {
|
||||
+ continue;
|
||||
+ }
|
||||
+ slot.score_suffix_pending = false;
|
||||
+
|
||||
+ if (!slot.is_processing() || !slot.task || slot.task->type != SERVER_TASK_TYPE_SCORE) {
|
||||
+ continue; // the task was aborted mid-iteration
|
||||
+ }
|
||||
+
|
||||
+ if (decode_score_suffixes(slot)) {
|
||||
+ send_score(slot);
|
||||
+ } else {
|
||||
+ send_error(slot, "failed to decode score candidate suffixes", ERROR_TYPE_SERVER);
|
||||
+ }
|
||||
+ slot.release();
|
||||
+ }
|
||||
+ }
|
||||
+
|
||||
//
|
||||
// Functions to process the task
|
||||
//
|
||||
@@ -2341,6 +2597,7 @@ private:
|
||||
case SERVER_TASK_TYPE_INFILL:
|
||||
case SERVER_TASK_TYPE_EMBEDDING:
|
||||
case SERVER_TASK_TYPE_RERANK:
|
||||
+ case SERVER_TASK_TYPE_SCORE:
|
||||
{
|
||||
// special case: if input is provided via CLI, tokenize it first
|
||||
// otherwise, no need to tokenize as it's already done inside the HTTP thread
|
||||
@@ -2832,6 +3089,13 @@ private:
|
||||
break; // stop any further processing
|
||||
}
|
||||
}
|
||||
+
|
||||
+ try {
|
||||
+ update_score_suffixes();
|
||||
+ } catch (const std::exception & e) {
|
||||
+ SRV_ERR("update_score_suffixes() failed: %s\n", e.what());
|
||||
+ abort_all_slots("update_score_suffixes() failed: " + std::string(e.what()));
|
||||
+ }
|
||||
}
|
||||
|
||||
void pre_decode() {
|
||||
@@ -3154,6 +3418,16 @@ private:
|
||||
n_past = std::min(n_past, slot.alora_invocation_start - 1);
|
||||
}
|
||||
|
||||
+ // score tasks need the logits that predict the first candidate
|
||||
+ // token, so the last shared-prompt token must be (re-)decoded
|
||||
+ // even when the cache already covers it
|
||||
+ if (slot.task->type == SERVER_TASK_TYPE_SCORE) {
|
||||
+ n_past = std::min(n_past, std::max(0, slot.task->n_score_prompt - 1));
|
||||
+ // remember the divergence point before the checkpoint
|
||||
+ // logic below possibly resets n_past to 0
|
||||
+ slot.score_divergence = n_past;
|
||||
+ }
|
||||
+
|
||||
const auto n_cache_reuse = slot.task->params.n_cache_reuse;
|
||||
|
||||
const bool can_cache_reuse =
|
||||
@@ -3395,8 +3669,12 @@ private:
|
||||
|
||||
bool do_checkpoint = params_base.n_ctx_checkpoints > 0;
|
||||
|
||||
- // make checkpoints only for completion tasks
|
||||
- do_checkpoint = do_checkpoint && slot.task->type == SERVER_TASK_TYPE_COMPLETION;
|
||||
+ // make checkpoints for completion tasks, and for score tasks at the
|
||||
+ // shared-prompt boundary: models whose memory cannot be partially
|
||||
+ // rewound (SWA/hybrid/recurrent) would otherwise re-process the whole
|
||||
+ // prompt for every candidate of a scoring call
|
||||
+ do_checkpoint = do_checkpoint && (slot.task->type == SERVER_TASK_TYPE_COMPLETION ||
|
||||
+ slot.task->type == SERVER_TASK_TYPE_SCORE);
|
||||
|
||||
// make a checkpoint of the parts of the memory that cannot be rolled back.
|
||||
// checkpoints are created only if:
|
||||
@@ -3463,10 +3741,17 @@ private:
|
||||
// embedding requires all tokens in the batch to be output;
|
||||
// MTP also wants logits at every prompt position so the
|
||||
// streaming hook can mirror t_h_nextn into ctx_dft.
|
||||
+ // score tasks need outputs at the positions that predict
|
||||
+ // each candidate token (the token at index i predicts the
|
||||
+ // task token at index i+1).
|
||||
+ const bool need_score_logit =
|
||||
+ slot.task->type == SERVER_TASK_TYPE_SCORE &&
|
||||
+ slot.prompt.n_tokens() + 1 >= slot.task->n_score_prompt &&
|
||||
+ slot.prompt.n_tokens() + 1 < slot.task->n_tokens();
|
||||
add_ok &= batch.add(slot.id,
|
||||
cur_tok,
|
||||
slot.prompt.tokens.pos_next(),
|
||||
- slot.need_embd());
|
||||
+ slot.need_embd() || need_score_logit);
|
||||
slot.prompt.tokens.push_back(cur_tok);
|
||||
|
||||
slot.n_prompt_tokens_processed++;
|
||||
@@ -3481,6 +3766,32 @@ private:
|
||||
}
|
||||
}
|
||||
|
||||
+ // score tasks: break at the shared-prompt boundary so the checkpoint
|
||||
+ // below lands exactly there — the other candidates of the same
|
||||
+ // scoring call re-process only their own tokens. Also break at the
|
||||
+ // point where this task diverged from the previous cache: after a
|
||||
+ // forced re-prefill a checkpoint there serves the next scoring call
|
||||
+ // over the same stable prefix (e.g. a classifier's option list).
|
||||
+ // The caller-declared stable-prefix boundary is the strongest of
|
||||
+ // these: a checkpoint there is at or before every future task's
|
||||
+ // divergence within the same option list, so it always survives
|
||||
+ // and always restores.
|
||||
+ if (do_checkpoint && slot.task->type == SERVER_TASK_TYPE_SCORE &&
|
||||
+ (slot.prompt.n_tokens() == slot.task->n_score_prompt - 1 ||
|
||||
+ (slot.task->n_stable_prompt > 0 &&
|
||||
+ slot.prompt.n_tokens() == slot.task->n_stable_prompt &&
|
||||
+ slot.prompt.n_tokens() < slot.task->n_score_prompt - 1) ||
|
||||
+ (slot.prompt.n_tokens() == slot.score_divergence &&
|
||||
+ slot.prompt.n_tokens() < slot.task->n_score_prompt - 1))) {
|
||||
+ bool have_ckpt = false;
|
||||
+ for (const auto & ckpt : slot.prompt.checkpoints) {
|
||||
+ have_ckpt |= ckpt.n_tokens == slot.prompt.n_tokens();
|
||||
+ }
|
||||
+ if (!have_ckpt) {
|
||||
+ break;
|
||||
+ }
|
||||
+ }
|
||||
+
|
||||
// process the last few tokens of the prompt separately in order to allow for a checkpoint to be created.
|
||||
// create checkpoints that many tokens before the end of the prompt:
|
||||
// - 4 + n_ubatch
|
||||
@@ -3513,6 +3824,15 @@ private:
|
||||
const bool is_user_start = spans.is_user_start(n_tokens_start);
|
||||
const bool is_last_user_message = n_tokens_start == last_user_pos;
|
||||
|
||||
+ // a batch starting at the score boundary or divergence point must
|
||||
+ // always checkpoint — min-step spacing would otherwise suppress it
|
||||
+ // and every candidate / next scoring call would re-process the prompt
|
||||
+ const bool is_score_boundary = slot.task->type == SERVER_TASK_TYPE_SCORE &&
|
||||
+ (n_tokens_start == slot.task->n_score_prompt - 1 ||
|
||||
+ (slot.task->n_stable_prompt > 0 &&
|
||||
+ n_tokens_start == slot.task->n_stable_prompt) ||
|
||||
+ n_tokens_start == slot.score_divergence);
|
||||
+
|
||||
// entire prompt has been processed
|
||||
if (slot.prompt.n_tokens() == slot.task->n_tokens()) {
|
||||
slot.state = SLOT_STATE_DONE_PROMPT;
|
||||
@@ -3528,8 +3848,8 @@ private:
|
||||
slot.init_sampler();
|
||||
} else {
|
||||
// skip ordinary mid-prompt checkpoints, unless the batch starts a user
|
||||
- // message or we are near the end of the prompt
|
||||
- if (!is_user_start && !near_prompt_end) {
|
||||
+ // message, the score boundary, or we are near the end of the prompt
|
||||
+ if (!is_user_start && !is_score_boundary && !near_prompt_end) {
|
||||
do_checkpoint = false;
|
||||
}
|
||||
}
|
||||
@@ -3546,10 +3866,10 @@ private:
|
||||
// do not checkpoint after mtmd chunks
|
||||
do_checkpoint = do_checkpoint && !has_mtmd;
|
||||
|
||||
- // no need to create checkpoints that are too close together, unless it's the last user message
|
||||
+ // no need to create checkpoints that are too close together, unless it's the last user message or the score boundary
|
||||
do_checkpoint = do_checkpoint && (
|
||||
slot.prompt.checkpoints.empty() ||
|
||||
- is_last_user_message || near_prompt_end ||
|
||||
+ is_last_user_message || near_prompt_end || is_score_boundary ||
|
||||
n_tokens_start > slot.prompt.checkpoints.back().n_tokens + params_base.checkpoint_min_step);
|
||||
SLT_DBG(slot, "main/do_checkpoint = %s, pos_min = %d, pos_max = %d\n", do_checkpoint ? "yes" : "no", pos_min, pos_max);
|
||||
|
||||
@@ -3703,6 +4023,13 @@ private:
|
||||
}
|
||||
}
|
||||
|
||||
+ // score slots harvest logprobs from every view that contains
|
||||
+ // their outputs, not just the one holding the final token
|
||||
+ if (slot.task && slot.task->type == SERVER_TASK_TYPE_SCORE &&
|
||||
+ (slot.state == SLOT_STATE_PROCESSING_PROMPT || slot.state == SLOT_STATE_DONE_PROMPT)) {
|
||||
+ collect_score_logprobs(slot, batch_view);
|
||||
+ }
|
||||
+
|
||||
if (!is_inside_view(slot.i_batch)) {
|
||||
// the required token not in this sub-batch, skip
|
||||
return;
|
||||
@@ -3724,6 +4051,25 @@ private:
|
||||
return;
|
||||
}
|
||||
|
||||
+ if (slot.task->type == SERVER_TASK_TYPE_SCORE) {
|
||||
+ // shared-prefix logprobs (and every candidate's first
|
||||
+ // suffix logprob) were accumulated per view above;
|
||||
+ // candidates with more suffix tokens still need the
|
||||
+ // forked decode at the end of update_slots()
|
||||
+ for (const auto & sfx : slot.task->score_suffixes) {
|
||||
+ if (sfx.size() > 1) {
|
||||
+ slot.score_suffix_pending = true;
|
||||
+ break;
|
||||
+ }
|
||||
+ }
|
||||
+ if (!slot.score_suffix_pending) {
|
||||
+ send_score(slot);
|
||||
+ slot.release();
|
||||
+ }
|
||||
+ slot.i_batch = -1;
|
||||
+ return;
|
||||
+ }
|
||||
+
|
||||
GGML_ASSERT(slot.task->need_sampling());
|
||||
|
||||
// prompt evaluated for next-token prediction
|
||||
diff --git a/tools/server/server-task.h b/tools/server/server-task.h
|
||||
index c3eea2e..fb3c178 100644
|
||||
--- a/tools/server/server-task.h
|
||||
+++ b/tools/server/server-task.h
|
||||
@@ -13,10 +13,25 @@
|
||||
|
||||
using json = nlohmann::ordered_json;
|
||||
|
||||
+// SERVER_TASK_TYPE_SCORE emits one logits output per candidate token (plus
|
||||
+// the forced last-token output), and the context's output budget
|
||||
+// (n_outputs_max) is reserved up front — so candidate length must be
|
||||
+// bounded. Raising this raises the worst-case compute-buffer reservation
|
||||
+// by ~n_vocab * 4 bytes per extra output.
|
||||
+constexpr int32_t SERVER_SCORE_MAX_CAND_TOKENS = 64;
|
||||
+
|
||||
+// Maximum sequences forked off the shared prefix in one score suffix
|
||||
+// decode. The context is created with this many seq ids (and
|
||||
+// recurrent-state cells) beyond the parallel slots — see
|
||||
+// common_params::n_seq_score_forks; candidates in excess of the budget
|
||||
+// are decoded in successive chunks.
|
||||
+constexpr int32_t SERVER_SCORE_FORK_SEQS = 16;
|
||||
+
|
||||
enum server_task_type {
|
||||
SERVER_TASK_TYPE_COMPLETION,
|
||||
SERVER_TASK_TYPE_EMBEDDING,
|
||||
SERVER_TASK_TYPE_RERANK,
|
||||
+ SERVER_TASK_TYPE_SCORE,
|
||||
SERVER_TASK_TYPE_INFILL,
|
||||
SERVER_TASK_TYPE_CANCEL,
|
||||
SERVER_TASK_TYPE_CONTROL,
|
||||
@@ -153,6 +168,18 @@ struct server_task {
|
||||
task_params params;
|
||||
server_tokens tokens;
|
||||
|
||||
+ // used by SERVER_TASK_TYPE_SCORE: `tokens` holds the shared prefix
|
||||
+ // (prompt + longest common candidate token prefix) and logprobs are
|
||||
+ // returned for its tokens from n_score_prompt onward. Each candidate's
|
||||
+ // tokens beyond the shared prefix ride a forked sequence.
|
||||
+ int32_t n_score_prompt = 0;
|
||||
+ std::vector<llama_tokens> score_suffixes;
|
||||
+ // token index where the caller-declared stable prompt prefix ends
|
||||
+ // (0 = no hint): the option-list system prompt that repeats across
|
||||
+ // scoring calls. A context checkpoint is forced there so models that
|
||||
+ // cannot rewind state re-process only the per-call tail next time.
|
||||
+ int32_t n_stable_prompt = 0;
|
||||
+
|
||||
// only used by CLI, this allow tokenizing CLI inputs on server side
|
||||
// we need this because mtmd_context and vocab are not accessible outside of server_context
|
||||
bool cli = false;
|
||||
@@ -197,6 +224,7 @@ struct server_task {
|
||||
switch (type) {
|
||||
case SERVER_TASK_TYPE_COMPLETION:
|
||||
case SERVER_TASK_TYPE_INFILL:
|
||||
+ case SERVER_TASK_TYPE_SCORE:
|
||||
return true;
|
||||
default:
|
||||
return false;
|
||||
@@ -494,6 +522,25 @@ struct server_task_result_rerank : server_task_result {
|
||||
virtual json to_json() override;
|
||||
};
|
||||
|
||||
+struct server_task_result_score : server_task_result {
|
||||
+ // log P(token | prefix) for the shared-prefix tokens after
|
||||
+ // n_score_prompt, in order; NaN marks positions the decode never
|
||||
+ // produced an output for
|
||||
+ std::vector<float> shared_logprobs;
|
||||
+
|
||||
+ // per candidate: logprobs of its suffix tokens, in task order (entry
|
||||
+ // 0 is the token right after the shared prefix, predicted by the last
|
||||
+ // shared token's logits)
|
||||
+ std::vector<std::vector<float>> cand_logprobs;
|
||||
+
|
||||
+ virtual json to_json() override {
|
||||
+ return json {
|
||||
+ {"shared_logprobs", shared_logprobs},
|
||||
+ {"cand_logprobs", cand_logprobs},
|
||||
+ };
|
||||
+ }
|
||||
+};
|
||||
+
|
||||
struct server_task_result_error : server_task_result {
|
||||
error_type err_type = ERROR_TYPE_SERVER;
|
||||
std::string err_msg;
|
||||
@@ -47,7 +47,6 @@ define turboquant-build
|
||||
# original under backend/cpp/llama-cpp/, so the stock llama-cpp build
|
||||
# stays compiling against vanilla upstream.
|
||||
bash $(CURRENT_MAKEFILE_DIR)/patch-grpc-server.sh $(CURRENT_MAKEFILE_DIR)/../turboquant-$(1)-build/grpc-server.cpp
|
||||
bash $(LLAMA_CPP_DIR)/disable-score-task.sh $(CURRENT_MAKEFILE_DIR)/../turboquant-$(1)-build/grpc-server.cpp
|
||||
$(info $(GREEN)I turboquant build info:$(1)$(RESET))
|
||||
LLAMA_REPO=$(LLAMA_REPO) LLAMA_VERSION=$(TURBOQUANT_VERSION) \
|
||||
$(MAKE) -C $(CURRENT_MAKEFILE_DIR)/../turboquant-$(1)-build llama.cpp
|
||||
@@ -85,7 +84,6 @@ turboquant-cpu-all:
|
||||
rm -rf $(CURRENT_MAKEFILE_DIR)/../turboquant-cpu-all-build/patches
|
||||
$(MAKE) -C $(CURRENT_MAKEFILE_DIR)/../turboquant-cpu-all-build purge
|
||||
bash $(CURRENT_MAKEFILE_DIR)/patch-grpc-server.sh $(CURRENT_MAKEFILE_DIR)/../turboquant-cpu-all-build/grpc-server.cpp
|
||||
bash $(LLAMA_CPP_DIR)/disable-score-task.sh $(CURRENT_MAKEFILE_DIR)/../turboquant-cpu-all-build/grpc-server.cpp
|
||||
$(info $(GREEN)I turboquant build info:cpu-all-variants$(RESET))
|
||||
LLAMA_REPO=$(LLAMA_REPO) LLAMA_VERSION=$(TURBOQUANT_VERSION) \
|
||||
$(MAKE) -C $(CURRENT_MAKEFILE_DIR)/../turboquant-cpu-all-build llama.cpp
|
||||
|
||||
@@ -32,9 +32,7 @@ import (
|
||||
type anthropicRequest struct {
|
||||
Model string `json:"model"`
|
||||
MaxTokens int32 `json:"max_tokens"`
|
||||
// System is `any`: a bare string normally, or []anthropicSystemBlock
|
||||
// when cache_prompt is on (the block form carries cache_control).
|
||||
System any `json:"system,omitempty"`
|
||||
System string `json:"system,omitempty"`
|
||||
Messages []anthropicMessage `json:"messages"`
|
||||
Stream bool `json:"stream,omitempty"`
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
@@ -54,30 +52,9 @@ type anthropicMessage struct {
|
||||
}
|
||||
|
||||
type anthropicTool struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
InputSchema json.RawMessage `json:"input_schema"`
|
||||
CacheControl *anthropicCacheControl `json:"cache_control,omitempty"`
|
||||
}
|
||||
|
||||
// anthropicCacheControl marks a prompt-cache breakpoint. Anthropic caches
|
||||
// everything up to and including a block tagged {"type":"ephemeral"} (5-min
|
||||
// TTL) and serves that prefix at the cache-read rate (0.1x input) on later
|
||||
// calls that share it — the win on agentic/multi-turn workloads.
|
||||
type anthropicCacheControl struct {
|
||||
Type string `json:"type"` // "ephemeral"
|
||||
}
|
||||
|
||||
// ephemeralCacheControl is the single reused breakpoint marker.
|
||||
var ephemeralCacheControl = &anthropicCacheControl{Type: "ephemeral"}
|
||||
|
||||
// anthropicSystemBlock is the block form of the top-level system field.
|
||||
// Anthropic accepts system as a bare string OR a list of text blocks; the
|
||||
// block form is required to attach cache_control to the system prompt.
|
||||
type anthropicSystemBlock struct {
|
||||
Type string `json:"type"` // "text"
|
||||
Text string `json:"text"`
|
||||
CacheControl *anthropicCacheControl `json:"cache_control,omitempty"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
InputSchema json.RawMessage `json:"input_schema"`
|
||||
}
|
||||
|
||||
// anthropicToolChoice mirrors the four shapes Anthropic accepts:
|
||||
@@ -104,9 +81,8 @@ type anthropicContentBlock struct {
|
||||
// Tool-result block fields. tool_result uses `content` (not
|
||||
// `text`) and pairs with `tool_use_id`; modelling them as
|
||||
// distinct fields avoids ambiguity at marshal time.
|
||||
ToolUseID string `json:"tool_use_id,omitempty"`
|
||||
ResultContent string `json:"content,omitempty"`
|
||||
CacheControl *anthropicCacheControl `json:"cache_control,omitempty"`
|
||||
ToolUseID string `json:"tool_use_id,omitempty"`
|
||||
ResultContent string `json:"content,omitempty"`
|
||||
}
|
||||
|
||||
type anthropicResponse struct {
|
||||
@@ -180,11 +156,6 @@ func buildAnthropicRequest(opts *pb.PredictOptions, cfg *proxyConfig, stream boo
|
||||
if req.ToolChoice != nil && req.ToolChoice.Type == anthropicToolChoiceNone {
|
||||
req.Tools, req.ToolChoice = nil, nil
|
||||
}
|
||||
// Prompt-cache breakpoint on the last tool: Anthropic caches the entire
|
||||
// tool block up to the marked tool — usually a large, fully stable prefix.
|
||||
if cfg.cachePrompt && len(req.Tools) > 0 {
|
||||
req.Tools[len(req.Tools)-1].CacheControl = ephemeralCacheControl
|
||||
}
|
||||
|
||||
var systemParts []string
|
||||
for _, m := range opts.GetMessages() {
|
||||
@@ -218,54 +189,15 @@ func buildAnthropicRequest(opts *pb.PredictOptions, cfg *proxyConfig, stream boo
|
||||
})
|
||||
}
|
||||
}
|
||||
// System: block form (with cache_control) when caching is on, else the
|
||||
// bare string. Only set when non-empty so `omitempty` still drops it.
|
||||
if len(systemParts) > 0 {
|
||||
joined := strings.Join(systemParts, "\n\n")
|
||||
if cfg.cachePrompt {
|
||||
req.System = []anthropicSystemBlock{{Type: "text", Text: joined, CacheControl: ephemeralCacheControl}}
|
||||
} else {
|
||||
req.System = joined
|
||||
}
|
||||
}
|
||||
req.System = strings.Join(systemParts, "\n\n")
|
||||
|
||||
if len(req.Messages) == 0 && opts.GetPrompt() != "" {
|
||||
req.Messages = []anthropicMessage{{Role: "user", Content: opts.GetPrompt()}}
|
||||
}
|
||||
|
||||
// Prompt-cache breakpoint on the final message block caches the whole
|
||||
// conversation prefix up to the newest turn. With the system + tools
|
||||
// breakpoints above, Anthropic serves the entire stable head at the
|
||||
// cache-read rate on the next agentic iteration (max 4 breakpoints; we
|
||||
// use at most 3, so we never exceed the limit).
|
||||
if cfg.cachePrompt {
|
||||
markLastMessageCacheable(req.Messages)
|
||||
}
|
||||
|
||||
return json.Marshal(req)
|
||||
}
|
||||
|
||||
// markLastMessageCacheable tags the final block of the last message with a
|
||||
// cache_control breakpoint. String content is promoted to a single text
|
||||
// block so the marker has somewhere to attach; block content gets the marker
|
||||
// on its last element.
|
||||
func markLastMessageCacheable(msgs []anthropicMessage) {
|
||||
if len(msgs) == 0 {
|
||||
return
|
||||
}
|
||||
last := &msgs[len(msgs)-1]
|
||||
switch c := last.Content.(type) {
|
||||
case string:
|
||||
if c != "" {
|
||||
last.Content = []anthropicContentBlock{{Type: "text", Text: c, CacheControl: ephemeralCacheControl}}
|
||||
}
|
||||
case []anthropicContentBlock:
|
||||
if len(c) > 0 {
|
||||
c[len(c)-1].CacheControl = ephemeralCacheControl
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// appendToolResult appends a tool_result block as a user message,
|
||||
// merging into a preceding user message that already carries blocks.
|
||||
// Anthropic concatenates consecutive same-role messages on its end,
|
||||
|
||||
@@ -328,62 +328,3 @@ func TestBuildAnthropic_RoundTripsAssistantToolCalls(t *testing.T) {
|
||||
g.Expect(r0["tool_use_id"]).To(Equal("call_abc"))
|
||||
g.Expect(r0["content"]).To(Equal(`{"models":["a","b"]}`))
|
||||
}
|
||||
|
||||
// TestPredict_Anthropic_PromptCache verifies that cache_prompt injects
|
||||
// exactly the intended cache_control breakpoints (system, last tool, last
|
||||
// message) when on, and none when off — asserting on the raw upstream body
|
||||
// because System becomes a block list that the typed struct hides.
|
||||
func TestPredict_Anthropic_PromptCache(t *testing.T) {
|
||||
g := NewWithT(t)
|
||||
|
||||
// run issues one translate Predict and returns the raw body the fake
|
||||
// Anthropic upstream received.
|
||||
run := func(cachePrompt bool) string {
|
||||
var rawBody string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
rawBody = string(b)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = io.WriteString(w, `{"id":"m","type":"message","role":"assistant","content":[{"type":"text","text":"ok"}],"model":"claude-3-5-sonnet-20241022","usage":{"input_tokens":5,"output_tokens":2}}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
t.Setenv("CLOUD_PROXY_ANTHROPIC_FAKE", "sk-ant-fake")
|
||||
cp := NewCloudProxy()
|
||||
err := cp.Load(&pb.ModelOptions{
|
||||
Model: "claude-local",
|
||||
Proxy: &pb.ProxyOptions{
|
||||
UpstreamUrl: srv.URL,
|
||||
Mode: modeTranslate,
|
||||
Provider: providerAnthropic,
|
||||
ApiKeyEnv: "CLOUD_PROXY_ANTHROPIC_FAKE",
|
||||
UpstreamModel: "claude-3-5-sonnet-20241022",
|
||||
CachePrompt: cachePrompt,
|
||||
},
|
||||
})
|
||||
g.Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
_, err = cp.Predict(&pb.PredictOptions{
|
||||
Messages: []*pb.Message{
|
||||
{Role: "system", Content: "be brief"},
|
||||
{Role: "user", Content: "hello"},
|
||||
},
|
||||
Tools: `[{"type":"function","function":{"name":"t","parameters":{"type":"object"}}}]`,
|
||||
Tokens: 32,
|
||||
})
|
||||
g.Expect(err).NotTo(HaveOccurred())
|
||||
return rawBody
|
||||
}
|
||||
|
||||
// cache_prompt ON: three ephemeral breakpoints (system + last tool +
|
||||
// last message), and system is emitted in block form.
|
||||
on := run(true)
|
||||
g.Expect(strings.Count(on, `"cache_control":{"type":"ephemeral"}`)).To(Equal(3),
|
||||
"expected 3 breakpoints (system, tool, last message); body=%s", on)
|
||||
g.Expect(on).To(ContainSubstring(`"system":[{"type":"text","text":"be brief"`))
|
||||
|
||||
// cache_prompt OFF: no breakpoints, system stays a bare string.
|
||||
off := run(false)
|
||||
g.Expect(off).NotTo(ContainSubstring("cache_control"))
|
||||
g.Expect(off).To(ContainSubstring(`"system":"be brief"`))
|
||||
}
|
||||
|
||||
@@ -48,7 +48,6 @@ type proxyConfig struct {
|
||||
upstreamModel string
|
||||
localModel string // ModelOptions.Model — fallback when upstream_model is unset
|
||||
apiKey string // resolved at Load time
|
||||
cachePrompt bool // inject Anthropic prompt-cache breakpoints (translate+anthropic)
|
||||
}
|
||||
|
||||
func NewCloudProxy() *CloudProxy {
|
||||
@@ -107,7 +106,6 @@ func (c *CloudProxy) Load(opts *pb.ModelOptions) error {
|
||||
upstreamModel: po.GetUpstreamModel(),
|
||||
localModel: opts.GetModel(),
|
||||
apiKey: key,
|
||||
cachePrompt: po.GetCachePrompt(),
|
||||
})
|
||||
xlog.Info("cloud-proxy: ready",
|
||||
"upstream", po.GetUpstreamUrl(),
|
||||
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# CrispASR version (release tag)
|
||||
CRISPASR_REPO?=https://github.com/CrispStrobe/CrispASR
|
||||
CRISPASR_VERSION?=754b67289cf1137e3ed722885705f94132fc614f
|
||||
CRISPASR_VERSION?=306faee45fab641d54f9f941f075de1e9c0d3278
|
||||
SO_TARGET?=libgocrispasr.so
|
||||
|
||||
CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF
|
||||
|
||||
@@ -14,7 +14,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
# It is kept alive by the upstream tag da2-support (survives a squash-merge);
|
||||
# repoint to the master merge commit once mudler/depth-anything.cpp PR #1 lands.
|
||||
DEPTHANYTHING_REPO?=https://github.com/mudler/depth-anything.cpp.git
|
||||
DEPTHANYTHING_VERSION?=2028b47ac75a8659c6a9aa617baf09be193eb55f
|
||||
DEPTHANYTHING_VERSION?=f4e17dea695dd12ae76bea98ba58030996b98118
|
||||
|
||||
ifeq ($(NATIVE),false)
|
||||
CMAKE_ARGS+=-DGGML_NATIVE=OFF
|
||||
|
||||
10
backend/go/dllm/.gitignore
vendored
10
backend/go/dllm/.gitignore
vendored
@@ -1,10 +0,0 @@
|
||||
.cache/
|
||||
sources/
|
||||
build/
|
||||
package/
|
||||
dllm-grpc
|
||||
# build artifacts staged in-tree by the Makefile (cp from sources/) or
|
||||
# symlinked for local dev; the real sources live in dllm.cpp upstream.
|
||||
*.so
|
||||
*.so.*
|
||||
compile_commands.json
|
||||
@@ -1,101 +0,0 @@
|
||||
# dllm backend Makefile.
|
||||
#
|
||||
# Upstream pin lives below as DLLM_VERSION?=<sha> so .github/bump_deps.sh
|
||||
# can find and update it - matches the whisper.cpp / parakeet-cpp / ds4
|
||||
# convention.
|
||||
#
|
||||
# Local dev shortcut: if you already have an out-of-tree dllm.cpp build,
|
||||
# you can symlink the .so into this directory and skip the clone/cmake
|
||||
# steps entirely, e.g.:
|
||||
#
|
||||
# ln -sf /path/to/dllm.cpp/build/libdllm.so .
|
||||
# go build -o dllm-grpc .
|
||||
#
|
||||
# That's what the gated C-ABI binding smoke uses (DLLM_TEST_LIBRARY). The
|
||||
# default target below does the proper clone-at-pin + cmake build so CI
|
||||
# doesn't need a side-checkout.
|
||||
#
|
||||
# NOTE: github.com/mudler/dllm.cpp is still private (publishing is planned);
|
||||
# until then the anonymous clone below fails. Use the symlink shortcut above
|
||||
# with a local checkout, or a git credential helper with access to the repo.
|
||||
|
||||
# The pin below is the P5 performance-parity head (device-resident
|
||||
# self-conditioning, full-GPU placement at ngl >= n_layer, graph reuse,
|
||||
# device-side EB reductions: ~8x per-step on GB10, see dllm.cpp
|
||||
# docs/validation.md section 10). C-ABI unchanged (still version 1). It
|
||||
# also carries the multimodal entry points (dllm_capi_generate_mm /
|
||||
# dllm_capi_generate_stream_mm) the image-input path probes for; older
|
||||
# libs still load, but image requests then fail with "library predates
|
||||
# the multimodal entry points".
|
||||
DLLM_VERSION?=320b57756efc3460169b8ea9e8c782867198f2a5
|
||||
DLLM_REPO?=https://github.com/mudler/dllm.cpp
|
||||
|
||||
GOCMD?=go
|
||||
GO_TAGS?=
|
||||
JOBS?=$(shell nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4)
|
||||
|
||||
BUILD_TYPE?=
|
||||
NATIVE?=false
|
||||
|
||||
# libdllm.so is self-contained: dllm.cpp's CMakeLists statically absorbs ggml
|
||||
# (BUILD_SHARED_LIBS=OFF + PIC) into the shared lib, so dlopen needs no
|
||||
# libggml*.so alongside it, only system libs (libstdc++/libgomp/libc) the
|
||||
# runtime image already provides. Tests/CLI are upstream-only concerns.
|
||||
CMAKE_ARGS?=-DCMAKE_BUILD_TYPE=Release -DDLLM_BUILD_TESTS=OFF
|
||||
|
||||
ifeq ($(NATIVE),false)
|
||||
CMAKE_ARGS+=-DGGML_NATIVE=OFF
|
||||
endif
|
||||
|
||||
# Same arch set the sibling ggml backends (acestep/vibevoice/qwen3-tts) bake
|
||||
# for their cublas images; override for a native build.
|
||||
CUDA_ARCHITECTURES?=75-virtual;80-virtual;86-real;89-real
|
||||
|
||||
# dllm.cpp gates CUDA behind DLLM_CUDA (set(GGML_CUDA ... CACHE FORCE)), so
|
||||
# forward that instead of a bare -DGGML_CUDA=ON.
|
||||
ifeq ($(BUILD_TYPE),cublas)
|
||||
CMAKE_ARGS+=-DDLLM_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES="$(CUDA_ARCHITECTURES)"
|
||||
endif
|
||||
|
||||
.PHONY: dllm-grpc package build clean purge test all
|
||||
|
||||
all: dllm-grpc
|
||||
|
||||
# Clone the upstream dllm.cpp source at the pinned commit (ggml comes in as
|
||||
# a submodule). Directory acts as the target so make only re-clones when
|
||||
# missing. After a DLLM_VERSION bump, run 'make purge && make' to refetch.
|
||||
sources/dllm.cpp:
|
||||
mkdir -p sources/dllm.cpp
|
||||
cd sources/dllm.cpp && \
|
||||
git init -q && \
|
||||
git remote add origin $(DLLM_REPO) && \
|
||||
git fetch --depth 1 origin $(DLLM_VERSION) && \
|
||||
git checkout FETCH_HEAD && \
|
||||
git submodule update --init --recursive --depth 1 --single-branch
|
||||
|
||||
# Build the shared lib out-of-tree, then stage it next to the Go sources so
|
||||
# purego.Dlopen("libdllm.so") and the packaging step both pick it up.
|
||||
libdllm.so: sources/dllm.cpp
|
||||
cmake -B sources/dllm.cpp/build -S sources/dllm.cpp $(CMAKE_ARGS)
|
||||
cmake --build sources/dllm.cpp/build --config Release -j$(JOBS)
|
||||
cp -fv sources/dllm.cpp/build/libdllm.so ./
|
||||
|
||||
dllm-grpc: libdllm.so main.go capi.go
|
||||
CGO_ENABLED=0 $(GOCMD) build -tags "$(GO_TAGS)" -o dllm-grpc .
|
||||
|
||||
package: dllm-grpc
|
||||
bash package.sh
|
||||
|
||||
build: package
|
||||
|
||||
# Test target. The C-ABI binding smoke is gated on DLLM_TEST_LIBRARY +
|
||||
# DLLM_TEST_TINY_MODEL; without them the gated specs auto-skip and only the
|
||||
# pure-Go helper specs run.
|
||||
test:
|
||||
LD_LIBRARY_PATH=$(CURDIR):$$LD_LIBRARY_PATH $(GOCMD) test ./... -count=1
|
||||
|
||||
clean: purge
|
||||
rm -rf libdllm.so* package dllm-grpc
|
||||
|
||||
purge:
|
||||
rm -rf sources/dllm.cpp
|
||||
@@ -1,326 +0,0 @@
|
||||
package main
|
||||
|
||||
// Typed Go wrappers over dllm.cpp's flat C-ABI (include/dllm_capi.h, ABI v1).
|
||||
//
|
||||
// Contract highlights the wrappers encode (see the header + src/capi.cpp):
|
||||
// - tokenize_json/generate return malloc'd char* the CALLER owns: bound as
|
||||
// uintptr, copied with goStringFromCPtr, released via dllm_capi_free_string.
|
||||
// - last_error returns a BORROWED pointer (valid until the next call on the
|
||||
// same ctx): bound as a plain string (purego copies), never freed, and only
|
||||
// read AFTER the failing call has returned - reading it while a generate is
|
||||
// in flight on the same ctx violates the per-ctx serialization contract.
|
||||
// - All entry points except dllm_capi_cancel must be externally serialized
|
||||
// per ctx (one ctx = one concurrent generate/tokenize). Cancel only flips
|
||||
// an atomic and may be called from any goroutine mid-generate.
|
||||
// - No C++ exception crosses the boundary; failures land in last_error.
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"unsafe"
|
||||
|
||||
"github.com/ebitengine/purego"
|
||||
)
|
||||
|
||||
// dllmABIVersion is the DLLM_CAPI_ABI_VERSION this binding was written
|
||||
// against; main.go refuses to start against a libdllm.so reporting another.
|
||||
const dllmABIVersion = 1
|
||||
|
||||
// purego-bound entry points from libdllm.so. Names match dllm_capi.h
|
||||
// exactly; loadCAPI (main.go) fills these in at boot.
|
||||
var (
|
||||
cppAbiVersion func() int32
|
||||
cppLoad func(ggufPath, paramsJSON string) uintptr
|
||||
cppFree func(ctx uintptr)
|
||||
cppLastError func(ctx uintptr) string // borrowed pointer: purego copies, do NOT free
|
||||
cppFreeString func(s uintptr)
|
||||
// malloc'd char* returns, hence uintptr (see loadCAPI's doc comment).
|
||||
cppTokenizeJSON func(ctx uintptr, text string) uintptr
|
||||
cppGenerate func(ctx uintptr, prompt, optsJSON string) uintptr
|
||||
// on_block/on_step are C function pointers produced by purego.NewCallback;
|
||||
// userData carries the streamCallStates registry key.
|
||||
cppGenerateStream func(ctx uintptr, prompt, optsJSON string, onBlock, onStep, userData uintptr) int32
|
||||
cppCancel func(ctx uintptr)
|
||||
)
|
||||
|
||||
// Optional multimodal entry points (dllm_capi.h's P4 surface). The ABI
|
||||
// version stays 1: presence is detected by PROBING the symbols with Dlsym at
|
||||
// boot (loadCAPI, mirroring the parakeet-cpp optional-symbol pattern). nil
|
||||
// means the loaded libdllm.so predates the mm surface; the wrappers below
|
||||
// then fail with errMMUnsupported instead of crashing on a nil call.
|
||||
var (
|
||||
cppGenerateMM func(ctx uintptr, prompt, imagesJSON, optsJSON string) uintptr
|
||||
cppGenerateStreamMM func(ctx uintptr, prompt, imagesJSON, optsJSON string, onBlock, onStep, userData uintptr) int32
|
||||
)
|
||||
|
||||
// mmImageMarker is the literal placeholder dllm_capi_generate_mm expands to
|
||||
// <boi> + soft-token placeholders + <eoi> (dllm_capi.h placeholder contract;
|
||||
// capi.cpp MM_MARKER). The prompt must carry exactly one marker per
|
||||
// images_json entry, in image order.
|
||||
const mmImageMarker = "<image>"
|
||||
|
||||
// errMMUnsupported is returned for image-bearing requests against an old
|
||||
// text-only libdllm.so (the Dlsym probe found no mm symbols).
|
||||
var errMMUnsupported = errors.New(
|
||||
"dllm: image input requires libdllm.so with the multimodal entry points (dllm_capi_generate_mm), but the loaded library predates them - rebuild/upgrade the dllm backend to use images")
|
||||
|
||||
// cMMSupported reports whether the loaded libdllm.so carries the multimodal
|
||||
// generate pair. Both symbols ship together (same dllm.cpp commit), but the
|
||||
// guard requires both anyway so a half-present surface can never dispatch.
|
||||
func cMMSupported() bool {
|
||||
return cppGenerateMM != nil && cppGenerateStreamMM != nil
|
||||
}
|
||||
|
||||
// cAbiVersion returns the library's DLLM_CAPI_ABI_VERSION.
|
||||
func cAbiVersion() int32 {
|
||||
return cppAbiVersion()
|
||||
}
|
||||
|
||||
// cLoad opens the GGUF at path with the flat params JSON (e.g.
|
||||
// {"n_gpu_layers":99}). Returns 0 on failure; per the header contract there
|
||||
// is no ctx to carry the reason, the C side logs it to stderr (and
|
||||
// cLastError(0) only yields the static NULL-ctx message).
|
||||
func cLoad(path, paramsJSON string) uintptr {
|
||||
return cppLoad(path, paramsJSON)
|
||||
}
|
||||
|
||||
// cFree releases a ctx; safe on 0 (delete nullptr).
|
||||
func cFree(h uintptr) {
|
||||
cppFree(h)
|
||||
}
|
||||
|
||||
// cLastError returns the ctx's last error message (or the static NULL-ctx
|
||||
// message for h==0). The C pointer is borrowed and only valid until the next
|
||||
// call on the same ctx; purego's string return copies it immediately, so the
|
||||
// returned Go string is safe to keep. Must not be called while another call
|
||||
// on the same ctx is in flight.
|
||||
func cLastError(h uintptr) string {
|
||||
return cppLastError(h)
|
||||
}
|
||||
|
||||
// lastErrorOr is cLastError with a fallback for the empty-message case, so
|
||||
// wrapped errors never end in ": ".
|
||||
func lastErrorOr(h uintptr, fallback string) string {
|
||||
if msg := cLastError(h); msg != "" {
|
||||
return msg
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
// cTokenizeJSON tokenizes text (the C side prepends bos per vocab.add_bos)
|
||||
// and returns the token ids as a JSON array string, e.g. "[2,18]".
|
||||
func cTokenizeJSON(h uintptr, text string) (string, error) {
|
||||
ret := cppTokenizeJSON(h, text)
|
||||
if ret == 0 {
|
||||
return "", fmt.Errorf("dllm: tokenize failed: %s", lastErrorOr(h, "unknown error"))
|
||||
}
|
||||
out := goStringFromCPtr(ret)
|
||||
cppFreeString(ret)
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// cGenerate runs a blocking generation and returns the detokenized text.
|
||||
// optsJSON must be a FLAT JSON object of scalars (use buildOptsJSON); the C
|
||||
// parser rejects nested objects/arrays. NULL return -> last_error (read only
|
||||
// after the call returned, per the serialization contract); a cancelled call
|
||||
// surfaces as the "cancelled" message.
|
||||
func cGenerate(h uintptr, prompt, optsJSON string) (string, error) {
|
||||
ret := cppGenerate(h, prompt, optsJSON)
|
||||
if ret == 0 {
|
||||
return "", fmt.Errorf("dllm: generate failed: %s", lastErrorOr(h, "unknown error"))
|
||||
}
|
||||
out := goStringFromCPtr(ret)
|
||||
cppFreeString(ret)
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// cGenerateMM is cGenerate's multimodal counterpart. imagesJSON is the flat
|
||||
// JSON array of image entries (data: base64 URIs here; the C side also takes
|
||||
// file paths) and the prompt must carry one mmImageMarker per entry - the
|
||||
// engine enforces the 1:1 match and reports mismatches through last_error.
|
||||
func cGenerateMM(h uintptr, prompt, imagesJSON, optsJSON string) (string, error) {
|
||||
if !cMMSupported() {
|
||||
return "", errMMUnsupported
|
||||
}
|
||||
ret := cppGenerateMM(h, prompt, imagesJSON, optsJSON)
|
||||
if ret == 0 {
|
||||
return "", fmt.Errorf("dllm: generate_mm failed: %s", lastErrorOr(h, "unknown error"))
|
||||
}
|
||||
out := goStringFromCPtr(ret)
|
||||
cppFreeString(ret)
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// streamCallState carries the Go callbacks for one in-flight
|
||||
// cGenerateStream call; the registry key travels through C as user_data.
|
||||
// The map shape mirrors the whisper backend's streamCallStates: only one
|
||||
// entry per ctx is ever live (the C-ABI is serialized per ctx), but keying
|
||||
// by call survives multiple models/processes sharing the package.
|
||||
type streamCallState struct {
|
||||
onBlock func(text string)
|
||||
onStep func(step, total int, preview string)
|
||||
}
|
||||
|
||||
var (
|
||||
streamCallStates sync.Map // uint64 -> *streamCallState
|
||||
streamCallSeq atomic.Uint64
|
||||
|
||||
// purego.NewCallback allocates a finite, never-released callback slot, so
|
||||
// the two trampolines are created exactly once and reused across calls.
|
||||
streamCbOnce sync.Once
|
||||
blockCbPtr uintptr
|
||||
stepCbPtr uintptr
|
||||
)
|
||||
|
||||
// onBlockTrampoline is the Go side of dllm_block_cb. It runs on the C
|
||||
// calling thread, mid-generate: keep it tiny and non-blocking (callers that
|
||||
// bridge to goroutines must hand off via buffered channels). The text
|
||||
// pointer is only valid for the duration of the invocation, so it is copied
|
||||
// to a Go string immediately.
|
||||
func onBlockTrampoline(text uintptr, userData uintptr) {
|
||||
v, ok := streamCallStates.Load(uint64(userData))
|
||||
if !ok {
|
||||
return // call already torn down
|
||||
}
|
||||
state := v.(*streamCallState)
|
||||
if state.onBlock != nil {
|
||||
state.onBlock(goStringFromCPtr(text))
|
||||
}
|
||||
}
|
||||
|
||||
// onStepTrampoline is the Go side of dllm_step_cb; same threading and
|
||||
// lifetime caveats as onBlockTrampoline.
|
||||
func onStepTrampoline(step int32, totalSteps int32, canvasPreview uintptr, userData uintptr) {
|
||||
v, ok := streamCallStates.Load(uint64(userData))
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
state := v.(*streamCallState)
|
||||
if state.onStep != nil {
|
||||
state.onStep(int(step), int(totalSteps), goStringFromCPtr(canvasPreview))
|
||||
}
|
||||
}
|
||||
|
||||
// withStreamCallbacks registers onBlock/onStep in the trampoline registry
|
||||
// for the duration of one streaming C call and invokes call with the C
|
||||
// function pointers (NULL for absent callbacks, so the C side skips the
|
||||
// per-block / per-step detokenize work entirely) plus the registry key to
|
||||
// pass as user_data. Shared by the text and multimodal stream wrappers.
|
||||
func withStreamCallbacks(onBlock func(text string), onStep func(step, total int, preview string), call func(blockPtr, stepPtr, userData uintptr) int32) int32 {
|
||||
streamCbOnce.Do(func() {
|
||||
blockCbPtr = purego.NewCallback(onBlockTrampoline)
|
||||
stepCbPtr = purego.NewCallback(onStepTrampoline)
|
||||
})
|
||||
|
||||
id := streamCallSeq.Add(1)
|
||||
streamCallStates.Store(id, &streamCallState{onBlock: onBlock, onStep: onStep})
|
||||
defer streamCallStates.Delete(id)
|
||||
|
||||
var blockPtr, stepPtr uintptr
|
||||
if onBlock != nil {
|
||||
blockPtr = blockCbPtr
|
||||
}
|
||||
if onStep != nil {
|
||||
stepPtr = stepCbPtr
|
||||
}
|
||||
return call(blockPtr, stepPtr, uintptr(id))
|
||||
}
|
||||
|
||||
// cGenerateStream runs a generation with per-committed-block (onBlock) and
|
||||
// per-denoising-step (onStep) callbacks; either may be nil. The callbacks
|
||||
// run on the C thread (see the trampoline docs). Returns an error carrying
|
||||
// last_error on failure; cancellation surfaces as the "cancelled" message.
|
||||
func cGenerateStream(h uintptr, prompt, optsJSON string, onBlock func(text string), onStep func(step, total int, preview string)) error {
|
||||
rc := withStreamCallbacks(onBlock, onStep, func(blockPtr, stepPtr, userData uintptr) int32 {
|
||||
return cppGenerateStream(h, prompt, optsJSON, blockPtr, stepPtr, userData)
|
||||
})
|
||||
if rc != 0 {
|
||||
return fmt.Errorf("dllm: generate_stream failed: %s", lastErrorOr(h, "unknown error"))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// cGenerateStreamMM is cGenerateStream's multimodal counterpart; see
|
||||
// cGenerateMM for the imagesJSON/marker contract.
|
||||
func cGenerateStreamMM(h uintptr, prompt, imagesJSON, optsJSON string, onBlock func(text string), onStep func(step, total int, preview string)) error {
|
||||
if !cMMSupported() {
|
||||
return errMMUnsupported
|
||||
}
|
||||
rc := withStreamCallbacks(onBlock, onStep, func(blockPtr, stepPtr, userData uintptr) int32 {
|
||||
return cppGenerateStreamMM(h, prompt, imagesJSON, optsJSON, blockPtr, stepPtr, userData)
|
||||
})
|
||||
if rc != 0 {
|
||||
return fmt.Errorf("dllm: generate_stream_mm failed: %s", lastErrorOr(h, "unknown error"))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// cCancel requests cancellation of the in-flight generate on h. This is the
|
||||
// ONE entry point safe to call from any goroutine while a generate runs (it
|
||||
// only flips an atomic). Note the cancel-reset race from the header: each
|
||||
// generate resets the flag on entry, so a watchdog should re-issue cancel if
|
||||
// the call has not returned.
|
||||
func cCancel(h uintptr) {
|
||||
cppCancel(h)
|
||||
}
|
||||
|
||||
// buildOptsJSON renders generation options as the flat JSON object the
|
||||
// C-ABI expects (known keys: n_predict, blocks, seed, eb_*, kv_cache). The
|
||||
// C-side scanner only understands scalar number/string values and rejects
|
||||
// nested objects/arrays loudly; bools are rejected here too because the
|
||||
// scanner has no concept of them. Fail loud rather than let an option be
|
||||
// silently misread.
|
||||
//
|
||||
// CAVEAT: json.Marshal HTML-escapes <, > and & inside string values (e.g.
|
||||
// "<" becomes the six-byte \u003c sequence). None of the known string-valued keys
|
||||
// (kv_cache: auto|on|off) can contain those bytes today; if one ever does,
|
||||
// switch to an Encoder with SetEscapeHTML(false) like gemma4JSONString.
|
||||
func buildOptsJSON(opts map[string]any) (string, error) {
|
||||
if len(opts) == 0 {
|
||||
return "{}", nil
|
||||
}
|
||||
for k, v := range opts {
|
||||
switch v.(type) {
|
||||
case string,
|
||||
int, int8, int16, int32, int64,
|
||||
uint, uint8, uint16, uint32, uint64,
|
||||
float32, float64,
|
||||
json.Number:
|
||||
// scalar: fine
|
||||
default:
|
||||
return "", fmt.Errorf("dllm: opts key %q has non-scalar value %T (the C-ABI only accepts flat number/string scalars)", k, v)
|
||||
}
|
||||
}
|
||||
b, err := json.Marshal(opts)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("dllm: marshal opts: %w", err)
|
||||
}
|
||||
return string(b), nil
|
||||
}
|
||||
|
||||
// goStringFromCPtr copies a NUL-terminated C string into Go memory. cptr is
|
||||
// the raw pointer returned by purego from the C-ABI (a malloc'd buffer the
|
||||
// caller owns, or a callback argument only valid during the invocation);
|
||||
// owning callers must free it via cppFreeString after the copy lands.
|
||||
//
|
||||
// A direct unsafe.Pointer(cptr) conversion trips go vet's unsafeptr check,
|
||||
// which can't distinguish a C-owned heap pointer from Go-managed memory (the
|
||||
// parakeet-cpp and whisper backends tolerate that warning). Reinterpreting
|
||||
// through &cptr below is equivalent at runtime and keeps plain `go vet`
|
||||
// clean. It is safe either way: the pointer addresses C memory the Go GC
|
||||
// neither tracks nor moves, and we dereference it immediately to copy the
|
||||
// bytes out.
|
||||
func goStringFromCPtr(cptr uintptr) string {
|
||||
if cptr == 0 {
|
||||
return ""
|
||||
}
|
||||
p := *(*unsafe.Pointer)(unsafe.Pointer(&cptr)) // C-owned buffer, not Go-GC memory (see doc above)
|
||||
n := 0
|
||||
for *(*byte)(unsafe.Add(p, n)) != 0 {
|
||||
n++
|
||||
}
|
||||
return string(unsafe.Slice((*byte)(p), n))
|
||||
}
|
||||
@@ -1,622 +0,0 @@
|
||||
package main
|
||||
|
||||
// LocalAI gRPC backend for dllm.cpp (DiffusionGemma block-diffusion models).
|
||||
//
|
||||
// Wiring overview:
|
||||
// - Load opens the GGUF via dllm_capi_load and starts the per-model worker
|
||||
// goroutine that serializes every C call (see submit).
|
||||
// - PredictRich / PredictStreamRich implement grpc.AIModelRich: when the
|
||||
// request carries raw messages (use_tokenizer_template), the backend owns
|
||||
// templating (RenderGemma4) and output parsing (Gemma4Parser) and replies
|
||||
// with ChatDeltas, like the llama.cpp autoparser and the ds4 backend.
|
||||
// - The legacy Predict / PredictStream methods delegate to the rich pair
|
||||
// (cloud-proxy precedent); the gRPC server prefers the rich path anyway.
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"unicode/utf8"
|
||||
|
||||
grpc "github.com/mudler/LocalAI/pkg/grpc"
|
||||
"github.com/mudler/LocalAI/pkg/grpc/base"
|
||||
"github.com/mudler/LocalAI/pkg/grpc/grpcerrors"
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// The gRPC server cancels in-flight generations on client disconnect only
|
||||
// for backends advertising the Cancellable capability; keep Dllm pinned to
|
||||
// it so a signature drift fails the build, not the disconnect path.
|
||||
var _ grpc.Cancellable = (*Dllm)(nil)
|
||||
|
||||
// generator is the seam between the backend wiring and the dllm.cpp C-ABI:
|
||||
// the real implementation (capiGenerator) wraps the cGenerate/cTokenizeJSON
|
||||
// family, while tests substitute a fake to exercise prompt construction,
|
||||
// parsing and serialization without libdllm.so.
|
||||
type generator interface {
|
||||
generate(prompt, optsJSON string) (string, error)
|
||||
// generateStream invokes onBlock once per committed diffusion block, on
|
||||
// the thread running the C call, before returning.
|
||||
generateStream(prompt, optsJSON string, onBlock func(text string)) error
|
||||
// generateMM / generateStreamMM are the multimodal counterparts:
|
||||
// imagesJSON is a flat JSON array of data: base64 URIs and the prompt
|
||||
// carries one mmImageMarker per entry (dllm_capi.h placeholder
|
||||
// contract). Against an old text-only libdllm.so they fail with
|
||||
// errMMUnsupported.
|
||||
generateMM(prompt, imagesJSON, optsJSON string) (string, error)
|
||||
generateStreamMM(prompt, imagesJSON, optsJSON string, onBlock func(text string)) error
|
||||
tokenizeJSON(text string) (string, error)
|
||||
// cancel is the ONE entry point safe to call concurrently with an
|
||||
// in-flight generate on the same ctx (dllm_capi.h: it only flips an
|
||||
// atomic; everything else must be externally serialized per ctx).
|
||||
cancel()
|
||||
free()
|
||||
}
|
||||
|
||||
// capiGenerator is the production generator over one dllm_ctx handle.
|
||||
type capiGenerator struct {
|
||||
h uintptr
|
||||
}
|
||||
|
||||
func (g *capiGenerator) generate(prompt, optsJSON string) (string, error) {
|
||||
return cGenerate(g.h, prompt, optsJSON)
|
||||
}
|
||||
|
||||
func (g *capiGenerator) generateStream(prompt, optsJSON string, onBlock func(text string)) error {
|
||||
// on_step (per-denoise-step canvas preview, dllm.cpp's --visual) is
|
||||
// passed as nil for now: a future progress hook for the React UI can
|
||||
// plumb it through without touching the C binding.
|
||||
return cGenerateStream(g.h, prompt, optsJSON, onBlock, nil)
|
||||
}
|
||||
|
||||
func (g *capiGenerator) generateMM(prompt, imagesJSON, optsJSON string) (string, error) {
|
||||
return cGenerateMM(g.h, prompt, imagesJSON, optsJSON)
|
||||
}
|
||||
|
||||
func (g *capiGenerator) generateStreamMM(prompt, imagesJSON, optsJSON string, onBlock func(text string)) error {
|
||||
// on_step is nil for the same reason as generateStream.
|
||||
return cGenerateStreamMM(g.h, prompt, imagesJSON, optsJSON, onBlock, nil)
|
||||
}
|
||||
|
||||
func (g *capiGenerator) tokenizeJSON(text string) (string, error) {
|
||||
return cTokenizeJSON(g.h, text)
|
||||
}
|
||||
|
||||
func (g *capiGenerator) cancel() {
|
||||
cCancel(g.h)
|
||||
}
|
||||
|
||||
func (g *capiGenerator) free() {
|
||||
cFree(g.h)
|
||||
}
|
||||
|
||||
// Dllm is the gRPC backend instance: one per loaded model (LocalAI starts
|
||||
// one backend process per model).
|
||||
type Dllm struct {
|
||||
base.Base
|
||||
|
||||
gen generator
|
||||
// genOpts holds the model-level generation overrides parsed from
|
||||
// ModelOptions.Options at Load (eb_*, blocks, kv_cache). The C-ABI takes
|
||||
// them per-generate, not per-load, so they are merged into every
|
||||
// request's opts JSON (requestOptsJSON).
|
||||
genOpts map[string]any
|
||||
|
||||
// jobs is the per-model worker queue. dllm_capi.h requires every entry
|
||||
// point EXCEPT dllm_capi_cancel to be externally serialized per ctx (one
|
||||
// ctx = one concurrent generate/tokenize; last_error is unsafe to read
|
||||
// while a call is in flight). A single goroutine owning all C calls makes
|
||||
// that contract structural instead of relying on lock discipline.
|
||||
jobs chan func()
|
||||
workerWG sync.WaitGroup
|
||||
|
||||
// genMu guards gen against Free racing in-flight requests: requests hold
|
||||
// the read lock for their full duration (they stay concurrent with each
|
||||
// other - the worker still serializes the C calls), Free takes the write
|
||||
// lock so it can only run when no request is in flight.
|
||||
genMu sync.RWMutex
|
||||
}
|
||||
|
||||
func (d *Dllm) startWorker() {
|
||||
d.jobs = make(chan func())
|
||||
d.workerWG.Add(1)
|
||||
go func() {
|
||||
defer d.workerWG.Done()
|
||||
for job := range d.jobs {
|
||||
job()
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// submit runs job on the worker goroutine and waits for it to finish.
|
||||
// Concurrent gRPC requests therefore queue up and execute one at a time
|
||||
// against the single dllm_ctx.
|
||||
func (d *Dllm) submit(job func()) {
|
||||
done := make(chan struct{})
|
||||
d.jobs <- func() {
|
||||
defer close(done)
|
||||
job()
|
||||
}
|
||||
<-done
|
||||
}
|
||||
|
||||
// Load opens the GGUF and prepares the worker. Load-time engine parameters
|
||||
// travel as the flat params JSON of dllm_capi_load; generation overrides
|
||||
// from Options are stored for per-request opts JSON instead (the C-ABI has
|
||||
// no per-load sampler state).
|
||||
func (d *Dllm) Load(opts *pb.ModelOptions) error {
|
||||
if d.gen != nil {
|
||||
return errors.New("dllm: model already loaded")
|
||||
}
|
||||
|
||||
params := map[string]any{
|
||||
"n_gpu_layers": opts.GetNGPULayers(),
|
||||
}
|
||||
if opts.GetThreads() > 0 {
|
||||
params["n_threads"] = opts.GetThreads()
|
||||
}
|
||||
if opts.GetContextSize() > 0 {
|
||||
params["ctx_len"] = opts.GetContextSize()
|
||||
}
|
||||
paramsJSON, err := buildOptsJSON(params)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
d.genOpts = parseModelGenOpts(opts.GetOptions())
|
||||
|
||||
h := cLoad(opts.GetModelFile(), paramsJSON)
|
||||
if h == 0 {
|
||||
// No ctx exists on load failure, so last_error(NULL) only carries the
|
||||
// static NULL-ctx message; the real reason is on the backend's stderr.
|
||||
return fmt.Errorf("dllm: load %q failed: %s (see backend log for details)",
|
||||
opts.GetModelFile(), lastErrorOr(0, "unknown error"))
|
||||
}
|
||||
d.gen = &capiGenerator{h: h}
|
||||
d.startWorker()
|
||||
xlog.Info("dllm: model loaded", "model", opts.GetModelFile(), "params", paramsJSON, "gen_opts", d.genOpts)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Free releases the dllm ctx and stops the worker. Safe when never loaded.
|
||||
//
|
||||
// The write lock is essential: the gRPC server (pkg/grpc/server.go, see the
|
||||
// model-unload path around line 764) calls Free with no locking of its own,
|
||||
// and base.Base provides none either. Without it a request racing Free would
|
||||
// panic sending on the closed jobs channel - or worse, generate on a freed C
|
||||
// ctx. Holding genMu until gen is nil also turns post-Free requests into a
|
||||
// clean "model not loaded" error instead of a crash.
|
||||
func (d *Dllm) Free() error {
|
||||
d.genMu.Lock()
|
||||
defer d.genMu.Unlock()
|
||||
if d.gen == nil {
|
||||
return nil
|
||||
}
|
||||
d.submit(d.gen.free)
|
||||
close(d.jobs)
|
||||
d.workerWG.Wait()
|
||||
d.gen = nil
|
||||
return nil
|
||||
}
|
||||
|
||||
// Cancel requests cancellation of the in-flight generate (the
|
||||
// grpc.Cancellable capability). The gRPC server arms it via
|
||||
// context.AfterFunc on the request/stream context, so a client
|
||||
// disconnect or timeout aborts the generation server-side - the same
|
||||
// semantics the llama.cpp C++ backend gets from polling IsCancelled().
|
||||
// It deliberately bypasses the worker queue: dllm_capi_cancel is the one
|
||||
// call the C-ABI allows from any goroutine mid-generate (it only flips
|
||||
// an atomic).
|
||||
//
|
||||
// Note dllm_capi.h's cancel-reset race: each generate resets the flag on
|
||||
// entry, so a Cancel racing a NEW generate on the same ctx can be lost
|
||||
// (and, with requests queued on the worker, it aborts whichever generate
|
||||
// is currently running). The single-flag granularity is acceptable here
|
||||
// because the server de-registers the hook on normal completion and one
|
||||
// backend process serves one model.
|
||||
func (d *Dllm) Cancel() {
|
||||
// RLock so a server-side AfterFunc firing in the window between a
|
||||
// request finishing and a model unload cannot touch a freed C ctx
|
||||
// (Free holds the write lock while tearing gen down). cancel() is the
|
||||
// one C call that is safe concurrently with an in-flight generate, so
|
||||
// taking a read lock here cannot deadlock against request holders.
|
||||
d.genMu.RLock()
|
||||
defer d.genMu.RUnlock()
|
||||
if d.gen != nil {
|
||||
d.gen.cancel()
|
||||
}
|
||||
}
|
||||
|
||||
// dllmGenOptKeys are the ModelOptions.Options keys this backend forwards to
|
||||
// the engine. Options is a shared free-form bag (other layers put their own
|
||||
// entries there), so unknown keys are skipped with a warning, not an error.
|
||||
var dllmGenOptKeys = map[string]bool{
|
||||
"blocks": true,
|
||||
"kv_cache": true, // "auto"|"on"|"off"; honored by the engine from P3
|
||||
}
|
||||
|
||||
// parseModelGenOpts parses "key:value" Options entries into the flat scalar
|
||||
// map merged into every generate's opts JSON. eb_* (Entropy-Bound sampler
|
||||
// knobs) and the keys in dllmGenOptKeys are recognized; values are typed by
|
||||
// first successful parse (int, then float, else string) to match the C
|
||||
// scanner's number/string scalars.
|
||||
func parseModelGenOpts(options []string) map[string]any {
|
||||
out := map[string]any{}
|
||||
for _, o := range options {
|
||||
key, val, found := strings.Cut(o, ":")
|
||||
if !found {
|
||||
xlog.Warn("dllm: ignoring malformed option (want key:value)", "option", o)
|
||||
continue
|
||||
}
|
||||
if !strings.HasPrefix(key, "eb_") && !dllmGenOptKeys[key] {
|
||||
xlog.Debug("dllm: ignoring unrecognized option", "key", key)
|
||||
continue
|
||||
}
|
||||
out[key] = parseScalarOpt(val)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func parseScalarOpt(v string) any {
|
||||
if iv, err := strconv.ParseInt(v, 10, 64); err == nil {
|
||||
return iv
|
||||
}
|
||||
if fv, err := strconv.ParseFloat(v, 64); err == nil {
|
||||
return fv
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
// metadataEnableThinking reads the enable_thinking gate. Unlike ds4 (default
|
||||
// ON, matching ds4-server), dllm defaults OFF: DiffusionGemma's chat
|
||||
// template guards every thinking branch with `enable_thinking is defined and
|
||||
// enable_thinking`, i.e. thinking is opt-in for this model family, and the
|
||||
// no-thinking render pre-closes an empty thought channel that the OFF
|
||||
// default must produce.
|
||||
func metadataEnableThinking(opts *pb.PredictOptions) bool {
|
||||
v := opts.GetMetadata()["enable_thinking"]
|
||||
return v == "true" || v == "1"
|
||||
}
|
||||
|
||||
// buildPrompt resolves the prompt for a request. With use_tokenizer_template
|
||||
// and raw messages the backend owns templating (RenderGemma4, including the
|
||||
// mmImageMarker injection for opts.Images) and the output is in the known
|
||||
// gemma4 format, so parse=true. Without it the caller templated the prompt
|
||||
// themselves (LocalAI's Go templates + PEG fallback, or a bare completion):
|
||||
// the prompt passes through verbatim - for image requests it must already
|
||||
// carry one literal mmImageMarker per image (the engine enforces the 1:1
|
||||
// match) - and the output is NOT gemma4-parsed - it is emitted as plain
|
||||
// content and the Go side's extraction applies, as for any non-autoparsing
|
||||
// backend.
|
||||
func buildPrompt(opts *pb.PredictOptions) (prompt string, parse bool, err error) {
|
||||
if opts.GetUseTokenizerTemplate() && len(opts.GetMessages()) > 0 {
|
||||
prompt, err = RenderGemma4(opts.GetMessages(), opts.GetTools(), len(opts.GetImages()), metadataEnableThinking(opts), true)
|
||||
return prompt, true, err
|
||||
}
|
||||
return opts.GetPrompt(), false, nil
|
||||
}
|
||||
|
||||
// imagesJSON renders opts.Images as the flat JSON array of data: URIs the mm
|
||||
// C-ABI expects, or "" when the request carries no images. The entries arrive
|
||||
// as RAW base64 payloads: LocalAI's OpenAI layer decodes every image_url /
|
||||
// image content part (URL download or data: URI) to plain base64 via
|
||||
// utils.GetContentURIAsBase64 (core/http/middleware/request.go) and core
|
||||
// flattens them into PredictOptions.Images (core/backend/llm.go). The
|
||||
// hardcoded image/jpeg mime mirrors the llama.cpp backend's re-wrapping
|
||||
// convention (grpc-server.cpp, "data:image/jpeg;base64," + images(i)); the
|
||||
// engine ignores the declared mime and sniffs the real format from the
|
||||
// decoded bytes (stb_image), so PNG/BMP payloads work through it too.
|
||||
func imagesJSON(images []string) (string, error) {
|
||||
if len(images) == 0 {
|
||||
return "", nil
|
||||
}
|
||||
uris := make([]string, len(images))
|
||||
for i, img := range images {
|
||||
// dllm_capi.h: array entries are read VERBATIM up to the closing
|
||||
// quote, with NO escape handling. json.Marshal would escape these
|
||||
// bytes and the C side would misparse the entry, so fail loud (they
|
||||
// can never appear in genuine base64 anyway).
|
||||
if strings.ContainsAny(img, "\"\\") {
|
||||
return "", fmt.Errorf("dllm: image %d is not base64 (contains a quote or backslash; PredictOptions.Images entries must be raw base64 payloads)", i)
|
||||
}
|
||||
uris[i] = "data:image/jpeg;base64," + img
|
||||
}
|
||||
b, err := json.Marshal(uris)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("dllm: marshal images: %w", err)
|
||||
}
|
||||
return string(b), nil
|
||||
}
|
||||
|
||||
// requestOptsJSON merges the model-level overrides with the request's
|
||||
// sampling fields into the flat opts JSON for one generate call.
|
||||
func (d *Dllm) requestOptsJSON(opts *pb.PredictOptions) (string, error) {
|
||||
m := make(map[string]any, len(d.genOpts)+2)
|
||||
for k, v := range d.genOpts {
|
||||
m[k] = v
|
||||
}
|
||||
if n := opts.GetTokens(); n > 0 {
|
||||
// The engine rounds n_predict UP to a whole number of diffusion
|
||||
// blocks (the canvas is denoised block-wise), so the completion may
|
||||
// run slightly past the requested budget. Tokens==0 omits the key so
|
||||
// the C-ABI default of 256 applies (hardcoded in capi.cpp's
|
||||
// parse_gen_opts, independent of canvas_length).
|
||||
m["n_predict"] = n
|
||||
}
|
||||
if s := opts.GetSeed(); s > 0 {
|
||||
// The engine seeds mt19937 with explicit non-negative seeds. Seed<=0
|
||||
// is omitted: proto3 cannot distinguish 0 from unset, and negative
|
||||
// values conventionally mean "random" across LocalAI backends.
|
||||
m["seed"] = s
|
||||
}
|
||||
return buildOptsJSON(m)
|
||||
}
|
||||
|
||||
// prepareRequest is the shared prologue of the rich methods: resolve the
|
||||
// prompt (and whether the output gets gemma4-parsed) and build the per-call
|
||||
// opts JSON plus the images JSON ("" for text-only requests, which routes
|
||||
// the call through the text generate entry points).
|
||||
func (d *Dllm) prepareRequest(opts *pb.PredictOptions) (prompt string, parse bool, optsJSON, imgJSON string, err error) {
|
||||
// Fail loud on media the engine has no path for, instead of silently
|
||||
// generating from a prompt that ignores them.
|
||||
if len(opts.GetVideos()) > 0 || len(opts.GetAudios()) > 0 {
|
||||
return "", false, "", "", errors.New("dllm: video/audio input is not supported (images only)")
|
||||
}
|
||||
prompt, parse, err = buildPrompt(opts)
|
||||
if err != nil {
|
||||
return "", false, "", "", err
|
||||
}
|
||||
optsJSON, err = d.requestOptsJSON(opts)
|
||||
if err != nil {
|
||||
return "", false, "", "", err
|
||||
}
|
||||
imgJSON, err = imagesJSON(opts.GetImages())
|
||||
if err != nil {
|
||||
return "", false, "", "", err
|
||||
}
|
||||
return prompt, parse, optsJSON, imgJSON, nil
|
||||
}
|
||||
|
||||
// sanitizeUTF8 makes s safe for a proto3 string field. Block-boundary
|
||||
// detokenization and byte-fallback tokens can produce invalid UTF-8, and
|
||||
// grpc-go refuses to marshal it ("string field contains invalid UTF-8"), so
|
||||
// every string destined for a Reply/ChatDelta must pass through here (or
|
||||
// through splitValidUTF8, which calls it). Lone malformed bytes are genuinely
|
||||
// undecodable: replace with U+FFFD rather than crash the stream.
|
||||
func sanitizeUTF8(s string) string {
|
||||
if utf8.ValidString(s) {
|
||||
return s
|
||||
}
|
||||
return strings.ToValidUTF8(s, "<22>")
|
||||
}
|
||||
|
||||
// utf8SeqLen returns the declared sequence length of a UTF-8 leading byte
|
||||
// (1 for bytes that can never lead a multi-byte sequence, so they are never
|
||||
// held back and fall through to sanitizeUTF8's replacement).
|
||||
func utf8SeqLen(b byte) int {
|
||||
switch {
|
||||
case b&0xE0 == 0xC0:
|
||||
return 2
|
||||
case b&0xF0 == 0xE0:
|
||||
return 3
|
||||
case b&0xF8 == 0xF0:
|
||||
return 4
|
||||
default:
|
||||
return 1
|
||||
}
|
||||
}
|
||||
|
||||
// splitValidUTF8 prepends the previous block's carry to the new block and
|
||||
// splits the result into text safe to emit now and a trailing INCOMPLETE
|
||||
// UTF-8 sequence (at most utf8.UTFMax-1 bytes) to carry into the next block:
|
||||
// the per-block detokenize can split a multi-byte character across block
|
||||
// boundaries (llama.cpp's grpc-server holds back the same way). Only a
|
||||
// suffix that can still become a valid rune is withheld; bytes that are
|
||||
// already undecodable are replaced immediately so the carry stays bounded.
|
||||
func splitValidUTF8(carry, block string) (emit, newCarry string) {
|
||||
s := carry + block
|
||||
cut := len(s)
|
||||
for i := len(s) - 1; i >= 0 && len(s)-i < utf8.UTFMax; i-- {
|
||||
b := s[i]
|
||||
if b < utf8.RuneSelf {
|
||||
break // ASCII: everything before the tail scan is complete
|
||||
}
|
||||
if !utf8.RuneStart(b) {
|
||||
continue // continuation byte: keep looking for its leading byte
|
||||
}
|
||||
// Leading byte: hold the sequence back iff it declares more bytes
|
||||
// than the stream has produced so far (it may complete next block).
|
||||
if utf8SeqLen(b) > len(s)-i {
|
||||
cut = i
|
||||
}
|
||||
break
|
||||
}
|
||||
return sanitizeUTF8(s[:cut]), s[cut:]
|
||||
}
|
||||
|
||||
// PredictRich is the non-streaming inference path (grpc.AIModelRich).
|
||||
// Returns one Reply whose Message is the aggregated assistant content and
|
||||
// whose ChatDeltas carry the parsed content/reasoning/tool-call events.
|
||||
func (d *Dllm) PredictRich(opts *pb.PredictOptions) (*pb.Reply, error) {
|
||||
d.genMu.RLock()
|
||||
defer d.genMu.RUnlock()
|
||||
if d.gen == nil {
|
||||
return nil, grpcerrors.ModelNotLoaded("dllm")
|
||||
}
|
||||
prompt, parse, optsJSON, imgJSON, err := d.prepareRequest(opts)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var out string
|
||||
var genErr error
|
||||
d.submit(func() {
|
||||
if imgJSON != "" {
|
||||
out, genErr = d.gen.generateMM(prompt, imgJSON, optsJSON)
|
||||
} else {
|
||||
out, genErr = d.gen.generate(prompt, optsJSON)
|
||||
}
|
||||
})
|
||||
if genErr != nil {
|
||||
return nil, genErr
|
||||
}
|
||||
// Byte-fallback tokens can detokenize to invalid UTF-8; proto3 strings
|
||||
// must be valid or grpc-go fails the whole reply at marshal time.
|
||||
out = sanitizeUTF8(out)
|
||||
|
||||
if !parse {
|
||||
// Raw-prompt mode: plain content, no gemma4 parsing (see buildPrompt).
|
||||
return &pb.Reply{Message: []byte(out), ChatDeltas: []*pb.ChatDelta{{Content: out}}}, nil
|
||||
}
|
||||
|
||||
// The prompt renders with add_generation_prompt; both thinking modes
|
||||
// leave the model starting in content state (see the Gemma4Parser header
|
||||
// comment), hence NewGemma4Parser(false).
|
||||
parser := NewGemma4Parser(false)
|
||||
if reply := replyFromDeltas(append(parser.Feed(out), parser.Close()...)); reply != nil {
|
||||
return reply, nil
|
||||
}
|
||||
// Everything was markers (or out was empty): an empty but non-nil Reply.
|
||||
return &pb.Reply{}, nil
|
||||
}
|
||||
|
||||
// PredictStreamRich is the streaming counterpart (grpc.AIModelRich): one
|
||||
// Reply per committed diffusion block that produced deltas. Per the
|
||||
// interface contract the channel is only sent into here - the gRPC server
|
||||
// closes it after this returns (opposite to legacy PredictStream).
|
||||
func (d *Dllm) PredictStreamRich(opts *pb.PredictOptions, results chan<- *pb.Reply) error {
|
||||
d.genMu.RLock()
|
||||
defer d.genMu.RUnlock()
|
||||
if d.gen == nil {
|
||||
return grpcerrors.ModelNotLoaded("dllm")
|
||||
}
|
||||
prompt, parse, optsJSON, imgJSON, err := d.prepareRequest(opts)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var parser *Gemma4Parser
|
||||
if parse {
|
||||
parser = NewGemma4Parser(false)
|
||||
}
|
||||
// emit runs inside onBlock, i.e. on the thread driving the C generate.
|
||||
// Sending on results can block on a slow consumer, but the server-side
|
||||
// pump (pkg/grpc/server.go PredictStream) drains continuously and drops
|
||||
// undeliverable sends, so this backpressure is brief and bounded - and
|
||||
// pausing the diffusion loop under it is the desired behavior anyway.
|
||||
emit := func(text string) {
|
||||
if !parse {
|
||||
if text != "" {
|
||||
results <- &pb.Reply{Message: []byte(text), ChatDeltas: []*pb.ChatDelta{{Content: text}}}
|
||||
}
|
||||
return
|
||||
}
|
||||
deltas := parser.Feed(text)
|
||||
if reply := replyFromDeltas(deltas); reply != nil {
|
||||
results <- reply
|
||||
}
|
||||
}
|
||||
// onBlock guards emit (and through it the parser) against invalid UTF-8:
|
||||
// a multi-byte character split across block boundaries is held back until
|
||||
// it completes (see splitValidUTF8), so proto3 marshaling never fails.
|
||||
var carry string
|
||||
onBlock := func(block string) {
|
||||
var text string
|
||||
text, carry = splitValidUTF8(carry, block)
|
||||
emit(text)
|
||||
}
|
||||
|
||||
var genErr error
|
||||
d.submit(func() {
|
||||
if imgJSON != "" {
|
||||
genErr = d.gen.generateStreamMM(prompt, imgJSON, optsJSON, onBlock)
|
||||
} else {
|
||||
genErr = d.gen.generateStream(prompt, optsJSON, onBlock)
|
||||
}
|
||||
})
|
||||
if genErr != nil {
|
||||
return genErr
|
||||
}
|
||||
if carry != "" {
|
||||
// The stream ended mid-sequence: the held-back bytes can no longer
|
||||
// complete, so flush them through the U+FFFD last resort.
|
||||
emit(sanitizeUTF8(carry))
|
||||
}
|
||||
if parse {
|
||||
if reply := replyFromDeltas(parser.Close()); reply != nil {
|
||||
results <- reply
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// replyFromDeltas wraps one batch of parsed deltas into a streaming Reply,
|
||||
// or nil when the batch is empty (markers consumed, nothing emitted yet).
|
||||
// Message mirrors the batch's content text so legacy chan-string consumers
|
||||
// see exactly the displayed tokens.
|
||||
func replyFromDeltas(deltas []*pb.ChatDelta) *pb.Reply {
|
||||
if len(deltas) == 0 {
|
||||
return nil
|
||||
}
|
||||
var content strings.Builder
|
||||
for _, delta := range deltas {
|
||||
content.WriteString(delta.GetContent())
|
||||
}
|
||||
return &pb.Reply{Message: []byte(content.String()), ChatDeltas: deltas}
|
||||
}
|
||||
|
||||
// Predict is the legacy (string, error) signature; the gRPC server prefers
|
||||
// PredictRich, this exists for non-rich callers (cloud-proxy precedent).
|
||||
func (d *Dllm) Predict(opts *pb.PredictOptions) (string, error) {
|
||||
reply, err := d.PredictRich(opts)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(reply.GetMessage()), nil
|
||||
}
|
||||
|
||||
// PredictStream is the legacy chan-string path: rich replies reduced to
|
||||
// their content text. Note the inverted channel ownership - the LEGACY
|
||||
// contract requires the impl to close the channel.
|
||||
func (d *Dllm) PredictStream(opts *pb.PredictOptions, results chan string) error {
|
||||
defer close(results)
|
||||
richCh := make(chan *pb.Reply)
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- d.PredictStreamRich(opts, richCh)
|
||||
close(richCh)
|
||||
}()
|
||||
for reply := range richCh {
|
||||
if msg := reply.GetMessage(); len(msg) > 0 {
|
||||
results <- string(msg)
|
||||
}
|
||||
}
|
||||
return <-errCh
|
||||
}
|
||||
|
||||
// TokenizeString tokenizes opts.Prompt via dllm_capi_tokenize_json (the C
|
||||
// side prepends bos per the vocab) and decodes the returned id array.
|
||||
func (d *Dllm) TokenizeString(opts *pb.PredictOptions) (pb.TokenizationResponse, error) {
|
||||
d.genMu.RLock()
|
||||
defer d.genMu.RUnlock()
|
||||
if d.gen == nil {
|
||||
return pb.TokenizationResponse{}, grpcerrors.ModelNotLoaded("dllm")
|
||||
}
|
||||
var out string
|
||||
var tokErr error
|
||||
d.submit(func() {
|
||||
out, tokErr = d.gen.tokenizeJSON(opts.GetPrompt())
|
||||
})
|
||||
if tokErr != nil {
|
||||
return pb.TokenizationResponse{}, tokErr
|
||||
}
|
||||
var tokens []int32
|
||||
if err := json.Unmarshal([]byte(out), &tokens); err != nil {
|
||||
return pb.TokenizationResponse{}, fmt.Errorf("dllm: decode tokenize result %q: %w", out, err)
|
||||
}
|
||||
return pb.TokenizationResponse{Length: int32(len(tokens)), Tokens: tokens}, nil
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,562 +0,0 @@
|
||||
// Gemma4 (DiffusionGemma) streaming output parser: raw model text, fed in
|
||||
// arbitrary fragments (per committed diffusion block; a fragment can split
|
||||
// anywhere, including mid-marker and mid-payload), is turned into
|
||||
// pb.ChatDelta events (content / reasoning_content / tool_calls).
|
||||
//
|
||||
// Normative sources:
|
||||
// - The chat template embedded at the top of gemma4_renderer.go ("tpl L<n>"
|
||||
// citations below refer to its numbered lines). The OUTPUT format mirrors
|
||||
// what the template renders for assistant history: thought channels
|
||||
// (<|channel>thought\n ... <channel|>, tpl L240), tool calls
|
||||
// (<|tool_call>call:name{...}<tool_call|>, tpl L246-L257) and turn ends
|
||||
// (<turn|>, tpl L351).
|
||||
// - vLLM PR #45163: vllm/tool_parsers/gemma4_tool_parser.py (marker
|
||||
// handling, the call:name{...} argument grammar and its decoder, ported
|
||||
// below) and vllm/reasoning/gemma4_reasoning_parser.py (channel markers,
|
||||
// the "thought\n" role label, is_reasoning_end semantics).
|
||||
//
|
||||
// Initial state (derived from the generation prompt, tpl L356-L362, see
|
||||
// RenderGemma4):
|
||||
// - enable_thinking=false: the prompt ends with "<|turn>model\n" +
|
||||
// "<|channel>thought\n<channel|>" - an EMPTY thought channel, pre-opened
|
||||
// AND pre-closed by the template. The model's output therefore starts in
|
||||
// plain content. Use NewGemma4Parser(false).
|
||||
// - enable_thinking=true: the prompt ends at "<|turn>model\n" and the model
|
||||
// opens and closes its own thought channel in the OUTPUT
|
||||
// ("<|channel>thought\n...reasoning...<channel|>final answer", per the
|
||||
// vLLM Gemma4ReasoningParser docstring). The parser still starts in
|
||||
// content state - the channel markers in the output drive the switch.
|
||||
// Use NewGemma4Parser(false) here too.
|
||||
// - NewGemma4Parser(true) is for callers that pre-open the thought channel
|
||||
// in the prompt themselves (appending "<|channel>thought\n" after the
|
||||
// generation prompt to force thinking): the output then begins mid-thought
|
||||
// and everything is reasoning until the first <channel|>.
|
||||
//
|
||||
// State diagram (markers are consumed, never emitted):
|
||||
//
|
||||
// <|channel> \n (channel name dropped: the
|
||||
// [content] --------------> [chan-header] ----> [thought] "thought\n" role
|
||||
// ^ | <channel|> (stray close: swallowed, label, stripped
|
||||
// +-+ strip_thinking semantics, tpl L148-L158) like vLLM does)
|
||||
// ^ <channel|>
|
||||
// +----------------------------------------- [thought]
|
||||
// ^ <tool_call|> | <|tool_call> (implicit
|
||||
// +-------------- [tool-call] <-------------------+ reasoning end, vLLM
|
||||
// | <|tool_call> ^ is_reasoning_end)
|
||||
// +-------------------+
|
||||
// [content]/[thought] --- <turn|> ---> [done] (everything after is dropped)
|
||||
//
|
||||
// Buffering rules:
|
||||
// - content/thought states hold back at most len(longest marker)-1 bytes:
|
||||
// the longest tail that is still a proper prefix of a watched marker.
|
||||
// Content is otherwise emitted immediately (no unbounded buffering).
|
||||
// - the tool-call state buffers the whole payload until <tool_call|>. This
|
||||
// is unbounded in principle but bounded in practice by the model's
|
||||
// diffusion canvas, and is required because the call:name{...} payload
|
||||
// only becomes decodable (and trustworthy) once complete - the same
|
||||
// reason vLLM's parser accumulates before parsing.
|
||||
// - Close() flushes whatever is still held: partial markers come out as
|
||||
// content/reasoning (per the state that held them); an unterminated
|
||||
// channel header or tool-call payload is re-emitted RAW (including its
|
||||
// opening marker) as content - malformed output is never silently
|
||||
// dropped (mirrors vLLM extract_tool_calls returning the raw text as
|
||||
// content when its regex does not match).
|
||||
//
|
||||
// Streaming granularity DIVERGENCE from vLLM: vLLM re-parses the partial
|
||||
// payload on every token and streams argument-JSON diffs (its `partial=True`
|
||||
// decoder mode plus withholding logic exist only for that). Our fragments are
|
||||
// whole committed diffusion blocks, so each completed tool call is emitted
|
||||
// once, as a single ToolCallDelta carrying index + id + name + the full
|
||||
// arguments JSON - exactly the shape backend/python/vllm/backend.py emits
|
||||
// per call and pkg/functions.ToolCallsFromChatDeltas re-accumulates.
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
)
|
||||
|
||||
// gemma4CallRE is vLLM's tool_call_regex
|
||||
// (`<\|tool_call>call:([\w\-\.]+)\{(.*?)\}<tool_call\|>`, DOTALL) anchored to
|
||||
// a single already-extracted payload: name charset [\w\-.], braces mandatory.
|
||||
var gemma4CallRE = regexp.MustCompile(`(?s)^call:([\w\-.]+)\{(.*)\}$`)
|
||||
|
||||
type g4State int
|
||||
|
||||
const (
|
||||
g4Content g4State = iota
|
||||
g4ChanHeader
|
||||
g4Thought
|
||||
g4ToolCall
|
||||
g4Done
|
||||
)
|
||||
|
||||
// Markers watched per emitting state. A stray <tool_call|> outside a tool
|
||||
// call is deliberately NOT watched: it passes through verbatim, consistent
|
||||
// with the malformed-payload fallback re-emitting it as content.
|
||||
var (
|
||||
gemma4ContentMarkers = []string{gemma4ChannelOpen, gemma4ChannelClose, gemma4ToolCallOpen, gemma4TurnEnd}
|
||||
gemma4ThoughtMarkers = []string{gemma4ChannelClose, gemma4ToolCallOpen, gemma4TurnEnd}
|
||||
)
|
||||
|
||||
type Gemma4Parser struct {
|
||||
state g4State
|
||||
// held is the per-state carry-over between Feed calls: a partial marker
|
||||
// (content/thought), a partial channel header (chan-header) or the
|
||||
// payload accumulated so far (tool-call).
|
||||
held string
|
||||
toolIdx int
|
||||
}
|
||||
|
||||
// NewGemma4Parser returns a parser positioned per the initial-state rules in
|
||||
// the header comment: startInThought=true only when the caller pre-opened a
|
||||
// thought channel in the prompt.
|
||||
func NewGemma4Parser(startInThought bool) *Gemma4Parser {
|
||||
state := g4Content
|
||||
if startInThought {
|
||||
state = g4Thought
|
||||
}
|
||||
return &Gemma4Parser{state: state}
|
||||
}
|
||||
|
||||
// Feed consumes the next output fragment and returns the deltas it completes.
|
||||
func (p *Gemma4Parser) Feed(text string) []*pb.ChatDelta {
|
||||
if text == "" || p.state == g4Done {
|
||||
return nil
|
||||
}
|
||||
pending := p.held + text
|
||||
p.held = ""
|
||||
var em g4Emitter
|
||||
for pending != "" {
|
||||
switch p.state {
|
||||
case g4Content, g4Thought:
|
||||
markers := gemma4ContentMarkers
|
||||
if p.state == g4Thought {
|
||||
markers = gemma4ThoughtMarkers
|
||||
}
|
||||
idx, marker := findEarliestGemma4Marker(pending, markers)
|
||||
if idx == -1 {
|
||||
hold := gemma4MarkerHoldback(pending, markers)
|
||||
p.emitText(&em, pending[:len(pending)-hold])
|
||||
p.held = pending[len(pending)-hold:]
|
||||
pending = ""
|
||||
continue
|
||||
}
|
||||
p.emitText(&em, pending[:idx])
|
||||
pending = pending[idx+len(marker):]
|
||||
switch marker {
|
||||
case gemma4ChannelOpen:
|
||||
p.state = g4ChanHeader
|
||||
case gemma4ChannelClose:
|
||||
// In thought: channel ends. In content: stray close,
|
||||
// swallowed (strip_thinking keeps both sides, tpl L148-L158).
|
||||
p.state = g4Content
|
||||
case gemma4ToolCallOpen:
|
||||
p.state = g4ToolCall
|
||||
case gemma4TurnEnd:
|
||||
p.state = g4Done
|
||||
}
|
||||
case g4ChanHeader:
|
||||
// The channel header is "<name>\n"; the template only ever writes
|
||||
// "thought" (tpl L240/L360) and the label is structural, so it is
|
||||
// dropped, not emitted (vLLM strips the same "thought\n" prefix).
|
||||
nl := strings.IndexByte(pending, '\n')
|
||||
if nl == -1 {
|
||||
p.held = pending
|
||||
pending = ""
|
||||
continue
|
||||
}
|
||||
pending = pending[nl+1:]
|
||||
p.state = g4Thought
|
||||
case g4ToolCall:
|
||||
end := strings.Index(pending, gemma4ToolCallClose)
|
||||
if end == -1 {
|
||||
p.held = pending
|
||||
pending = ""
|
||||
continue
|
||||
}
|
||||
p.emitToolCall(&em, pending[:end])
|
||||
pending = pending[end+len(gemma4ToolCallClose):]
|
||||
p.state = g4Content
|
||||
case g4Done:
|
||||
pending = ""
|
||||
}
|
||||
}
|
||||
return em.deltas
|
||||
}
|
||||
|
||||
// Close flushes held-back partials. Incomplete structures (open channel
|
||||
// header, unterminated tool payload) are re-emitted raw as content rather
|
||||
// than dropped. The parser is finished afterwards.
|
||||
func (p *Gemma4Parser) Close() []*pb.ChatDelta {
|
||||
var em g4Emitter
|
||||
switch p.state {
|
||||
case g4Content:
|
||||
em.content(p.held)
|
||||
case g4Thought:
|
||||
em.reasoning(p.held)
|
||||
case g4ChanHeader:
|
||||
em.content(gemma4ChannelOpen + p.held)
|
||||
case g4ToolCall:
|
||||
em.content(gemma4ToolCallOpen + p.held)
|
||||
case g4Done:
|
||||
}
|
||||
p.held = ""
|
||||
p.state = g4Done
|
||||
return em.deltas
|
||||
}
|
||||
|
||||
func (p *Gemma4Parser) emitText(em *g4Emitter, s string) {
|
||||
if p.state == g4Thought {
|
||||
em.reasoning(s)
|
||||
return
|
||||
}
|
||||
em.content(s)
|
||||
}
|
||||
|
||||
// emitToolCall decodes one complete <|tool_call>...<tool_call|> payload. On a
|
||||
// payload that does not match call:name{...} the raw text (markers included)
|
||||
// is emitted as content, mirroring vLLM's extract_tool_calls fallback.
|
||||
func (p *Gemma4Parser) emitToolCall(em *g4Emitter, payload string) {
|
||||
m := gemma4CallRE.FindStringSubmatch(payload)
|
||||
if m == nil {
|
||||
em.content(gemma4ToolCallOpen + payload + gemma4ToolCallClose)
|
||||
return
|
||||
}
|
||||
// Index-based ids: deterministic (the split-invariance property relies
|
||||
// on it) and matching the call_<n> convention of pkg/grpc/rich_test.go;
|
||||
// core only needs ids to be non-empty and unique within the response.
|
||||
em.tool(p.toolIdx, "call_"+strconv.Itoa(p.toolIdx), m[1], decodeGemma4Args(m[2], 0))
|
||||
p.toolIdx++
|
||||
}
|
||||
|
||||
// g4Emitter collects ChatDeltas; empty text events are dropped.
|
||||
type g4Emitter struct {
|
||||
deltas []*pb.ChatDelta
|
||||
}
|
||||
|
||||
func (e *g4Emitter) content(s string) {
|
||||
if s != "" {
|
||||
e.deltas = append(e.deltas, &pb.ChatDelta{Content: s})
|
||||
}
|
||||
}
|
||||
|
||||
func (e *g4Emitter) reasoning(s string) {
|
||||
if s != "" {
|
||||
e.deltas = append(e.deltas, &pb.ChatDelta{ReasoningContent: s})
|
||||
}
|
||||
}
|
||||
|
||||
func (e *g4Emitter) tool(index int, id, name, argsJSON string) {
|
||||
e.deltas = append(e.deltas, &pb.ChatDelta{ToolCalls: []*pb.ToolCallDelta{{
|
||||
Index: int32(index),
|
||||
Id: id,
|
||||
Name: name,
|
||||
Arguments: argsJSON,
|
||||
}}})
|
||||
}
|
||||
|
||||
// findEarliestGemma4Marker returns the position and value of the first
|
||||
// complete marker occurrence, or (-1, "").
|
||||
func findEarliestGemma4Marker(s string, markers []string) (int, string) {
|
||||
best, bestMarker := -1, ""
|
||||
for _, m := range markers {
|
||||
if idx := strings.Index(s, m); idx >= 0 && (best == -1 || idx < best) {
|
||||
best, bestMarker = idx, m
|
||||
}
|
||||
}
|
||||
return best, bestMarker
|
||||
}
|
||||
|
||||
// gemma4MarkerHoldback returns the length of the longest suffix of s that is
|
||||
// a proper prefix of a watched marker - the only bytes that may still grow
|
||||
// into a marker and therefore must not be emitted yet (bounded by the
|
||||
// longest marker, so content is never buffered unboundedly).
|
||||
func gemma4MarkerHoldback(s string, markers []string) int {
|
||||
maxHold := 0
|
||||
for _, m := range markers {
|
||||
if len(m)-1 > maxHold {
|
||||
maxHold = len(m) - 1
|
||||
}
|
||||
}
|
||||
if len(s) < maxHold {
|
||||
maxHold = len(s)
|
||||
}
|
||||
for k := maxHold; k >= 1; k-- {
|
||||
tail := s[len(s)-k:]
|
||||
for _, m := range markers {
|
||||
if strings.HasPrefix(m, tail) {
|
||||
return k
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// call:name{...} argument decoder
|
||||
//
|
||||
// Port of vLLM's _parse_gemma4_args / _parse_gemma4_array /
|
||||
// _parse_gemma4_value (gemma4_tool_parser.py) in non-partial mode only: this
|
||||
// parser decodes exclusively COMPLETE payloads (incomplete ones fall back to
|
||||
// raw content at Close), so vLLM's partial-withholding machinery
|
||||
// (trailing-dot floats, withheld bare tails) is intentionally not ported.
|
||||
//
|
||||
// Grammar (inverse of the renderer's formatGemma4Argument, tpl L118-L147):
|
||||
//
|
||||
// args := pair (',' pair)*
|
||||
// pair := key ':' value (keys unquoted, up to the first ':')
|
||||
// value := string | object | array | bare
|
||||
// string := '<|"|>' ... '<|"|>' (no escapes; unterminated -> rest)
|
||||
// object := '{' args '}' (delimited strings skipped when
|
||||
// array := '[' value,* ']' counting braces/brackets)
|
||||
// bare := true | false | null/none/nil | number | bare-string
|
||||
//
|
||||
// Output is a JSON object/array string with keys in payload order (Python
|
||||
// dict insertion order), built with HTML escaping off so payload text
|
||||
// survives byte-for-byte.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func isGemma4Space(c byte) bool { return c == ' ' || c == '\n' || c == '\t' }
|
||||
|
||||
// gemma4MaxArgsDepth caps the mutual recursion between decodeGemma4Args and
|
||||
// decodeGemma4Array. Defense against model-generated deep nesting: a Go stack
|
||||
// overflow is a fatal process kill, not a recoverable error, so past the cap
|
||||
// a nested body gracefully degrades to a JSON string of its raw text.
|
||||
const gemma4MaxArgsDepth = 100
|
||||
|
||||
// decodeGemma4Args decodes one args body (the text between the outer braces
|
||||
// of call:name{...}) into a JSON object string. depth is the current nesting
|
||||
// level (0 at the payload root); see gemma4MaxArgsDepth.
|
||||
func decodeGemma4Args(s string, depth int) string {
|
||||
if depth > gemma4MaxArgsDepth {
|
||||
return gemma4JSONString(s)
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString("{")
|
||||
first := true
|
||||
pair := func(key, val string) {
|
||||
if !first {
|
||||
b.WriteString(",")
|
||||
}
|
||||
first = false
|
||||
b.WriteString(gemma4JSONString(key))
|
||||
b.WriteString(":")
|
||||
b.WriteString(val)
|
||||
}
|
||||
i, n := 0, len(s)
|
||||
for i < n {
|
||||
for i < n && (isGemma4Space(s[i]) || s[i] == ',') {
|
||||
i++
|
||||
}
|
||||
if i >= n {
|
||||
break
|
||||
}
|
||||
keyStart := i
|
||||
for i < n && s[i] != ':' {
|
||||
i++
|
||||
}
|
||||
if i >= n {
|
||||
break // no ':' -> trailing junk, dropped (vLLM does the same)
|
||||
}
|
||||
key := strings.TrimSpace(s[keyStart:i])
|
||||
i++ // skip ':'
|
||||
for i < n && isGemma4Space(s[i]) {
|
||||
i++
|
||||
}
|
||||
if i >= n {
|
||||
pair(key, `""`) // "key:" with nothing after -> empty string
|
||||
break
|
||||
}
|
||||
switch {
|
||||
case strings.HasPrefix(s[i:], gemma4StringDelim):
|
||||
i += len(gemma4StringDelim)
|
||||
if end := strings.Index(s[i:], gemma4StringDelim); end == -1 {
|
||||
pair(key, gemma4JSONString(s[i:])) // unterminated -> take rest
|
||||
i = n
|
||||
} else {
|
||||
pair(key, gemma4JSONString(s[i:i+end]))
|
||||
i += end + len(gemma4StringDelim)
|
||||
}
|
||||
case s[i] == '{':
|
||||
inner, next := scanGemma4Balanced(s, i, '{', '}')
|
||||
pair(key, decodeGemma4Args(inner, depth+1))
|
||||
i = next
|
||||
case s[i] == '[':
|
||||
inner, next := scanGemma4Balanced(s, i, '[', ']')
|
||||
pair(key, decodeGemma4Array(inner, depth+1))
|
||||
i = next
|
||||
default:
|
||||
valStart := i
|
||||
for i < n && s[i] != ',' && s[i] != '}' && s[i] != ']' {
|
||||
i++
|
||||
}
|
||||
if i == valStart {
|
||||
// No progress (value starts on a stray '}'/']'): abort on
|
||||
// malformed input rather than loop, like vLLM.
|
||||
i = n
|
||||
continue
|
||||
}
|
||||
pair(key, decodeGemma4Bare(s[valStart:i]))
|
||||
}
|
||||
}
|
||||
b.WriteString("}")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// decodeGemma4Array decodes one array body (the text between '[' and ']')
|
||||
// into a JSON array string. depth is the current nesting level; see
|
||||
// gemma4MaxArgsDepth.
|
||||
func decodeGemma4Array(s string, depth int) string {
|
||||
if depth > gemma4MaxArgsDepth {
|
||||
return gemma4JSONString(s)
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString("[")
|
||||
first := true
|
||||
item := func(val string) {
|
||||
if !first {
|
||||
b.WriteString(",")
|
||||
}
|
||||
first = false
|
||||
b.WriteString(val)
|
||||
}
|
||||
i, n := 0, len(s)
|
||||
for i < n {
|
||||
for i < n && (isGemma4Space(s[i]) || s[i] == ',') {
|
||||
i++
|
||||
}
|
||||
if i >= n {
|
||||
break
|
||||
}
|
||||
switch {
|
||||
case strings.HasPrefix(s[i:], gemma4StringDelim):
|
||||
i += len(gemma4StringDelim)
|
||||
if end := strings.Index(s[i:], gemma4StringDelim); end == -1 {
|
||||
item(gemma4JSONString(s[i:]))
|
||||
i = n
|
||||
} else {
|
||||
item(gemma4JSONString(s[i : i+end]))
|
||||
i += end + len(gemma4StringDelim)
|
||||
}
|
||||
case s[i] == '{':
|
||||
inner, next := scanGemma4Balanced(s, i, '{', '}')
|
||||
item(decodeGemma4Args(inner, depth+1))
|
||||
i = next
|
||||
case s[i] == '[':
|
||||
inner, next := scanGemma4Balanced(s, i, '[', ']')
|
||||
item(decodeGemma4Array(inner, depth+1))
|
||||
i = next
|
||||
default:
|
||||
valStart := i
|
||||
for i < n && s[i] != ',' && s[i] != ']' {
|
||||
i++
|
||||
}
|
||||
if i == valStart {
|
||||
i = n // no progress: abort on malformed input, like vLLM
|
||||
continue
|
||||
}
|
||||
item(decodeGemma4Bare(s[valStart:i]))
|
||||
}
|
||||
}
|
||||
b.WriteString("]")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// scanGemma4Balanced scans a brace/bracket-balanced span starting at the
|
||||
// opener s[start], skipping over <|"|>-delimited strings so structural
|
||||
// characters inside them do not count (vLLM's depth scan). Returns the inner
|
||||
// text and the index just past the closer; an unterminated span yields the
|
||||
// rest of the string (the inner decoder still extracts what is there - this
|
||||
// path is only reachable from genuinely malformed complete payloads).
|
||||
func scanGemma4Balanced(s string, start int, open, close byte) (string, int) {
|
||||
depth := 1
|
||||
i := start + 1
|
||||
innerStart := i
|
||||
n := len(s)
|
||||
for i < n && depth > 0 {
|
||||
if strings.HasPrefix(s[i:], gemma4StringDelim) {
|
||||
i += len(gemma4StringDelim)
|
||||
if nd := strings.Index(s[i:], gemma4StringDelim); nd == -1 {
|
||||
i = n
|
||||
} else {
|
||||
i += nd + len(gemma4StringDelim)
|
||||
}
|
||||
continue
|
||||
}
|
||||
switch s[i] {
|
||||
case open:
|
||||
depth++
|
||||
case close:
|
||||
depth--
|
||||
}
|
||||
i++
|
||||
}
|
||||
if depth > 0 {
|
||||
return s[innerStart:], n
|
||||
}
|
||||
return s[innerStart : i-1], i
|
||||
}
|
||||
|
||||
// decodeGemma4Bare maps an undelimited value to its JSON form: booleans,
|
||||
// null aliases (null/none/nil, case-insensitive - the renderer writes
|
||||
// Python None as "None", tpl L144-L145 via format_argument's else branch),
|
||||
// numbers (vLLM's rule: a '.' tries float, otherwise int; anything that
|
||||
// fails parses as a bare string).
|
||||
func decodeGemma4Bare(raw string) string {
|
||||
v := strings.TrimSpace(raw)
|
||||
if v == "" {
|
||||
return `""`
|
||||
}
|
||||
if v == "true" || v == "false" {
|
||||
return v
|
||||
}
|
||||
switch strings.ToLower(v) {
|
||||
case "null", "none", "nil":
|
||||
return "null"
|
||||
}
|
||||
if strings.Contains(v, ".") {
|
||||
if f, err := strconv.ParseFloat(v, 64); err == nil {
|
||||
return formatGemma4Float(f)
|
||||
}
|
||||
} else if iv, err := strconv.ParseInt(v, 10, 64); err == nil {
|
||||
return strconv.FormatInt(iv, 10)
|
||||
}
|
||||
return gemma4JSONString(v)
|
||||
}
|
||||
|
||||
// formatGemma4Float renders like Python's json.dumps(float): integral floats
|
||||
// keep a ".0" suffix ("108." decodes to 108.0, not 108), so the arguments
|
||||
// JSON matches what vLLM would have produced for the same payload.
|
||||
func formatGemma4Float(f float64) string {
|
||||
s := strconv.FormatFloat(f, 'g', -1, 64)
|
||||
if !strings.ContainsAny(s, ".eE") {
|
||||
s += ".0"
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// gemma4JSONString encodes a JSON string WITHOUT HTML escaping (json.Marshal
|
||||
// would escape the angle brackets in "<div>" to \u003c / \u003e sequences;
|
||||
// payload text should survive
|
||||
// byte-for-byte, like Python's json.dumps(ensure_ascii=False)).
|
||||
func gemma4JSONString(s string) string {
|
||||
var sb strings.Builder
|
||||
enc := json.NewEncoder(&sb)
|
||||
enc.SetEscapeHTML(false)
|
||||
if err := enc.Encode(s); err != nil {
|
||||
// Unreachable for plain strings; fall back to default escaping
|
||||
// rather than emitting invalid JSON.
|
||||
b, mErr := json.Marshal(s)
|
||||
if mErr != nil {
|
||||
return `""`
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
// Encode appends a trailing newline.
|
||||
return strings.TrimSuffix(sb.String(), "\n")
|
||||
}
|
||||
@@ -1,592 +0,0 @@
|
||||
package main
|
||||
|
||||
// Parser specs for Gemma4Parser (model output text -> pb.ChatDelta events).
|
||||
//
|
||||
// Fixture provenance:
|
||||
// - Entries marked "vLLM: <name>" are direct ports of the named test from
|
||||
// vLLM PR #45163, tests/tool_parsers/test_gemma4_tool_parser.py (the
|
||||
// authoritative test-suite for the gemma4 tool-call wire format). The
|
||||
// streaming tests' chunk lists are reused verbatim as Feed fragments.
|
||||
// - Decoder entries port the TestParseGemma4Args / TestParseGemma4Array
|
||||
// classes from the same file (non-partial mode only; this parser never
|
||||
// decodes partial payloads, see the divergence note in gemma4_parser.go).
|
||||
// - Channel/turn-marker expectations come from the chat template embedded
|
||||
// in gemma4_renderer.go (tpl L356-L362 generation prompt, L148-L158
|
||||
// strip_thinking) and vLLM's Gemma4ReasoningParser
|
||||
// (vllm/reasoning/gemma4_reasoning_parser.py).
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
)
|
||||
|
||||
// flatGemma4Tool is one accumulated tool call, mirroring how LocalAI core
|
||||
// folds ToolCallDelta streams (pkg/functions/chat_deltas.go
|
||||
// ToolCallsFromChatDeltas: name/id latch on first non-empty, arguments
|
||||
// concatenate per index). Tests flatten through the same rules so they
|
||||
// assert exactly what core will reconstruct.
|
||||
type flatGemma4Tool struct {
|
||||
id string
|
||||
name string
|
||||
args string
|
||||
}
|
||||
|
||||
func flattenGemma4Deltas(deltas []*pb.ChatDelta) (string, string, []flatGemma4Tool) {
|
||||
var content, reasoning strings.Builder
|
||||
byIndex := map[int32]*flatGemma4Tool{}
|
||||
maxIdx := int32(-1)
|
||||
for _, d := range deltas {
|
||||
content.WriteString(d.GetContent())
|
||||
reasoning.WriteString(d.GetReasoningContent())
|
||||
for _, tc := range d.GetToolCalls() {
|
||||
acc, ok := byIndex[tc.GetIndex()]
|
||||
if !ok {
|
||||
acc = &flatGemma4Tool{}
|
||||
byIndex[tc.GetIndex()] = acc
|
||||
}
|
||||
if tc.GetName() != "" {
|
||||
acc.name = tc.GetName()
|
||||
}
|
||||
if tc.GetId() != "" {
|
||||
acc.id = tc.GetId()
|
||||
}
|
||||
acc.args += tc.GetArguments()
|
||||
if tc.GetIndex() > maxIdx {
|
||||
maxIdx = tc.GetIndex()
|
||||
}
|
||||
}
|
||||
}
|
||||
var tools []flatGemma4Tool
|
||||
for i := int32(0); i <= maxIdx; i++ {
|
||||
if acc, ok := byIndex[i]; ok {
|
||||
tools = append(tools, *acc)
|
||||
}
|
||||
}
|
||||
return content.String(), reasoning.String(), tools
|
||||
}
|
||||
|
||||
type wantGemma4Tool struct {
|
||||
name string
|
||||
argsJSON string // compared with MatchJSON (key order irrelevant)
|
||||
}
|
||||
|
||||
type parseGemma4Case struct {
|
||||
startInThought bool
|
||||
fragments []string
|
||||
wantContent string
|
||||
wantReasoning string
|
||||
wantTools []wantGemma4Tool
|
||||
}
|
||||
|
||||
func parseGemma4Fragments(startInThought bool, fragments []string) []*pb.ChatDelta {
|
||||
p := NewGemma4Parser(startInThought)
|
||||
var all []*pb.ChatDelta
|
||||
for _, f := range fragments {
|
||||
all = append(all, p.Feed(f)...)
|
||||
}
|
||||
return append(all, p.Close()...)
|
||||
}
|
||||
|
||||
var _ = Describe("Gemma4Parser", func() {
|
||||
DescribeTable("parses streamed gemma4 output into ChatDeltas",
|
||||
func(c parseGemma4Case) {
|
||||
content, reasoning, tools := flattenGemma4Deltas(parseGemma4Fragments(c.startInThought, c.fragments))
|
||||
Expect(content).To(Equal(c.wantContent))
|
||||
Expect(reasoning).To(Equal(c.wantReasoning))
|
||||
Expect(tools).To(HaveLen(len(c.wantTools)))
|
||||
seenIDs := map[string]bool{}
|
||||
for i, want := range c.wantTools {
|
||||
Expect(tools[i].name).To(Equal(want.name), "tool %d name", i)
|
||||
Expect(tools[i].args).To(MatchJSON(want.argsJSON), "tool %d arguments", i)
|
||||
Expect(tools[i].id).ToNot(BeEmpty(), "tool %d id", i)
|
||||
Expect(seenIDs).ToNot(HaveKey(tools[i].id), "tool %d id must be unique", i)
|
||||
seenIDs[tools[i].id] = true
|
||||
}
|
||||
},
|
||||
|
||||
// --- (1) pure content -------------------------------------------------
|
||||
// vLLM: test_no_tool_calls
|
||||
Entry("pure content, single fragment", parseGemma4Case{
|
||||
fragments: []string{"Hello, how can I help you today?"},
|
||||
wantContent: "Hello, how can I help you today?",
|
||||
}),
|
||||
|
||||
// --- (2) thought -> final transition ----------------------------------
|
||||
// enable_thinking render: prompt ends at <|turn>model\n and the model
|
||||
// opens/closes its own thought channel in the OUTPUT (vLLM
|
||||
// Gemma4ReasoningParser docstring; tpl L356-L362). The "thought\n"
|
||||
// role label after <|channel> is structural and must be stripped
|
||||
// (vLLM _THOUGHT_PREFIX handling).
|
||||
Entry("thought channel then final content", parseGemma4Case{
|
||||
fragments: []string{"<|channel>thought\nLet me think about this.\n<channel|>The answer is 42."},
|
||||
wantReasoning: "Let me think about this.\n",
|
||||
wantContent: "The answer is 42.",
|
||||
}),
|
||||
|
||||
// --- (3) startInThought both ways -------------------------------------
|
||||
Entry("startInThought=true routes initial text to reasoning until <channel|>", parseGemma4Case{
|
||||
startInThought: true,
|
||||
fragments: []string{"I am thinking hard.<channel|>Done."},
|
||||
wantReasoning: "I am thinking hard.",
|
||||
wantContent: "Done.",
|
||||
}),
|
||||
// A stray <channel|> with no open channel is swallowed, matching the
|
||||
// template's strip_thinking (tpl L148-L158: the marker is dropped,
|
||||
// text on both sides is kept).
|
||||
Entry("startInThought=false keeps the same text as content, stray <channel|> swallowed", parseGemma4Case{
|
||||
startInThought: false,
|
||||
fragments: []string{"I am thinking hard.<channel|>Done."},
|
||||
wantContent: "I am thinking hard.Done.",
|
||||
}),
|
||||
|
||||
// --- (4) one tool call, full payload type zoo --------------------------
|
||||
Entry("single tool call: strings, numbers, bools, null, nested object and array", parseGemma4Case{
|
||||
fragments: []string{`<|tool_call>call:complex_function{text:<|"|>with, comma and {braces}<|"|>,count:42,score:3.14,yes:true,no:false,nothing:null,obj:{inner:<|"|>v<|"|>,k:1},arr:[<|"|>a<|"|>,2,true]}<tool_call|>`},
|
||||
wantTools: []wantGemma4Tool{{
|
||||
name: "complex_function",
|
||||
argsJSON: `{"text":"with, comma and {braces}","count":42,"score":3.14,"yes":true,"no":false,"nothing":null,"obj":{"inner":"v","k":1},"arr":["a",2,true]}`,
|
||||
}},
|
||||
}),
|
||||
|
||||
// --- (5) payload split across 3 fragments ------------------------------
|
||||
Entry("tool-call payload split across three fragments", parseGemma4Case{
|
||||
fragments: []string{
|
||||
"<|tool_call>call:get_weather{loc",
|
||||
`ation:<|"|>Paris, Fra`,
|
||||
`nce<|"|>}<tool_call|>`,
|
||||
},
|
||||
wantTools: []wantGemma4Tool{{name: "get_weather", argsJSON: `{"location":"Paris, France"}`}},
|
||||
}),
|
||||
|
||||
// --- (6) marker split across fragments ----------------------------------
|
||||
Entry("tool-call open marker split across fragments", parseGemma4Case{
|
||||
fragments: []string{
|
||||
"<|tool_ca",
|
||||
`ll>call:get_weather{location:<|"|>London<|"|>}<tool_call|>`,
|
||||
},
|
||||
wantTools: []wantGemma4Tool{{name: "get_weather", argsJSON: `{"location":"London"}`}},
|
||||
}),
|
||||
Entry("channel open marker split across fragments", parseGemma4Case{
|
||||
fragments: []string{
|
||||
"<|chan",
|
||||
"nel>thought\ndeep thought<channel|>final",
|
||||
},
|
||||
wantReasoning: "deep thought",
|
||||
wantContent: "final",
|
||||
}),
|
||||
|
||||
// --- (7) trailing partial marker held, flushed by Close -----------------
|
||||
Entry("trailing partial marker is held back and flushed by Close", parseGemma4Case{
|
||||
fragments: []string{"Hello <|tool"},
|
||||
wantContent: "Hello <|tool",
|
||||
}),
|
||||
|
||||
// --- (8) malformed/incomplete payload -> content fallback ---------------
|
||||
// vLLM: test_incomplete_tool_call (no end marker: the whole text stays
|
||||
// content, never silently dropped).
|
||||
Entry("incomplete tool payload at Close is emitted as raw content", parseGemma4Case{
|
||||
fragments: []string{`<|tool_call>call:get_weather{location:<|"|>London`},
|
||||
wantContent: `<|tool_call>call:get_weather{location:<|"|>London`,
|
||||
}),
|
||||
Entry("malformed complete payload is emitted as raw content, parsing continues", parseGemma4Case{
|
||||
fragments: []string{"<|tool_call>oops no call syntax<tool_call|> done"},
|
||||
wantContent: "<|tool_call>oops no call syntax<tool_call|> done",
|
||||
}),
|
||||
|
||||
// --- (9) <turn|> ends the turn -------------------------------------------
|
||||
Entry("text after <turn|> is ignored, including later fragments", parseGemma4Case{
|
||||
fragments: []string{
|
||||
"before<turn|>after",
|
||||
`more <|tool_call>call:f{}<tool_call|>`,
|
||||
},
|
||||
wantContent: "before",
|
||||
}),
|
||||
Entry("<turn|> inside a thought channel ends the turn", parseGemma4Case{
|
||||
startInThought: true,
|
||||
fragments: []string{"thinking<turn|>ignored"},
|
||||
wantReasoning: "thinking",
|
||||
}),
|
||||
|
||||
// --- (10) ported vLLM non-streaming cases ---------------------------------
|
||||
// vLLM: test_single_tool_call
|
||||
Entry("vLLM: test_single_tool_call", parseGemma4Case{
|
||||
fragments: []string{`<|tool_call>call:get_weather{location:<|"|>London<|"|>}<tool_call|>`},
|
||||
wantTools: []wantGemma4Tool{{name: "get_weather", argsJSON: `{"location":"London"}`}},
|
||||
}),
|
||||
// vLLM: test_multiple_arguments
|
||||
Entry("vLLM: test_multiple_arguments", parseGemma4Case{
|
||||
fragments: []string{`<|tool_call>call:get_weather{location:<|"|>San Francisco<|"|>,unit:<|"|>celsius<|"|>}<tool_call|>`},
|
||||
wantTools: []wantGemma4Tool{{name: "get_weather", argsJSON: `{"location":"San Francisco","unit":"celsius"}`}},
|
||||
}),
|
||||
// vLLM: test_text_before_tool_call. DIVERGENCE: vLLM's non-streaming
|
||||
// extractor trims the content ("...you."); a streaming parser cannot
|
||||
// retroactively trim already-emitted text, so the trailing space is
|
||||
// kept (vLLM's own streaming path keeps it too, see
|
||||
// test_streaming_text_before_tool_call which only checks a prefix).
|
||||
Entry("vLLM: test_text_before_tool_call (streaming semantics: no trim)", parseGemma4Case{
|
||||
fragments: []string{`Let me check the weather for you. <|tool_call>call:get_weather{location:<|"|>Paris<|"|>}<tool_call|>`},
|
||||
wantContent: "Let me check the weather for you. ",
|
||||
wantTools: []wantGemma4Tool{{name: "get_weather", argsJSON: `{"location":"Paris"}`}},
|
||||
}),
|
||||
// vLLM: test_multiple_tool_calls (also covers case 11: multi-tool sequence)
|
||||
Entry("vLLM: test_multiple_tool_calls", parseGemma4Case{
|
||||
fragments: []string{`<|tool_call>call:get_weather{location:<|"|>London<|"|>}<tool_call|><|tool_call>call:get_time{location:<|"|>London<|"|>}<tool_call|>`},
|
||||
wantTools: []wantGemma4Tool{
|
||||
{name: "get_weather", argsJSON: `{"location":"London"}`},
|
||||
{name: "get_time", argsJSON: `{"location":"London"}`},
|
||||
},
|
||||
}),
|
||||
// vLLM: test_nested_arguments
|
||||
Entry("vLLM: test_nested_arguments", parseGemma4Case{
|
||||
fragments: []string{`<|tool_call>call:complex_function{nested:{inner:<|"|>value<|"|>},list:[<|"|>a<|"|>,<|"|>b<|"|>]}<tool_call|>`},
|
||||
wantTools: []wantGemma4Tool{{name: "complex_function", argsJSON: `{"nested":{"inner":"value"},"list":["a","b"]}`}},
|
||||
}),
|
||||
// vLLM: test_tool_call_with_number_and_boolean
|
||||
Entry("vLLM: test_tool_call_with_number_and_boolean", parseGemma4Case{
|
||||
fragments: []string{`<|tool_call>call:set_status{is_active:true,count:42,score:3.14}<tool_call|>`},
|
||||
wantTools: []wantGemma4Tool{{name: "set_status", argsJSON: `{"is_active":true,"count":42,"score":3.14}`}},
|
||||
}),
|
||||
// vLLM: test_hyphenated_function_name
|
||||
Entry("vLLM: test_hyphenated_function_name", parseGemma4Case{
|
||||
fragments: []string{`<|tool_call>call:get-weather{location:<|"|>London<|"|>}<tool_call|>`},
|
||||
wantTools: []wantGemma4Tool{{name: "get-weather", argsJSON: `{"location":"London"}`}},
|
||||
}),
|
||||
// vLLM: test_dotted_function_name
|
||||
Entry("vLLM: test_dotted_function_name", parseGemma4Case{
|
||||
fragments: []string{`<|tool_call>call:weather.get{location:<|"|>London<|"|>}<tool_call|>`},
|
||||
wantTools: []wantGemma4Tool{{name: "weather.get", argsJSON: `{"location":"London"}`}},
|
||||
}),
|
||||
// vLLM: test_no_arguments
|
||||
Entry("vLLM: test_no_arguments", parseGemma4Case{
|
||||
fragments: []string{"<|tool_call>call:get_status{}<tool_call|>"},
|
||||
wantTools: []wantGemma4Tool{{name: "get_status", argsJSON: `{}`}},
|
||||
}),
|
||||
|
||||
// --- ported vLLM streaming cases (chunk lists reused as fragments) --------
|
||||
// vLLM: test_basic_streaming_single_tool
|
||||
Entry("vLLM: test_basic_streaming_single_tool", parseGemma4Case{
|
||||
fragments: []string{
|
||||
"<|tool_call>",
|
||||
"call:get_weather{",
|
||||
`location:<|"|>Paris`,
|
||||
", France",
|
||||
`<|"|>}`,
|
||||
"<tool_call|>",
|
||||
},
|
||||
wantTools: []wantGemma4Tool{{name: "get_weather", argsJSON: `{"location":"Paris, France"}`}},
|
||||
}),
|
||||
// vLLM: test_streaming_multi_arg
|
||||
Entry("vLLM: test_streaming_multi_arg", parseGemma4Case{
|
||||
fragments: []string{
|
||||
"<|tool_call>",
|
||||
"call:get_weather{",
|
||||
`location:<|"|>Tokyo<|"|>,`,
|
||||
`unit:<|"|>celsius<|"|>}`,
|
||||
"<tool_call|>",
|
||||
},
|
||||
wantTools: []wantGemma4Tool{{name: "get_weather", argsJSON: `{"location":"Tokyo","unit":"celsius"}`}},
|
||||
}),
|
||||
// vLLM: test_streaming_text_before_tool_call
|
||||
Entry("vLLM: test_streaming_text_before_tool_call", parseGemma4Case{
|
||||
fragments: []string{
|
||||
"Let me check ",
|
||||
"the weather. ",
|
||||
"<|tool_call>",
|
||||
"call:get_weather{",
|
||||
`location:<|"|>London<|"|>}`,
|
||||
"<tool_call|>",
|
||||
},
|
||||
wantContent: "Let me check the weather. ",
|
||||
wantTools: []wantGemma4Tool{{name: "get_weather", argsJSON: `{"location":"London"}`}},
|
||||
}),
|
||||
// vLLM: test_streaming_numeric_args
|
||||
Entry("vLLM: test_streaming_numeric_args", parseGemma4Case{
|
||||
fragments: []string{
|
||||
"<|tool_call>",
|
||||
"call:set_config{",
|
||||
"count:42,",
|
||||
"active:true}",
|
||||
"<tool_call|>",
|
||||
},
|
||||
wantTools: []wantGemma4Tool{{name: "set_config", argsJSON: `{"count":42,"active":true}`}},
|
||||
}),
|
||||
// vLLM: test_streaming_boolean_split_across_chunks
|
||||
Entry("vLLM: test_streaming_boolean_split_across_chunks", parseGemma4Case{
|
||||
fragments: []string{
|
||||
"<|tool_call>",
|
||||
"call:search{input:{all:tru",
|
||||
"e}}",
|
||||
"<tool_call|>",
|
||||
},
|
||||
wantTools: []wantGemma4Tool{{name: "search", argsJSON: `{"input":{"all":true}}`}},
|
||||
}),
|
||||
// vLLM: test_streaming_false_split_across_chunks
|
||||
Entry("vLLM: test_streaming_false_split_across_chunks", parseGemma4Case{
|
||||
fragments: []string{
|
||||
"<|tool_call>",
|
||||
"call:set{flag:fals",
|
||||
"e}",
|
||||
"<tool_call|>",
|
||||
},
|
||||
wantTools: []wantGemma4Tool{{name: "set", argsJSON: `{"flag":false}`}},
|
||||
}),
|
||||
// vLLM: test_streaming_number_split_across_chunks
|
||||
Entry("vLLM: test_streaming_number_split_across_chunks", parseGemma4Case{
|
||||
fragments: []string{
|
||||
"<|tool_call>",
|
||||
"call:set{count:4",
|
||||
"2}",
|
||||
"<tool_call|>",
|
||||
},
|
||||
wantTools: []wantGemma4Tool{{name: "set", argsJSON: `{"count":42}`}},
|
||||
}),
|
||||
// vLLM: test_streaming_empty_args
|
||||
Entry("vLLM: test_streaming_empty_args", parseGemma4Case{
|
||||
fragments: []string{
|
||||
"<|tool_call>",
|
||||
"call:get_status{}",
|
||||
"<tool_call|>",
|
||||
},
|
||||
wantTools: []wantGemma4Tool{{name: "get_status", argsJSON: `{}`}},
|
||||
}),
|
||||
// vLLM: test_streaming_split_delimiter_no_invalid_json (string
|
||||
// delimiter <|"|> split across fragments must not leak fragments).
|
||||
Entry("vLLM: test_streaming_split_delimiter_no_invalid_json", parseGemma4Case{
|
||||
fragments: []string{
|
||||
"<|tool_call>",
|
||||
"call:todowrite{",
|
||||
`content:<|"|>Buy milk<|`,
|
||||
`"|>}`,
|
||||
"<tool_call|>",
|
||||
},
|
||||
wantTools: []wantGemma4Tool{{name: "todowrite", argsJSON: `{"content":"Buy milk"}`}},
|
||||
}),
|
||||
// vLLM: test_streaming_does_not_duplicate_plain_text_after_tool_call
|
||||
Entry("vLLM: test_streaming_does_not_duplicate_plain_text_after_tool_call", parseGemma4Case{
|
||||
fragments: []string{
|
||||
"<|tool_call>",
|
||||
"call:get_weather{",
|
||||
`location:<|"|>Paris<|"|>}`,
|
||||
"<tool_call|><",
|
||||
"div>",
|
||||
},
|
||||
wantContent: "<div>",
|
||||
wantTools: []wantGemma4Tool{{name: "get_weather", argsJSON: `{"location":"Paris"}`}},
|
||||
}),
|
||||
// vLLM: test_streaming_html_argument_does_not_duplicate_tag_prefixes
|
||||
Entry("vLLM: test_streaming_html_argument_does_not_duplicate_tag_prefixes", parseGemma4Case{
|
||||
fragments: []string{
|
||||
"<|tool_call>",
|
||||
"call:write_file{",
|
||||
`path:<|"|>index.html<|"|>,`,
|
||||
`content:<|"|><!DOCTYPE html>` + "\n<",
|
||||
`html lang="zh-CN">` + "\n<",
|
||||
"head>\n <",
|
||||
`meta charset="UTF-8">` + "\n <",
|
||||
`meta name="viewport" content="width=device-width">` + "\n",
|
||||
`<|"|>}`,
|
||||
"<tool_call|>",
|
||||
},
|
||||
wantTools: []wantGemma4Tool{{
|
||||
name: "write_file",
|
||||
argsJSON: `{"path":"index.html","content":"<!DOCTYPE html>\n<html lang=\"zh-CN\">\n<head>\n <meta charset=\"UTF-8\">\n <meta name=\"viewport\" content=\"width=device-width\">\n"}`,
|
||||
}},
|
||||
}),
|
||||
// vLLM: test_streaming_single_chunk_complete_tool_call
|
||||
Entry("vLLM: test_streaming_single_chunk_complete_tool_call", parseGemma4Case{
|
||||
fragments: []string{`<|tool_call>call:name_a_color{color_hex:<|"|>00ff11<|"|>}<tool_call|>`},
|
||||
wantTools: []wantGemma4Tool{{name: "name_a_color", argsJSON: `{"color_hex":"00ff11"}`}},
|
||||
}),
|
||||
// vLLM: test_streaming_multi_chunk_batched_tool_calls (two complete
|
||||
// calls in ONE fragment; both must come out with distinct indices)
|
||||
Entry("vLLM: test_streaming_multi_chunk_batched_tool_calls", parseGemma4Case{
|
||||
fragments: []string{
|
||||
`<|tool_call>call:get_weather{location:<|"|>London<|"|>}<tool_call|>` +
|
||||
`<|tool_call>call:get_time{timezone:<|"|>GMT<|"|>}<tool_call|>`,
|
||||
},
|
||||
wantTools: []wantGemma4Tool{
|
||||
{name: "get_weather", argsJSON: `{"location":"London"}`},
|
||||
{name: "get_time", argsJSON: `{"timezone":"GMT"}`},
|
||||
},
|
||||
}),
|
||||
// vLLM: test_streaming_trailing_bare_bool_not_duplicated
|
||||
Entry("vLLM: test_streaming_trailing_bare_bool_not_duplicated", parseGemma4Case{
|
||||
fragments: []string{
|
||||
"<|tool_call>",
|
||||
"call:Edit{",
|
||||
`file_path:<|"|>src/env.py<|"|>,`,
|
||||
`old_string:<|"|>old_val<|"|>,`,
|
||||
`new_string:<|"|>new_val<|"|>,`,
|
||||
"replace_all:",
|
||||
"false}",
|
||||
"<tool_call|>",
|
||||
},
|
||||
wantTools: []wantGemma4Tool{{
|
||||
name: "Edit",
|
||||
argsJSON: `{"file_path":"src/env.py","old_string":"old_val","new_string":"new_val","replace_all":false}`,
|
||||
}},
|
||||
}),
|
||||
|
||||
// --- implicit reasoning end on <|tool_call> (vLLM is_reasoning_end:
|
||||
// a tool_call token means reasoning is over) -----------------------------
|
||||
Entry("tool call inside an open thought channel ends the reasoning", parseGemma4Case{
|
||||
startInThought: true,
|
||||
fragments: []string{`need the weather<|tool_call>call:get_weather{location:<|"|>Rome<|"|>}<tool_call|>`},
|
||||
wantReasoning: "need the weather",
|
||||
wantTools: []wantGemma4Tool{{name: "get_weather", argsJSON: `{"location":"Rome"}`}},
|
||||
}),
|
||||
|
||||
// --- (12) empty fragments are no-ops --------------------------------------
|
||||
Entry("empty fragments are no-ops", parseGemma4Case{
|
||||
fragments: []string{"", "Hello", "", "", " world", ""},
|
||||
wantContent: "Hello world",
|
||||
}),
|
||||
)
|
||||
|
||||
It("returns no deltas for an empty fragment and after Close", func() {
|
||||
p := NewGemma4Parser(false)
|
||||
Expect(p.Feed("")).To(BeEmpty())
|
||||
Expect(p.Feed("hi")).ToNot(BeEmpty())
|
||||
Expect(p.Close()).To(BeEmpty()) // nothing held back
|
||||
// The parser is finished after Close: further input is dropped.
|
||||
Expect(p.Feed("more")).To(BeEmpty())
|
||||
Expect(p.Close()).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("generates index-based tool call ids (call_<index>)", func() {
|
||||
// Mirrors the index-based id convention of pkg/grpc/rich_test.go and
|
||||
// keeps ids deterministic for the split-invariance property below.
|
||||
deltas := parseGemma4Fragments(false, []string{
|
||||
`<|tool_call>call:a{}<tool_call|><|tool_call>call:b{}<tool_call|>`,
|
||||
})
|
||||
_, _, tools := flattenGemma4Deltas(deltas)
|
||||
Expect(tools).To(HaveLen(2))
|
||||
Expect(tools[0].id).To(Equal("call_0"))
|
||||
Expect(tools[1].id).To(Equal("call_1"))
|
||||
})
|
||||
|
||||
// Property: for a fixed full output, EVERY 2-split position must yield
|
||||
// exactly the same flattened result as the unsplit parse. This kills
|
||||
// fragment-boundary bugs (mid-marker, mid-delimiter, mid-payload splits).
|
||||
DescribeTable("2-split fragment invariance",
|
||||
func(startInThought bool, full string) {
|
||||
refContent, refReasoning, refTools := flattenGemma4Deltas(
|
||||
parseGemma4Fragments(startInThought, []string{full}))
|
||||
for i := 0; i <= len(full); i++ {
|
||||
content, reasoning, tools := flattenGemma4Deltas(
|
||||
parseGemma4Fragments(startInThought, []string{full[:i], full[i:]}))
|
||||
Expect(content).To(Equal(refContent), fmt.Sprintf("content diverged at split %d", i))
|
||||
Expect(reasoning).To(Equal(refReasoning), fmt.Sprintf("reasoning diverged at split %d", i))
|
||||
Expect(tools).To(Equal(refTools), fmt.Sprintf("tool calls diverged at split %d", i))
|
||||
}
|
||||
},
|
||||
Entry("thought + content + two tool calls + turn end", false,
|
||||
"<|channel>thought\nPondering the request...\n<channel|>Sure - calling tools now. "+
|
||||
`<|tool_call>call:get_weather{location:<|"|>Paris, France<|"|>,unit:<|"|>celsius<|"|>,days:3,detailed:true}<tool_call|>`+
|
||||
`<|tool_call>call:get_time{timezone:<|"|>Europe/Lisbon<|"|>,nested:{flag:false,vals:[1,2.5,<|"|>x<|"|>]}}<tool_call|>`+
|
||||
"Done.<turn|>ignored tail"),
|
||||
Entry("startInThought + tool call + trailing partial marker", true,
|
||||
`Deep thought<channel|>final answer <|tool_call>call:noop{}<tool_call|> trailing <|tool`),
|
||||
Entry("malformed payload fallback", false,
|
||||
`pre <|tool_call>not a call<tool_call|> post`),
|
||||
)
|
||||
})
|
||||
|
||||
// Decoder-level ports of vLLM's TestParseGemma4Args / TestParseGemma4Array
|
||||
// (non-partial mode; the partial-withholding tests do not apply because this
|
||||
// parser only ever decodes COMPLETE payloads, see gemma4_parser.go).
|
||||
var _ = Describe("decodeGemma4Args", func() {
|
||||
DescribeTable("decodes the gemma4 call syntax into JSON arguments",
|
||||
func(in, wantJSON string) {
|
||||
Expect(decodeGemma4Args(in, 0)).To(MatchJSON(wantJSON))
|
||||
},
|
||||
// vLLM: test_empty_string / test_whitespace_only
|
||||
Entry("empty string", "", `{}`),
|
||||
Entry("whitespace only", " ", `{}`),
|
||||
// vLLM: test_single_string_value
|
||||
Entry("single string value", `location:<|"|>Paris<|"|>`, `{"location":"Paris"}`),
|
||||
// vLLM: test_string_value_with_comma
|
||||
Entry("string value with comma", `location:<|"|>Paris, France<|"|>`, `{"location":"Paris, France"}`),
|
||||
// vLLM: test_multiple_string_values
|
||||
Entry("multiple string values", `location:<|"|>San Francisco<|"|>,unit:<|"|>celsius<|"|>`, `{"location":"San Francisco","unit":"celsius"}`),
|
||||
// vLLM: test_integer_value / test_float_value
|
||||
Entry("integer value", "count:42", `{"count":42}`),
|
||||
Entry("float value", "score:3.14", `{"score":3.14}`),
|
||||
// vLLM: test_boolean_true / test_boolean_false
|
||||
Entry("boolean true", "flag:true", `{"flag":true}`),
|
||||
Entry("boolean false", "flag:false", `{"flag":false}`),
|
||||
// vLLM: test_null_value (bare null must become JSON null, not "null")
|
||||
Entry("null value", "param:null", `{"param":null}`),
|
||||
// vLLM: test_mixed_types
|
||||
Entry("mixed types", `name:<|"|>test<|"|>,count:42,active:true,score:3.14`,
|
||||
`{"name":"test","count":42,"active":true,"score":3.14}`),
|
||||
// vLLM: test_nested_object
|
||||
Entry("nested object", `nested:{inner:<|"|>value<|"|>}`, `{"nested":{"inner":"value"}}`),
|
||||
// vLLM: test_array_of_strings
|
||||
Entry("array of strings", `items:[<|"|>a<|"|>,<|"|>b<|"|>]`, `{"items":["a","b"]}`),
|
||||
// vLLM: test_unterminated_string (take everything after the delimiter)
|
||||
Entry("unterminated string", `key:<|"|>unterminated`, `{"key":"unterminated"}`),
|
||||
// vLLM: test_empty_value (key with no value after colon)
|
||||
Entry("empty value", "key:", `{"key":""}`),
|
||||
// vLLM: test_trailing_dot_float_partial_withheld, non-partial branch
|
||||
// (trailing-dot floats parse normally outside streaming).
|
||||
Entry("trailing dot float, complete payload", "left:108.,right:22.8", `{"left":108.0,"right":22.8}`),
|
||||
)
|
||||
|
||||
It("terminates and yields valid JSON on malformed input", func() {
|
||||
// vLLM: test_malformed_partial_array (the assertion there is only
|
||||
// "returns a dict without hanging"; ours is "valid JSON object").
|
||||
out := decodeGemma4Args(":[t:[]", 0)
|
||||
var v map[string]any
|
||||
Expect(json.Unmarshal([]byte(out), &v)).To(Succeed())
|
||||
})
|
||||
|
||||
It("degrades nesting beyond the recursion cap to a string value", func() {
|
||||
// 200 levels of a:{a:{...a:1...}}. Without the depth cap the mutual
|
||||
// recursion would grow the stack with the model's output; a Go stack
|
||||
// overflow is a fatal process kill, so levels past gemma4MaxArgsDepth
|
||||
// must gracefully fall back to the raw inner text as a JSON string.
|
||||
const depth = 200
|
||||
body := strings.Repeat("a:{", depth-1) + "a:1" + strings.Repeat("}", depth-1)
|
||||
out := decodeGemma4Args(body, 0)
|
||||
var v map[string]any
|
||||
Expect(json.Unmarshal([]byte(out), &v)).To(Succeed())
|
||||
levels := 0
|
||||
var cur any = v
|
||||
for {
|
||||
m, ok := cur.(map[string]any)
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
Expect(m).To(HaveKey("a"))
|
||||
cur = m["a"]
|
||||
levels++
|
||||
}
|
||||
Expect(levels).To(Equal(gemma4MaxArgsDepth + 1))
|
||||
Expect(cur).To(BeAssignableToTypeOf(""))
|
||||
Expect(cur).To(ContainSubstring("a:{"))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("decodeGemma4Array", func() {
|
||||
DescribeTable("decodes gemma4 array bodies into JSON arrays",
|
||||
func(in, wantJSON string) {
|
||||
Expect(decodeGemma4Array(in, 0)).To(MatchJSON(wantJSON))
|
||||
},
|
||||
// vLLM: test_string_array / test_empty_array / test_bare_values
|
||||
Entry("string array", `<|"|>a<|"|>,<|"|>b<|"|>`, `["a","b"]`),
|
||||
Entry("empty array", "", `[]`),
|
||||
Entry("bare values", "42,true,3.14", `[42,true,3.14]`),
|
||||
// vLLM: test_string_element_with_closing_bracket (a ']' inside a
|
||||
// delimited string must not close the array)
|
||||
Entry("string element with closing bracket", `[<|"|>a]b<|"|>,<|"|>c<|"|>],<|"|>tail<|"|>`, `[["a]b","c"],"tail"]`),
|
||||
// vLLM: test_stray_closing_bracket (no-progress abort, keep prefix)
|
||||
Entry("stray closing bracket", "42,]trailing", `[42]`),
|
||||
)
|
||||
})
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,406 +0,0 @@
|
||||
package main
|
||||
|
||||
// Renderer specs for RenderGemma4 against the canonical gemma4 chat template
|
||||
// (see the normative template comment in gemma4_renderer.go).
|
||||
//
|
||||
// Fixture provenance:
|
||||
// - "single user message" and "enable_thinking" are the EXACT expected
|
||||
// decodes from transformers tests/models/diffusion_gemma/
|
||||
// test_modeling_diffusion_gemma.py (test_diffusion_gemma_chat_template
|
||||
// and ..._with_thinking) with ONE difference: the transformers fixtures
|
||||
// start with "<bos>" because apply_chat_template tokenizes the rendered
|
||||
// text with add_bos. Our prompt goes through dllm_capi_generate, whose
|
||||
// run_generate already tokenizes with prepend_bos = vocab.add_bos
|
||||
// (dllm.cpp src/capi.cpp:230-231, true for gemma4), so the renderer must
|
||||
// NOT emit a literal <bos> (it would double) and every expected string
|
||||
// here drops that leading token.
|
||||
// - All other expected strings were produced by rendering the verbatim
|
||||
// GGUF template with jinja2 3.1.2 (bos_token="<bos>") and dropping the
|
||||
// leading "<bos>" for the same reason.
|
||||
|
||||
import (
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
)
|
||||
|
||||
// Two-function tools array used by the tool fixtures (OpenAI wire shape, as
|
||||
// LocalAI passes it through PredictOptions.Tools).
|
||||
const testToolsJSON = `[{"type":"function","function":{"name":"get_weather","description":"Get the current weather in a location.","parameters":{"type":"object","properties":{"location":{"type":"string","description":"The city name."},"unit":{"type":"string","enum":["celsius","fahrenheit"]}},"required":["location"]}}},{"type":"function","function":{"name":"get_time","description":"Get the current time in a timezone.","parameters":{"type":"object","properties":{"timezone":{"type":"string","description":"IANA timezone name."}},"required":["timezone"]}}}]`
|
||||
|
||||
// The <|tool>...<tool|> block the template renders for testToolsJSON inside
|
||||
// the system turn (jinja2-verified).
|
||||
const testToolsBlock = `<|tool>declaration:get_weather{description:<|"|>Get the current weather in a location.<|"|>,parameters:{properties:{location:{description:<|"|>The city name.<|"|>,type:<|"|>STRING<|"|>},unit:{enum:[<|"|>celsius<|"|>,<|"|>fahrenheit<|"|>],type:<|"|>STRING<|"|>}},required:[<|"|>location<|"|>],type:<|"|>OBJECT<|"|>}}<tool|><|tool>declaration:get_time{description:<|"|>Get the current time in a timezone.<|"|>,parameters:{properties:{timezone:{description:<|"|>IANA timezone name.<|"|>,type:<|"|>STRING<|"|>}},required:[<|"|>timezone<|"|>],type:<|"|>OBJECT<|"|>}}<tool|>`
|
||||
|
||||
// A single tool exercising the deep format_parameters branches: array items
|
||||
// (string-typed and nested-array), nullable, enum+nullable, nested object
|
||||
// properties/required, and a response declaration.
|
||||
const complexToolsJSON = `[{"type":"function","function":{"name":"complex_tool","description":"A complex tool.","parameters":{"type":"object","properties":{"tags":{"type":"array","description":"Tags.","items":{"type":"string"}},"matrix":{"type":"array","items":{"type":"array","items":{"type":"number"}}},"opts":{"type":"object","description":"Options.","properties":{"depth":{"type":"integer","nullable":true}},"required":["depth"]},"mode":{"type":"string","enum":["a","b"],"nullable":true}},"required":["tags","opts"]},"response":{"description":"The result.","type":"object"}}}]`
|
||||
|
||||
// jinja2-verified render of complexToolsJSON. Notable template quirks pinned
|
||||
// here: nested array items go through format_argument with ESCAPED keys and
|
||||
// an un-uppercased type (<|"|>type<|"|>:<|"|>number<|"|>), while direct item
|
||||
// types are uppercased; properties dictsort case-insensitively.
|
||||
const complexToolsBlock = `<|tool>declaration:complex_tool{description:<|"|>A complex tool.<|"|>,parameters:{properties:{matrix:{items:{items:{<|"|>type<|"|>:<|"|>number<|"|>},type:<|"|>ARRAY<|"|>},type:<|"|>ARRAY<|"|>},mode:{enum:[<|"|>a<|"|>,<|"|>b<|"|>],nullable:true,type:<|"|>STRING<|"|>},opts:{description:<|"|>Options.<|"|>,properties:{depth:{nullable:true,type:<|"|>INTEGER<|"|>}},required:[<|"|>depth<|"|>],type:<|"|>OBJECT<|"|>},tags:{description:<|"|>Tags.<|"|>,items:{type:<|"|>STRING<|"|>},type:<|"|>ARRAY<|"|>}},required:[<|"|>tags<|"|>,<|"|>opts<|"|>],type:<|"|>OBJECT<|"|>},response:{description:<|"|>The result.<|"|>,type:<|"|>OBJECT<|"|>}}<tool|>`
|
||||
|
||||
type renderGemma4Case struct {
|
||||
msgs []*pb.Message
|
||||
toolsJSON string
|
||||
// nImages mirrors len(PredictOptions.Images): the OpenAI layer strips
|
||||
// image content parts out of the messages, so the renderer re-injects
|
||||
// one engine marker per image on the last user message (see the IMAGE
|
||||
// NOTE on RenderGemma4).
|
||||
nImages int
|
||||
enableThinking bool
|
||||
noGenerationPrompt bool // inverted so the zero value is the common case
|
||||
expected string
|
||||
}
|
||||
|
||||
var _ = Describe("RenderGemma4", func() {
|
||||
DescribeTable("renders the canonical gemma4 prompt",
|
||||
func(c renderGemma4Case) {
|
||||
out, err := RenderGemma4(c.msgs, c.toolsJSON, c.nImages, c.enableThinking, !c.noGenerationPrompt)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out).To(Equal(c.expected))
|
||||
// The C-ABI generate prepends BOS itself: a literal <bos>
|
||||
// anywhere in the rendered prompt would double-encode it.
|
||||
Expect(out).ToNot(ContainSubstring("<bos>"))
|
||||
},
|
||||
|
||||
// transformers fixture (test_diffusion_gemma_chat_template), sans <bos>:
|
||||
// default thinking pre-opens an EMPTY thought channel in the
|
||||
// generation prompt.
|
||||
Entry("single user message, default (no thinking)", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "Write a long essay about Portugal."},
|
||||
},
|
||||
expected: "<|turn>user\nWrite a long essay about Portugal.<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
|
||||
}),
|
||||
|
||||
// transformers fixture (test_diffusion_gemma_chat_template_with_thinking),
|
||||
// sans <bos>: a system turn carrying <|think|> and NO auto-opened
|
||||
// thought channel.
|
||||
Entry("enable_thinking=true", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "Write a long essay about Portugal."},
|
||||
},
|
||||
enableThinking: true,
|
||||
expected: "<|turn>system\n<|think|>\n<turn|>\n<|turn>user\nWrite a long essay about Portugal.<turn|>\n<|turn>model\n",
|
||||
}),
|
||||
|
||||
Entry("multi-turn user/assistant/user", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "Hello, who are you?"},
|
||||
{Role: "assistant", Content: "I am Gemma, a helpful assistant."},
|
||||
{Role: "user", Content: "Tell me a joke."},
|
||||
},
|
||||
expected: "<|turn>user\nHello, who are you?<turn|>\n<|turn>model\nI am Gemma, a helpful assistant.<turn|>\n<|turn>user\nTell me a joke.<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
|
||||
}),
|
||||
|
||||
// tpl L178-L195: a leading system message is folded into the system
|
||||
// turn (trimmed) and consumed from the loop.
|
||||
Entry("system message folds into the system turn", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "system", Content: "You are a pirate."},
|
||||
{Role: "user", Content: "Hello!"},
|
||||
},
|
||||
expected: "<|turn>system\nYou are a pirate.<turn|>\n<|turn>user\nHello!<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
|
||||
}),
|
||||
|
||||
// tpl L182-L185: <|think|> goes at the very top of the SAME system
|
||||
// turn, before the system prompt text.
|
||||
Entry("system message with enable_thinking shares the turn", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "system", Content: "You are a pirate."},
|
||||
{Role: "user", Content: "Hello!"},
|
||||
},
|
||||
enableThinking: true,
|
||||
expected: "<|turn>system\n<|think|>\nYou are a pirate.<turn|>\n<|turn>user\nHello!<turn|>\n<|turn>model\n",
|
||||
}),
|
||||
|
||||
// tpl L196-L203: tool declarations render in the system turn, one
|
||||
// <|tool>declaration:...<tool|> block per tool, no separators.
|
||||
Entry("tools array (two functions)", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "What is the weather in Tokyo?"},
|
||||
},
|
||||
toolsJSON: testToolsJSON,
|
||||
expected: "<|turn>system\n" + testToolsBlock + "<turn|>\n<|turn>user\nWhat is the weather in Tokyo?<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
|
||||
}),
|
||||
|
||||
// format_parameters deep branches (tpl L1-L85) + response declaration
|
||||
// (tpl L106-L116).
|
||||
Entry("complex tool schema (array items, nullable, nested object, response)", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "go"},
|
||||
},
|
||||
toolsJSON: complexToolsJSON,
|
||||
expected: "<|turn>system\n" + complexToolsBlock + "<turn|>\n<|turn>user\ngo<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
|
||||
}),
|
||||
|
||||
// tpl L243-L313: assistant tool_calls render as
|
||||
// <|tool_call>call:name{args}<tool_call|>; the following role=tool
|
||||
// message renders inline as <|tool_response>response:name{value:..}
|
||||
// <tool_response|>; the model turn stays OPEN (no <turn|>, no new
|
||||
// generation prompt) so the model continues after the response.
|
||||
Entry("assistant tool_calls + role=tool result", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "What is the weather in Tokyo?"},
|
||||
{Role: "assistant", Content: "", ToolCalls: `[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"location\":\"Tokyo\",\"unit\":\"celsius\"}"}}]`},
|
||||
{Role: "tool", ToolCallId: "call_1", Content: "Sunny, 22 degrees celsius."},
|
||||
},
|
||||
toolsJSON: testToolsJSON,
|
||||
expected: "<|turn>system\n" + testToolsBlock + "<turn|>\n<|turn>user\nWhat is the weather in Tokyo?<turn|>\n<|turn>model\n" + `<|tool_call>call:get_weather{location:<|"|>Tokyo<|"|>,unit:<|"|>celsius<|"|>}<tool_call|><|tool_response>response:get_weather{value:<|"|>Sunny, 22 degrees celsius.<|"|>}<tool_response|>`,
|
||||
}),
|
||||
|
||||
// tpl L348-L349: a tool_calls turn with no rendered responses ends
|
||||
// on an OPEN <|tool_response> marker for the runtime to fill, and
|
||||
// add_generation_prompt adds nothing (tpl L357).
|
||||
Entry("assistant tool_calls without a result leaves <|tool_response> open", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "What is the weather in Tokyo?"},
|
||||
{Role: "assistant", Content: "", ToolCalls: `[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"location\":\"Tokyo\",\"unit\":\"celsius\"}"}}]`},
|
||||
},
|
||||
toolsJSON: testToolsJSON,
|
||||
expected: "<|turn>system\n" + testToolsBlock + "<turn|>\n<|turn>user\nWhat is the weather in Tokyo?<turn|>\n<|turn>model\n" + `<|tool_call>call:get_weather{location:<|"|>Tokyo<|"|>,unit:<|"|>celsius<|"|>}<tool_call|><|tool_response>`,
|
||||
}),
|
||||
|
||||
// tpl L237-L241: reasoning_content renders as a thought channel only
|
||||
// on a tool-calling turn after the last user message.
|
||||
Entry("reasoning_content with tool_calls renders the thought channel", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "weather?"},
|
||||
{Role: "assistant", Content: "", ReasoningContent: "I should call the tool", ToolCalls: `[{"index":0,"id":"c1","type":"function","function":{"name":"get_weather","arguments":"{\"location\":\"Tokyo\"}"}}]`},
|
||||
{Role: "tool", ToolCallId: "c1", Content: "Sunny"},
|
||||
},
|
||||
expected: "<|turn>user\nweather?<turn|>\n<|turn>model\n<|channel>thought\nI should call the tool\n<channel|>" + `<|tool_call>call:get_weather{location:<|"|>Tokyo<|"|>}<tool_call|><|tool_response>response:get_weather{value:<|"|>Sunny<|"|>}<tool_response|>`,
|
||||
}),
|
||||
|
||||
// tpl L220-L235: the assistant answer following its own tool round
|
||||
// continues the SAME model turn (no second <|turn>model).
|
||||
Entry("tool round then final assistant answer then user", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "weather?"},
|
||||
{Role: "assistant", Content: "", ToolCalls: `[{"index":0,"id":"c1","type":"function","function":{"name":"get_weather","arguments":"{\"location\":\"Tokyo\"}"}}]`},
|
||||
{Role: "tool", ToolCallId: "c1", Content: "Sunny"},
|
||||
{Role: "assistant", Content: "It is sunny."},
|
||||
{Role: "user", Content: "thanks"},
|
||||
},
|
||||
expected: "<|turn>user\nweather?<turn|>\n<|turn>model\n" + `<|tool_call>call:get_weather{location:<|"|>Tokyo<|"|>}<tool_call|><|tool_response>response:get_weather{value:<|"|>Sunny<|"|>}<tool_response|>` + "It is sunny.<turn|>\n<|turn>user\nthanks<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
|
||||
}),
|
||||
|
||||
// format_argument (tpl L118-L147): numbers keep their JSON literal,
|
||||
// booleans lower-case, nested maps have unquoted dictsorted keys,
|
||||
// arrays bracketed; top-level args are dictsorted case-insensitively.
|
||||
Entry("tool_call argument types (number/bool/nested/array)", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "go"},
|
||||
{Role: "assistant", Content: "", ToolCalls: `[{"index":0,"id":"c1","type":"function","function":{"name":"f","arguments":"{\"count\":42,\"ratio\":3.5,\"flag\":true,\"off\":false,\"nested\":{\"x\":\"y\",\"n\":7},\"list\":[\"a\",1,true]}"}}]`},
|
||||
},
|
||||
expected: "<|turn>user\ngo<turn|>\n<|turn>model\n" + `<|tool_call>call:f{count:42,flag:true,list:[<|"|>a<|"|>,1,true],nested:{n:7,x:<|"|>y<|"|>},off:false,ratio:3.5}<tool_call|><|tool_response>`,
|
||||
}),
|
||||
|
||||
// jinja dictsort is case-insensitive: alpha sorts before Beta.
|
||||
Entry("tool_call argument dictsort is case-insensitive", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "go"},
|
||||
{Role: "assistant", Content: "", ToolCalls: `[{"index":0,"id":"c1","type":"function","function":{"name":"f","arguments":"{\"Beta\":1,\"alpha\":2}"}}]`},
|
||||
},
|
||||
expected: "<|turn>user\ngo<turn|>\n<|turn>model\n<|tool_call>call:f{alpha:2,Beta:1}<tool_call|><|tool_response>",
|
||||
}),
|
||||
|
||||
// jinja renders Python None as "None" (round-trips through vLLM's
|
||||
// parser, which lowers "none" back to null).
|
||||
Entry("tool_call null argument renders as None", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "go"},
|
||||
{Role: "assistant", Content: "", ToolCalls: `[{"index":0,"id":"c1","type":"function","function":{"name":"f","arguments":"{\"maybe\":null}"}}]`},
|
||||
},
|
||||
expected: "<|turn>user\ngo<turn|>\n<|turn>model\n<|tool_call>call:f{maybe:None}<tool_call|><|tool_response>",
|
||||
}),
|
||||
|
||||
Entry("tool_call empty arguments render empty braces", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "go"},
|
||||
{Role: "assistant", Content: "", ToolCalls: `[{"index":0,"id":"c1","type":"function","function":{"name":"f","arguments":"{}"}}]`},
|
||||
},
|
||||
expected: "<|turn>user\ngo<turn|>\n<|turn>model\n<|tool_call>call:f{}<tool_call|><|tool_response>",
|
||||
}),
|
||||
|
||||
// tpl L253-L254: a non-object arguments string renders verbatim.
|
||||
Entry("tool_call non-object string arguments render verbatim", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "go"},
|
||||
{Role: "assistant", Content: "", ToolCalls: `[{"index":0,"id":"c1","type":"function","function":{"name":"f","arguments":"just text"}}]`},
|
||||
},
|
||||
expected: "<|turn>user\ngo<turn|>\n<|turn>model\n<|tool_call>call:f{just text}<tool_call|><|tool_response>",
|
||||
}),
|
||||
|
||||
// tpl L278-L285: unmatched tool_call_id falls back to the tool
|
||||
// message's own name.
|
||||
Entry("tool result name falls back when tool_call_id does not match", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "go"},
|
||||
{Role: "assistant", Content: "", ToolCalls: `[{"index":0,"id":"c1","type":"function","function":{"name":"f","arguments":"{}"}}]`},
|
||||
{Role: "tool", ToolCallId: "OTHER", Name: "named_tool", Content: "out"},
|
||||
},
|
||||
expected: "<|turn>user\ngo<turn|>\n<|turn>model\n" + `<|tool_call>call:f{}<tool_call|><|tool_response>response:named_tool{value:<|"|>out<|"|>}<tool_response|>`,
|
||||
}),
|
||||
|
||||
// strip_thinking (tpl L148-L158): historical assistant content loses
|
||||
// its <|channel>...<channel|> spans.
|
||||
Entry("assistant content thinking channels are stripped", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "hi"},
|
||||
{Role: "assistant", Content: "<|channel>thought\nsecret\n<channel|>visible answer"},
|
||||
{Role: "user", Content: "more"},
|
||||
},
|
||||
expected: "<|turn>user\nhi<turn|>\n<|turn>model\nvisible answer<turn|>\n<|turn>user\nmore<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
|
||||
}),
|
||||
|
||||
// tpl L220-L235: consecutive assistant messages suppress the second
|
||||
// <|turn>model (continuation), but each still closes with <turn|>.
|
||||
Entry("consecutive assistant messages continue the model turn", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "hi"},
|
||||
{Role: "assistant", Content: "part one"},
|
||||
{Role: "assistant", Content: "part two"},
|
||||
{Role: "user", Content: "ok"},
|
||||
},
|
||||
expected: "<|turn>user\nhi<turn|>\n<|turn>model\npart one<turn|>\npart two<turn|>\n<|turn>user\nok<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
|
||||
}),
|
||||
|
||||
Entry("add_generation_prompt=false renders no model turn", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "hi"},
|
||||
},
|
||||
noGenerationPrompt: true,
|
||||
expected: "<|turn>user\nhi<turn|>\n",
|
||||
}),
|
||||
|
||||
// One engine marker per image, appended directly after the user
|
||||
// text with no separator (tpl L323-L341 emits parts back-to-back;
|
||||
// "<image>" is dllm_capi.h's splice marker, not the template's
|
||||
// <|image|> text token - see the IMAGE NOTE on RenderGemma4).
|
||||
Entry("one image appends one engine marker to the user message", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "What is in this picture?"},
|
||||
},
|
||||
nImages: 1,
|
||||
expected: "<|turn>user\nWhat is in this picture?<image><turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
|
||||
}),
|
||||
|
||||
Entry("multiple images append markers in image order", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "Compare these."},
|
||||
},
|
||||
nImages: 3,
|
||||
expected: "<|turn>user\nCompare these.<image><image><image><turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
|
||||
}),
|
||||
|
||||
// Flattened delivery loses per-message attribution, so all images
|
||||
// attach to the LAST user message (llama.cpp grpc-server convention).
|
||||
Entry("images attach to the last user message in multi-turn", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: "hi"},
|
||||
{Role: "assistant", Content: "hello"},
|
||||
{Role: "user", Content: "and this?"},
|
||||
},
|
||||
nImages: 1,
|
||||
expected: "<|turn>user\nhi<turn|>\n<|turn>model\nhello<turn|>\n<|turn>user\nand this?<image><turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
|
||||
}),
|
||||
|
||||
// tpl L346: the markers count as captured_content, so an image-only
|
||||
// user message still has content and closes its turn normally.
|
||||
Entry("image with empty user text still closes the turn", renderGemma4Case{
|
||||
msgs: []*pb.Message{
|
||||
{Role: "user", Content: ""},
|
||||
},
|
||||
nImages: 1,
|
||||
expected: "<|turn>user\n<image><turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
|
||||
}),
|
||||
)
|
||||
|
||||
Describe("error handling", func() {
|
||||
It("fails loud on an unknown role", func() {
|
||||
_, err := RenderGemma4([]*pb.Message{
|
||||
{Role: "narrator", Content: "Meanwhile..."},
|
||||
}, "", 0, false, true)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring(`unknown role "narrator"`))
|
||||
})
|
||||
|
||||
It("fails on invalid tools JSON", func() {
|
||||
_, err := RenderGemma4([]*pb.Message{
|
||||
{Role: "user", Content: "hi"},
|
||||
}, "{not json", 0, false, true)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("tools JSON"))
|
||||
})
|
||||
|
||||
It("fails on invalid tool_calls JSON", func() {
|
||||
_, err := RenderGemma4([]*pb.Message{
|
||||
{Role: "user", Content: "hi"},
|
||||
{Role: "assistant", Content: "", ToolCalls: "{not json"},
|
||||
}, "", 0, false, true)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("tool_calls JSON"))
|
||||
})
|
||||
|
||||
It("fails on an orphan tool message, naming its index", func() {
|
||||
// A role:tool message with no preceding assistant tool_calls turn
|
||||
// would be silently dropped by the jinja; we fail loud instead.
|
||||
_, err := RenderGemma4([]*pb.Message{
|
||||
{Role: "user", Content: "hi"},
|
||||
{Role: "tool", Content: `{"temp": 20}`, ToolCallId: "call_1"},
|
||||
}, "", 0, false, true)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("orphan tool message 1"))
|
||||
})
|
||||
|
||||
It("fails on trailing garbage after the tools JSON array", func() {
|
||||
_, err := RenderGemma4([]*pb.Message{
|
||||
{Role: "user", Content: "hi"},
|
||||
}, "[] junk", 0, false, true)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("tools JSON"))
|
||||
})
|
||||
|
||||
It("fails when the tools JSON is not an array", func() {
|
||||
_, err := RenderGemma4([]*pb.Message{
|
||||
{Role: "user", Content: "hi"},
|
||||
}, `{"type":"function"}`, 0, false, true)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("tools JSON is not an array"))
|
||||
})
|
||||
|
||||
It("fails when a tools array element is not an object", func() {
|
||||
_, err := RenderGemma4([]*pb.Message{
|
||||
{Role: "user", Content: "hi"},
|
||||
}, `[42]`, 0, false, true)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("tools[0] is not an object"))
|
||||
})
|
||||
|
||||
It("rejects a nil message via the unknown-role check", func() {
|
||||
// Pins current behavior: pb getters are nil-safe, so a nil message
|
||||
// reads as role "" and trips the fail-loud unknown-role guard.
|
||||
_, err := RenderGemma4([]*pb.Message{nil}, "", 0, false, true)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring(`unknown role "" in message 0`))
|
||||
})
|
||||
|
||||
It("fails loud on images with no user message to attach them to", func() {
|
||||
// The engine would reject the markerless prompt anyway
|
||||
// (marker/image count mismatch); the renderer surfaces the bad
|
||||
// request with a usable message instead.
|
||||
_, err := RenderGemma4([]*pb.Message{
|
||||
{Role: "system", Content: "sys"},
|
||||
{Role: "assistant", Content: "hi"},
|
||||
}, "", 1, false, true)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("no user message"))
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -1,98 +0,0 @@
|
||||
package main
|
||||
|
||||
// Started internally by LocalAI - one gRPC server per loaded model.
|
||||
//
|
||||
// Loads libdllm.so via purego and registers the flat C-ABI declared in
|
||||
// dllm.cpp's include/dllm_capi.h (ABI v1): 9 mandatory symbols plus the
|
||||
// Dlsym-probed optional multimodal pair. The library name can
|
||||
// be overridden with DLLM_LIBRARY (mirrors the PARAKEET_LIBRARY /
|
||||
// WHISPER_LIBRARY convention in the sibling backends); the default looks
|
||||
// for the .so next to this binary (run.sh puts the package dir on
|
||||
// LD_LIBRARY_PATH).
|
||||
import (
|
||||
"flag"
|
||||
"fmt"
|
||||
"os"
|
||||
|
||||
"github.com/ebitengine/purego"
|
||||
grpc "github.com/mudler/LocalAI/pkg/grpc"
|
||||
)
|
||||
|
||||
var (
|
||||
addr = flag.String("addr", "localhost:50051", "the address to connect to")
|
||||
)
|
||||
|
||||
type LibFuncs struct {
|
||||
FuncPtr any
|
||||
Name string
|
||||
}
|
||||
|
||||
// loadCAPI dlopens libName and binds the 9 dllm_capi_* entry points 1:1 to
|
||||
// dllm_capi.h, so an `nm libdllm.so | grep dllm_capi` is enough to spot
|
||||
// drift. Shared with the test suite (ensureLibLoaded), which drives the
|
||||
// bridge without the gRPC server.
|
||||
//
|
||||
// The C-ABI returns malloc'd char* buffers from tokenize_json/generate; we
|
||||
// register those as uintptr so we get the raw pointer back and can call
|
||||
// dllm_capi_free_string on it (purego's string return would copy and forget
|
||||
// the original pointer, leaking it on every call). last_error returns a
|
||||
// BORROWED pointer instead, so it is registered as a plain string: purego
|
||||
// copies it and nothing must be freed.
|
||||
func loadCAPI(libName string) error {
|
||||
lib, err := purego.Dlopen(libName, purego.RTLD_NOW|purego.RTLD_GLOBAL)
|
||||
if err != nil {
|
||||
return fmt.Errorf("dllm: dlopen %q: %w", libName, err)
|
||||
}
|
||||
|
||||
libFuncs := []LibFuncs{
|
||||
{&cppAbiVersion, "dllm_capi_abi_version"},
|
||||
{&cppLoad, "dllm_capi_load"},
|
||||
{&cppFree, "dllm_capi_free"},
|
||||
{&cppLastError, "dllm_capi_last_error"},
|
||||
{&cppFreeString, "dllm_capi_free_string"},
|
||||
{&cppTokenizeJSON, "dllm_capi_tokenize_json"},
|
||||
{&cppGenerate, "dllm_capi_generate"},
|
||||
{&cppGenerateStream, "dllm_capi_generate_stream"},
|
||||
{&cppCancel, "dllm_capi_cancel"},
|
||||
}
|
||||
for _, lf := range libFuncs {
|
||||
purego.RegisterLibFunc(lf.FuncPtr, lib, lf.Name)
|
||||
}
|
||||
|
||||
// Multimodal entry points (dllm_capi.h's P4 surface). Additive: the ABI
|
||||
// version stays 1 and consumers detect the surface by probing the symbols
|
||||
// (the parakeet-cpp optional-symbol pattern), so the backend still loads
|
||||
// against an older text-only libdllm.so - image requests then fail with
|
||||
// errMMUnsupported instead of a boot failure.
|
||||
if sym, err := purego.Dlsym(lib, "dllm_capi_generate_mm"); err == nil && sym != 0 {
|
||||
purego.RegisterLibFunc(&cppGenerateMM, lib, "dllm_capi_generate_mm")
|
||||
}
|
||||
if sym, err := purego.Dlsym(lib, "dllm_capi_generate_stream_mm"); err == nil && sym != 0 {
|
||||
purego.RegisterLibFunc(&cppGenerateStreamMM, lib, "dllm_capi_generate_stream_mm")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func main() {
|
||||
libName := os.Getenv("DLLM_LIBRARY")
|
||||
if libName == "" {
|
||||
libName = "libdllm.so"
|
||||
}
|
||||
|
||||
if err := loadCAPI(libName); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
// Hard-fail on an ABI mismatch: the flat-pointer bindings above would
|
||||
// otherwise misbehave silently against a future libdllm.so.
|
||||
if v := cAbiVersion(); v != dllmABIVersion {
|
||||
panic(fmt.Errorf("dllm: libdllm.so ABI=%d, this backend speaks ABI=%d", v, dllmABIVersion))
|
||||
}
|
||||
fmt.Fprintf(os.Stderr, "[dllm] ABI=%d multimodal=%t\n", cAbiVersion(), cMMSupported())
|
||||
|
||||
flag.Parse()
|
||||
|
||||
if err := grpc.StartServer(*addr, &Dllm{}); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
@@ -1,24 +0,0 @@
|
||||
#!/bin/bash
|
||||
#
|
||||
# T1 packaging stub: copy the binary, run.sh and libdllm.so into package/.
|
||||
# The full ldd walk (libc, libstdc++, libgomp, GPU runtimes, arch
|
||||
# detection) lands with the registration task, mirroring
|
||||
# backend/go/whisper/package.sh.
|
||||
|
||||
set -e
|
||||
|
||||
CURDIR=$(dirname "$(realpath "$0")")
|
||||
|
||||
mkdir -p "$CURDIR/package/lib"
|
||||
|
||||
cp -avf "$CURDIR/dllm-grpc" "$CURDIR/package/"
|
||||
cp -avf "$CURDIR/run.sh" "$CURDIR/package/"
|
||||
|
||||
# libdllm.so + any soname symlinks, should upstream ever add them.
|
||||
cp -avf "$CURDIR"/libdllm.so* "$CURDIR/package/lib/" 2>/dev/null || {
|
||||
echo "ERROR: libdllm.so not found in $CURDIR, run 'make' first" >&2
|
||||
exit 1
|
||||
}
|
||||
|
||||
echo "T1 package layout (full ldd walk lands with registration):"
|
||||
ls -liah "$CURDIR/package/" "$CURDIR/package/lib/"
|
||||
@@ -1,16 +0,0 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
CURDIR=$(dirname "$(realpath "$0")")
|
||||
|
||||
export LD_LIBRARY_PATH="$CURDIR/lib:$CURDIR:${LD_LIBRARY_PATH:-}"
|
||||
|
||||
# If a self-contained ld.so was packaged, route through it so the
|
||||
# packaged libc / libstdc++ are used instead of the host's (matches the
|
||||
# whisper / parakeet-cpp backends' runtime layout).
|
||||
if [ -f "$CURDIR/lib/ld.so" ]; then
|
||||
echo "Using lib/ld.so"
|
||||
exec "$CURDIR/lib/ld.so" "$CURDIR/dllm-grpc" "$@"
|
||||
fi
|
||||
|
||||
exec "$CURDIR/dllm-grpc" "$@"
|
||||
@@ -1,6 +1,6 @@
|
||||
# parakeet-cpp backend Makefile.
|
||||
#
|
||||
# Upstream pin lives below as PARAKEET_VERSION?=e747acdaee69b916cef62263ae5f718bda9ff3f3
|
||||
# Upstream pin lives below as PARAKEET_VERSION?=1da853421de9710cbe894a0110711de5a0516486
|
||||
# (.github/bump_deps.sh) can find and update it - matches the
|
||||
# whisper.cpp / ds4 / vibevoice-cpp convention.
|
||||
#
|
||||
@@ -15,7 +15,7 @@
|
||||
# That's what the L0 smoke test uses. The default target below does the
|
||||
# proper clone-at-pin + cmake build so CI doesn't need a side-checkout.
|
||||
|
||||
PARAKEET_VERSION?=e747acdaee69b916cef62263ae5f718bda9ff3f3
|
||||
PARAKEET_VERSION?=1da853421de9710cbe894a0110711de5a0516486
|
||||
PARAKEET_REPO?=https://github.com/mudler/parakeet.cpp
|
||||
|
||||
GOCMD?=go
|
||||
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# stablediffusion.cpp (ggml)
|
||||
STABLEDIFFUSION_GGML_REPO?=https://github.com/leejet/stable-diffusion.cpp
|
||||
STABLEDIFFUSION_GGML_VERSION?=22516991cbdf725e69b0b4a87e52ca16cce07c2d
|
||||
STABLEDIFFUSION_GGML_VERSION?=2d0385ba85af358f7115dda608a63eafd9de7ffd
|
||||
|
||||
CMAKE_ARGS+=-DGGML_MAX_NAME=128
|
||||
|
||||
|
||||
6
backend/go/trellis2cpp/.gitignore
vendored
6
backend/go/trellis2cpp/.gitignore
vendored
@@ -1,6 +0,0 @@
|
||||
package/
|
||||
sources/
|
||||
.cache/
|
||||
build-*/
|
||||
variants/
|
||||
trellis2cpp
|
||||
@@ -1,132 +0,0 @@
|
||||
CMAKE_ARGS?=
|
||||
BUILD_TYPE?=
|
||||
NATIVE?=false
|
||||
|
||||
CURRENT_DIR=$(abspath ./)
|
||||
GOCMD?=go
|
||||
GO_TAGS?=
|
||||
JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# trellis2.cpp — C++/ggml port of Microsoft TRELLIS.2 (image -> 3D GLB).
|
||||
# The ggml submodule is pinned by trellis2cpp's .gitmodules and fetched via
|
||||
# --recursive. The commit pin lives here so bump_deps.yaml can update it.
|
||||
TRELLIS2CPP_REPO?=https://github.com/localai-org/trellis2cpp
|
||||
TRELLIS2CPP_VERSION?=73dfbe5dfc2cbefd0950853718086556c6d9b043
|
||||
|
||||
# libtrellis2 + ggml as shared libraries; no example/test binaries.
|
||||
CMAKE_ARGS+=-DCMAKE_BUILD_TYPE=Release
|
||||
CMAKE_ARGS+=-DBUILD_SHARED_LIBS=ON
|
||||
CMAKE_ARGS+=-DTRELLIS2_BUILD_EXAMPLES=OFF
|
||||
CMAKE_ARGS+=-DTRELLIS2_BUILD_TESTS=OFF
|
||||
# Print remeshing is part of trellis2cpp's ABI, so upstream owns the tested
|
||||
# CGAL/Boost versions, checksums, fetch logic, and update automation. LocalAI
|
||||
# only opts into that dependency set and pins the trellis2cpp commit above.
|
||||
CMAKE_ARGS+=-DTRELLIS2_FETCH_PRINT_REMESH_DEPS=ON
|
||||
CMAKE_ARGS+=-DTRELLIS2_PRINT_REMESH_DEPS_DIR=$(CURRENT_DIR)/sources/print-remesh-deps
|
||||
|
||||
ifeq ($(NATIVE),false)
|
||||
CMAKE_ARGS+=-DGGML_NATIVE=OFF
|
||||
endif
|
||||
|
||||
ifeq ($(BUILD_TYPE),cublas)
|
||||
CMAKE_ARGS+=-DGGML_CUDA=ON
|
||||
else ifeq ($(BUILD_TYPE),vulkan)
|
||||
CMAKE_ARGS+=-DGGML_VULKAN=ON
|
||||
else ifeq ($(BUILD_TYPE),hipblas)
|
||||
ROCM_HOME ?= /opt/rocm
|
||||
ROCM_PATH ?= /opt/rocm
|
||||
export CXX=$(ROCM_HOME)/llvm/bin/clang++
|
||||
export CC=$(ROCM_HOME)/llvm/bin/clang
|
||||
AMDGPU_TARGETS?=gfx908,gfx90a,gfx942,gfx950,gfx1030,gfx1100,gfx1101,gfx1102,gfx1200,gfx1201
|
||||
CMAKE_ARGS+=-DGGML_HIP=ON -DAMDGPU_TARGETS=$(AMDGPU_TARGETS)
|
||||
else ifeq ($(OS),Darwin)
|
||||
ifneq ($(BUILD_TYPE),metal)
|
||||
CMAKE_ARGS+=-DTRELLIS2_METAL=OFF -DGGML_METAL=OFF
|
||||
else
|
||||
# trellis2cpp turns on GGML_METAL(+EMBED_LIBRARY) itself when
|
||||
# TRELLIS2_METAL is enabled on Apple platforms.
|
||||
CMAKE_ARGS+=-DTRELLIS2_METAL=ON
|
||||
endif
|
||||
# Dependent libggml*.dylib resolve next to libtrellis2.dylib even
|
||||
# without DYLD_LIBRARY_PATH being exported.
|
||||
CMAKE_ARGS+=-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON -DCMAKE_INSTALL_RPATH=@loader_path
|
||||
endif
|
||||
|
||||
ifeq ($(BUILD_TYPE),sycl_f16)
|
||||
CMAKE_ARGS+=-DGGML_SYCL=ON \
|
||||
-DCMAKE_C_COMPILER=icx \
|
||||
-DCMAKE_CXX_COMPILER=icpx \
|
||||
-DGGML_SYCL_F16=ON
|
||||
endif
|
||||
|
||||
ifeq ($(BUILD_TYPE),sycl_f32)
|
||||
CMAKE_ARGS+=-DGGML_SYCL=ON \
|
||||
-DCMAKE_C_COMPILER=icx \
|
||||
-DCMAKE_CXX_COMPILER=icpx
|
||||
endif
|
||||
|
||||
sources/trellis2cpp:
|
||||
git clone --recursive $(TRELLIS2CPP_REPO) sources/trellis2cpp && \
|
||||
cd sources/trellis2cpp && \
|
||||
git checkout $(TRELLIS2CPP_VERSION) && \
|
||||
git submodule update --init --recursive --depth 1 --single-branch
|
||||
|
||||
# Detect OS
|
||||
UNAME_S := $(shell uname -s)
|
||||
UNAME_M := $(shell uname -m)
|
||||
|
||||
# The AVX variants are x86-only. ARM64 images use the portable fallback while
|
||||
# still enabling the selected GPU backend (Vulkan/CUDA) through CMAKE_ARGS.
|
||||
ifeq ($(UNAME_S),Linux)
|
||||
ifneq (,$(filter x86_64 amd64,$(UNAME_M)))
|
||||
VARIANTS = avx avx2 avx512 fallback
|
||||
else
|
||||
VARIANTS = fallback
|
||||
endif
|
||||
else
|
||||
# On non-Linux (e.g., Darwin), build only the fallback variant
|
||||
VARIANTS = fallback
|
||||
endif
|
||||
VARIANT_TARGETS = $(foreach v,$(VARIANTS),variants/$(v)/.built)
|
||||
|
||||
VARIANT_FLAGS_avx = -DGGML_AVX=on -DGGML_AVX2=off -DGGML_AVX512=off -DGGML_FMA=off -DGGML_F16C=off -DGGML_BMI2=off
|
||||
VARIANT_FLAGS_avx2 = -DGGML_AVX=on -DGGML_AVX2=on -DGGML_AVX512=off -DGGML_FMA=on -DGGML_F16C=on -DGGML_BMI2=on
|
||||
VARIANT_FLAGS_avx512 = -DGGML_AVX=on -DGGML_AVX2=on -DGGML_AVX512=on -DGGML_FMA=on -DGGML_F16C=on -DGGML_BMI2=on
|
||||
VARIANT_FLAGS_fallback = -DGGML_AVX=off -DGGML_AVX2=off -DGGML_AVX512=off -DGGML_FMA=off -DGGML_F16C=off -DGGML_BMI2=off
|
||||
|
||||
# libtrellis2 links libggml/libggml-base/libggml-cpu (+ the GPU backend) by
|
||||
# soname, and those sonames collide across SIMD variants — so each variant
|
||||
# lives in its own directory and run.sh selects one via LD_LIBRARY_PATH,
|
||||
# unlike stablediffusion-ggml's flat renamed-.so scheme.
|
||||
variants/%/.built: sources/trellis2cpp
|
||||
rm -rf build-$* variants/$*
|
||||
mkdir -p build-$* variants/$*
|
||||
cd build-$* && cmake ../sources/trellis2cpp $(CMAKE_ARGS) $(VARIANT_FLAGS_$*) && \
|
||||
cmake --build . --config Release -j$(JOBS)
|
||||
@for f in build-$*/libtrellis2.so build-$*/libtrellis2.dylib; do \
|
||||
if [ -e $$f ]; then cp -a $$f variants/$*/; fi; done
|
||||
find build-$*/ggml \( -name 'libggml*.so*' -o -name 'libggml*.dylib' \) -exec cp -a {} variants/$*/ \;
|
||||
rm -rf build-$*
|
||||
touch $@
|
||||
|
||||
trellis2cpp: main.go trellis2.go $(VARIANT_TARGETS)
|
||||
CGO_ENABLED=0 $(GOCMD) build -tags "$(GO_TAGS)" -o trellis2cpp ./
|
||||
|
||||
package: trellis2cpp
|
||||
bash package.sh
|
||||
|
||||
build: package
|
||||
|
||||
clean: purge
|
||||
rm -rf variants trellis2cpp package sources
|
||||
|
||||
purge:
|
||||
rm -rf build-*
|
||||
|
||||
# Weight-free by construction: pure-Go unit tests over model-path resolution,
|
||||
# validation, and request-parameter mapping. The multi-GB GGUF weights are
|
||||
# never downloaded in CI; end-to-end generation is exercised manually.
|
||||
test:
|
||||
$(GOCMD) test -v ./...
|
||||
|
||||
all: trellis2cpp package
|
||||
@@ -1,252 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
)
|
||||
|
||||
const (
|
||||
glbMagic = 0x46546c67
|
||||
glbJSONChunk = 0x4e4f534a
|
||||
glbBINChunk = 0x004e4942
|
||||
)
|
||||
|
||||
type glbAccessor struct {
|
||||
BufferView int `json:"bufferView"`
|
||||
ByteOffset int `json:"byteOffset"`
|
||||
ComponentType int `json:"componentType"`
|
||||
Count int `json:"count"`
|
||||
Type string `json:"type"`
|
||||
Normalized bool `json:"normalized"`
|
||||
}
|
||||
|
||||
type glbBufferView struct {
|
||||
Buffer int `json:"buffer"`
|
||||
ByteOffset int `json:"byteOffset"`
|
||||
ByteLength int `json:"byteLength"`
|
||||
ByteStride int `json:"byteStride"`
|
||||
}
|
||||
|
||||
type glbPrimitive struct {
|
||||
Attributes map[string]int `json:"attributes"`
|
||||
Indices *int `json:"indices"`
|
||||
}
|
||||
|
||||
type glbDocument struct {
|
||||
Accessors []glbAccessor `json:"accessors"`
|
||||
BufferViews []glbBufferView `json:"bufferViews"`
|
||||
Meshes []struct {
|
||||
Primitives []glbPrimitive `json:"primitives"`
|
||||
} `json:"meshes"`
|
||||
}
|
||||
|
||||
type glbVertexMesh struct {
|
||||
verts []float32
|
||||
tris []int32
|
||||
pbr []float32
|
||||
}
|
||||
|
||||
func glbLayout(componentType int, accessorType string) (componentBytes, components int, err error) {
|
||||
switch componentType {
|
||||
case 5121:
|
||||
componentBytes = 1
|
||||
case 5123:
|
||||
componentBytes = 2
|
||||
case 5125, 5126:
|
||||
componentBytes = 4
|
||||
default:
|
||||
return 0, 0, fmt.Errorf("unsupported GLB component type %d", componentType)
|
||||
}
|
||||
switch accessorType {
|
||||
case "SCALAR":
|
||||
components = 1
|
||||
case "VEC2":
|
||||
components = 2
|
||||
case "VEC3":
|
||||
components = 3
|
||||
case "VEC4":
|
||||
components = 4
|
||||
default:
|
||||
return 0, 0, fmt.Errorf("unsupported GLB accessor type %q", accessorType)
|
||||
}
|
||||
return componentBytes, components, nil
|
||||
}
|
||||
|
||||
func glbAccessorData(doc *glbDocument, binChunk []byte, index int) (glbAccessor, []byte, error) {
|
||||
if index < 0 || index >= len(doc.Accessors) {
|
||||
return glbAccessor{}, nil, fmt.Errorf("missing GLB accessor %d", index)
|
||||
}
|
||||
a := doc.Accessors[index]
|
||||
if a.BufferView < 0 || a.BufferView >= len(doc.BufferViews) {
|
||||
return glbAccessor{}, nil, fmt.Errorf("missing GLB buffer view %d", a.BufferView)
|
||||
}
|
||||
v := doc.BufferViews[a.BufferView]
|
||||
if v.Buffer != 0 || v.ByteStride != 0 {
|
||||
return glbAccessor{}, nil, fmt.Errorf("interleaved or external GLB buffers are unsupported")
|
||||
}
|
||||
componentBytes, components, err := glbLayout(a.ComponentType, a.Type)
|
||||
if err != nil {
|
||||
return glbAccessor{}, nil, err
|
||||
}
|
||||
if a.Count <= 0 || a.Count > math.MaxInt/(componentBytes*components) {
|
||||
return glbAccessor{}, nil, fmt.Errorf("invalid GLB accessor count %d", a.Count)
|
||||
}
|
||||
length := a.Count * componentBytes * components
|
||||
if v.ByteOffset < 0 || v.ByteLength < 0 || a.ByteOffset < 0 ||
|
||||
a.ByteOffset > v.ByteLength || length > v.ByteLength-a.ByteOffset ||
|
||||
length > len(binChunk) || v.ByteOffset > len(binChunk)-length-a.ByteOffset {
|
||||
return glbAccessor{}, nil, fmt.Errorf("GLB accessor %d is outside the BIN chunk", index)
|
||||
}
|
||||
start := v.ByteOffset + a.ByteOffset
|
||||
return a, binChunk[start : start+length], nil
|
||||
}
|
||||
|
||||
// parseVertexGLB reads the dense vertex-PBR form emitted by trellis2.cpp. GLB
|
||||
// coordinates and linear COLOR_0 values are converted back to the native
|
||||
// trellis coordinate/material convention before CGAL remeshing and rebaking.
|
||||
func parseVertexGLB(data []byte) (*glbVertexMesh, error) {
|
||||
if len(data) < 20 || binary.LittleEndian.Uint32(data[0:4]) != glbMagic {
|
||||
return nil, fmt.Errorf("input is not a GLB file")
|
||||
}
|
||||
if binary.LittleEndian.Uint32(data[4:8]) != 2 {
|
||||
return nil, fmt.Errorf("unsupported GLB version")
|
||||
}
|
||||
total := int(binary.LittleEndian.Uint32(data[8:12]))
|
||||
if total != len(data) {
|
||||
return nil, fmt.Errorf("invalid GLB length")
|
||||
}
|
||||
|
||||
var jsonChunk, binChunk []byte
|
||||
for offset := 12; offset <= len(data)-8; {
|
||||
length := int(binary.LittleEndian.Uint32(data[offset : offset+4]))
|
||||
chunkType := binary.LittleEndian.Uint32(data[offset+4 : offset+8])
|
||||
start := offset + 8
|
||||
if length < 0 || start > len(data)-length {
|
||||
return nil, fmt.Errorf("invalid GLB chunk length")
|
||||
}
|
||||
switch chunkType {
|
||||
case glbJSONChunk:
|
||||
if jsonChunk == nil {
|
||||
jsonChunk = data[start : start+length]
|
||||
}
|
||||
case glbBINChunk:
|
||||
if binChunk == nil {
|
||||
binChunk = data[start : start+length]
|
||||
}
|
||||
}
|
||||
offset = start + length
|
||||
}
|
||||
if jsonChunk == nil || binChunk == nil {
|
||||
return nil, fmt.Errorf("GLB must contain JSON and BIN chunks")
|
||||
}
|
||||
|
||||
var doc glbDocument
|
||||
if err := json.Unmarshal(jsonChunk, &doc); err != nil {
|
||||
return nil, fmt.Errorf("parsing GLB JSON: %w", err)
|
||||
}
|
||||
if len(doc.Meshes) != 1 || len(doc.Meshes[0].Primitives) != 1 {
|
||||
return nil, fmt.Errorf("GLB must contain one mesh primitive")
|
||||
}
|
||||
primitive := doc.Meshes[0].Primitives[0]
|
||||
positionIndex, ok := primitive.Attributes["POSITION"]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("GLB mesh has no POSITION attribute")
|
||||
}
|
||||
position, positionData, err := glbAccessorData(&doc, binChunk, positionIndex)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if position.ComponentType != 5126 || position.Type != "VEC3" {
|
||||
return nil, fmt.Errorf("GLB POSITION must be float32 VEC3")
|
||||
}
|
||||
|
||||
mesh := &glbVertexMesh{verts: make([]float32, position.Count*3)}
|
||||
for i := 0; i < position.Count; i++ {
|
||||
x := math.Float32frombits(binary.LittleEndian.Uint32(positionData[(i*3)*4:]))
|
||||
y := math.Float32frombits(binary.LittleEndian.Uint32(positionData[(i*3+1)*4:]))
|
||||
z := math.Float32frombits(binary.LittleEndian.Uint32(positionData[(i*3+2)*4:]))
|
||||
if math.IsNaN(float64(x)) || math.IsNaN(float64(y)) || math.IsNaN(float64(z)) ||
|
||||
math.IsInf(float64(x), 0) || math.IsInf(float64(y), 0) || math.IsInf(float64(z), 0) {
|
||||
return nil, fmt.Errorf("GLB POSITION contains a non-finite value")
|
||||
}
|
||||
mesh.verts[i*3] = x
|
||||
mesh.verts[i*3+1] = -z
|
||||
mesh.verts[i*3+2] = y
|
||||
}
|
||||
|
||||
if primitive.Indices == nil {
|
||||
if position.Count%3 != 0 {
|
||||
return nil, fmt.Errorf("unindexed GLB vertex count is not divisible by three")
|
||||
}
|
||||
mesh.tris = make([]int32, position.Count)
|
||||
for i := range mesh.tris {
|
||||
mesh.tris[i] = int32(i)
|
||||
}
|
||||
} else {
|
||||
indices, indexData, err := glbAccessorData(&doc, binChunk, *primitive.Indices)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if indices.Type != "SCALAR" || indices.Count%3 != 0 || (indices.ComponentType != 5123 && indices.ComponentType != 5125) {
|
||||
return nil, fmt.Errorf("GLB indices must be uint16/uint32 triangles")
|
||||
}
|
||||
mesh.tris = make([]int32, indices.Count)
|
||||
for i := range mesh.tris {
|
||||
var value uint32
|
||||
if indices.ComponentType == 5123 {
|
||||
value = uint32(binary.LittleEndian.Uint16(indexData[i*2:]))
|
||||
} else {
|
||||
value = binary.LittleEndian.Uint32(indexData[i*4:])
|
||||
}
|
||||
if value >= uint32(position.Count) || value > math.MaxInt32 {
|
||||
return nil, fmt.Errorf("GLB index %d is outside the vertex buffer", value)
|
||||
}
|
||||
mesh.tris[i] = int32(value)
|
||||
}
|
||||
}
|
||||
|
||||
colorIndex, hasColor := primitive.Attributes["COLOR_0"]
|
||||
if !hasColor {
|
||||
return mesh, nil
|
||||
}
|
||||
color, colorData, err := glbAccessorData(&doc, binChunk, colorIndex)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if color.ComponentType != 5123 || color.Type != "VEC4" || !color.Normalized || color.Count != position.Count {
|
||||
return nil, fmt.Errorf("GLB COLOR_0 must be normalized uint16 VEC4 aligned with POSITION")
|
||||
}
|
||||
metalRoughIndex, hasMetalRough := primitive.Attributes["_METALLIC_ROUGHNESS"]
|
||||
var metalRoughData []byte
|
||||
if hasMetalRough {
|
||||
metalRough, data, err := glbAccessorData(&doc, binChunk, metalRoughIndex)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if metalRough.ComponentType != 5121 || metalRough.Type != "VEC2" || !metalRough.Normalized || metalRough.Count != position.Count {
|
||||
return nil, fmt.Errorf("GLB _METALLIC_ROUGHNESS must be normalized uint8 VEC2 aligned with POSITION")
|
||||
}
|
||||
metalRoughData = data
|
||||
}
|
||||
|
||||
mesh.pbr = make([]float32, position.Count*6)
|
||||
for i := 0; i < position.Count; i++ {
|
||||
for channel := 0; channel < 3; channel++ {
|
||||
linear := float32(binary.LittleEndian.Uint16(colorData[(i*4+channel)*2:])) / 65535
|
||||
if linear <= 0.0031308 {
|
||||
mesh.pbr[i*6+channel] = linear * 12.92
|
||||
} else {
|
||||
mesh.pbr[i*6+channel] = 1.055*float32(math.Pow(float64(linear), 1.0/2.4)) - 0.055
|
||||
}
|
||||
}
|
||||
mesh.pbr[i*6+5] = float32(binary.LittleEndian.Uint16(colorData[(i*4+3)*2:])) / 65535
|
||||
mesh.pbr[i*6+4] = 0.6
|
||||
if hasMetalRough {
|
||||
mesh.pbr[i*6+3] = float32(metalRoughData[i*2]) / 255
|
||||
mesh.pbr[i*6+4] = float32(metalRoughData[i*2+1]) / 255
|
||||
}
|
||||
}
|
||||
return mesh, nil
|
||||
}
|
||||
@@ -1,74 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"math"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
func tinyVertexGLB() []byte {
|
||||
bin := make([]byte, 80)
|
||||
positions := []float32{1, 2, 3, 4, 5, 6, 7, 8, 9}
|
||||
for i, value := range positions {
|
||||
binary.LittleEndian.PutUint32(bin[i*4:], math.Float32bits(value))
|
||||
}
|
||||
colors := []uint16{
|
||||
65535, 0, 0, 65535,
|
||||
0, 65535, 0, 32768,
|
||||
0, 0, 65535, 65535,
|
||||
}
|
||||
for i, value := range colors {
|
||||
binary.LittleEndian.PutUint16(bin[36+i*2:], value)
|
||||
}
|
||||
copy(bin[60:], []byte{0, 153, 64, 128, 255, 32})
|
||||
for i, value := range []uint32{0, 1, 2} {
|
||||
binary.LittleEndian.PutUint32(bin[68+i*4:], value)
|
||||
}
|
||||
|
||||
jsonChunk := []byte(fmt.Sprintf(`{"asset":{"version":"2.0"},"meshes":[{"primitives":[{"attributes":{"POSITION":0,"COLOR_0":1,"_METALLIC_ROUGHNESS":2},"indices":3}]}],"accessors":[{"bufferView":0,"componentType":5126,"count":3,"type":"VEC3"},{"bufferView":1,"componentType":5123,"normalized":true,"count":3,"type":"VEC4"},{"bufferView":2,"componentType":5121,"normalized":true,"count":3,"type":"VEC2"},{"bufferView":3,"componentType":5125,"count":3,"type":"SCALAR"}],"bufferViews":[{"buffer":0,"byteOffset":0,"byteLength":36},{"buffer":0,"byteOffset":36,"byteLength":24},{"buffer":0,"byteOffset":60,"byteLength":6},{"buffer":0,"byteOffset":68,"byteLength":12}],"buffers":[{"byteLength":%d}]}`, len(bin)))
|
||||
for len(jsonChunk)%4 != 0 {
|
||||
jsonChunk = append(jsonChunk, ' ')
|
||||
}
|
||||
total := 12 + 8 + len(jsonChunk) + 8 + len(bin)
|
||||
glb := make([]byte, total)
|
||||
binary.LittleEndian.PutUint32(glb[0:], glbMagic)
|
||||
binary.LittleEndian.PutUint32(glb[4:], 2)
|
||||
binary.LittleEndian.PutUint32(glb[8:], uint32(total))
|
||||
binary.LittleEndian.PutUint32(glb[12:], uint32(len(jsonChunk)))
|
||||
binary.LittleEndian.PutUint32(glb[16:], glbJSONChunk)
|
||||
copy(glb[20:], jsonChunk)
|
||||
binHeader := 20 + len(jsonChunk)
|
||||
binary.LittleEndian.PutUint32(glb[binHeader:], uint32(len(bin)))
|
||||
binary.LittleEndian.PutUint32(glb[binHeader+4:], glbBINChunk)
|
||||
copy(glb[binHeader+8:], bin)
|
||||
return glb
|
||||
}
|
||||
|
||||
var _ = Describe("vertex GLB parsing for print remeshing", func() {
|
||||
It("restores trellis coordinates, topology, and PBR values", func() {
|
||||
mesh, err := parseVertexGLB(tinyVertexGLB())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(mesh.verts).To(Equal([]float32{1, -3, 2, 4, -6, 5, 7, -9, 8}))
|
||||
Expect(mesh.tris).To(Equal([]int32{0, 1, 2}))
|
||||
Expect(mesh.pbr).To(HaveLen(18))
|
||||
Expect(mesh.pbr[0]).To(BeNumerically("~", 1, 1e-5))
|
||||
Expect(mesh.pbr[3]).To(BeNumerically("~", 0, 1e-5))
|
||||
Expect(mesh.pbr[4]).To(BeNumerically("~", 0.6, 0.01))
|
||||
Expect(mesh.pbr[11]).To(BeNumerically("~", 32768.0/65535.0, 1e-5))
|
||||
})
|
||||
|
||||
It("rejects indices outside the source vertex buffer", func() {
|
||||
glb := tinyVertexGLB()
|
||||
binary.LittleEndian.PutUint32(glb[len(glb)-12:], 3)
|
||||
_, err := parseVertexGLB(glb)
|
||||
Expect(err).To(MatchError(ContainSubstring("outside the vertex buffer")))
|
||||
})
|
||||
|
||||
It("rejects non-GLB input", func() {
|
||||
_, err := parseVertexGLB([]byte("not a mesh"))
|
||||
Expect(err).To(MatchError("input is not a GLB file"))
|
||||
})
|
||||
})
|
||||
@@ -1,50 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"flag"
|
||||
"fmt"
|
||||
"os"
|
||||
"runtime"
|
||||
|
||||
"github.com/ebitengine/purego"
|
||||
grpc "github.com/mudler/LocalAI/pkg/grpc"
|
||||
)
|
||||
|
||||
var (
|
||||
addr = flag.String("addr", "localhost:50051", "the address to connect to")
|
||||
)
|
||||
|
||||
func registerLibFuncs(lib uintptr) {
|
||||
registerLibFuncsWith(func(fptr any, name string) {
|
||||
purego.RegisterLibFunc(fptr, lib, name)
|
||||
})
|
||||
}
|
||||
|
||||
func main() {
|
||||
// run.sh selects the CPU-variant directory and points TRELLIS2_LIBRARY at it.
|
||||
libName := os.Getenv("TRELLIS2_LIBRARY")
|
||||
if libName == "" {
|
||||
if runtime.GOOS == "darwin" {
|
||||
libName = "./variants/fallback/libtrellis2.dylib"
|
||||
} else {
|
||||
libName = "./variants/fallback/libtrellis2.so"
|
||||
}
|
||||
}
|
||||
|
||||
lib, err := purego.Dlopen(libName, purego.RTLD_NOW|purego.RTLD_GLOBAL)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
registerLibFuncs(lib)
|
||||
|
||||
if got := t2AbiVersion(); got != abiVersion {
|
||||
panic(fmt.Sprintf("trellis2 ABI mismatch: library reports %d, backend built for %d", got, abiVersion))
|
||||
}
|
||||
|
||||
flag.Parse()
|
||||
|
||||
if err := grpc.StartServer(*addr, &Trellis2{}); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
@@ -1,63 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Script to copy the appropriate libraries based on architecture
|
||||
# This script is used in the final stage of the Dockerfile
|
||||
|
||||
set -e
|
||||
|
||||
CURDIR=$(dirname "$(realpath $0)")
|
||||
REPO_ROOT="${CURDIR}/../../.."
|
||||
|
||||
# Create lib directory
|
||||
mkdir -p $CURDIR/package/lib
|
||||
|
||||
# Each CPU variant keeps its libtrellis2 + libggml* set in its own directory
|
||||
# (their sonames collide across variants); run.sh selects one at startup.
|
||||
cp -a $CURDIR/variants $CURDIR/package/
|
||||
cp -avf $CURDIR/trellis2cpp $CURDIR/package/
|
||||
cp -fv $CURDIR/run.sh $CURDIR/package/
|
||||
|
||||
# Detect architecture and copy appropriate libraries
|
||||
if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then
|
||||
# x86_64 architecture
|
||||
echo "Detected x86_64 architecture, copying x86_64 libraries..."
|
||||
cp -arfLv /lib64/ld-linux-x86-64.so.2 $CURDIR/package/lib/ld.so
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libc.so.6 $CURDIR/package/lib/libc.so.6
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libgcc_s.so.1 $CURDIR/package/lib/libgcc_s.so.1
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libstdc++.so.6 $CURDIR/package/lib/libstdc++.so.6
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libm.so.6 $CURDIR/package/lib/libm.so.6
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libgomp.so.1 $CURDIR/package/lib/libgomp.so.1
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libdl.so.2 $CURDIR/package/lib/libdl.so.2
|
||||
cp -arfLv /lib/x86_64-linux-gnu/librt.so.1 $CURDIR/package/lib/librt.so.1
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libpthread.so.0 $CURDIR/package/lib/libpthread.so.0
|
||||
elif [ -f "/lib/ld-linux-aarch64.so.1" ]; then
|
||||
# ARM64 architecture
|
||||
echo "Detected ARM64 architecture, copying ARM64 libraries..."
|
||||
cp -arfLv /lib/ld-linux-aarch64.so.1 $CURDIR/package/lib/ld.so
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libc.so.6 $CURDIR/package/lib/libc.so.6
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libgcc_s.so.1 $CURDIR/package/lib/libgcc_s.so.1
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libstdc++.so.6 $CURDIR/package/lib/libstdc++.so.6
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libm.so.6 $CURDIR/package/lib/libm.so.6
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libgomp.so.1 $CURDIR/package/lib/libgomp.so.1
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libdl.so.2 $CURDIR/package/lib/libdl.so.2
|
||||
cp -arfLv /lib/aarch64-linux-gnu/librt.so.1 $CURDIR/package/lib/librt.so.1
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libpthread.so.0 $CURDIR/package/lib/libpthread.so.0
|
||||
elif [ $(uname -s) = "Darwin" ]; then
|
||||
echo "Detected Darwin"
|
||||
else
|
||||
echo "Error: Could not detect architecture"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Package GPU libraries based on BUILD_TYPE
|
||||
# The GPU library packaging script will detect BUILD_TYPE and copy appropriate GPU libraries
|
||||
GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh"
|
||||
if [ -f "$GPU_LIB_SCRIPT" ]; then
|
||||
echo "Packaging GPU libraries for BUILD_TYPE=${BUILD_TYPE:-cpu}..."
|
||||
source "$GPU_LIB_SCRIPT" "$CURDIR/package/lib"
|
||||
package_gpu_libs
|
||||
fi
|
||||
|
||||
echo "Packaging completed successfully"
|
||||
ls -liah $CURDIR/package/
|
||||
ls -liah $CURDIR/package/lib/
|
||||
@@ -1,61 +0,0 @@
|
||||
#!/bin/bash
|
||||
set -ex
|
||||
|
||||
# Get the absolute current dir where the script is located
|
||||
CURDIR=$(dirname "$(realpath "$0")")
|
||||
|
||||
cd /
|
||||
|
||||
echo "CPU info:"
|
||||
if [ "$(uname)" != "Darwin" ]; then
|
||||
grep -e "model\sname" /proc/cpuinfo | head -1
|
||||
grep -e "flags" /proc/cpuinfo | head -1
|
||||
fi
|
||||
|
||||
# Each variant directory bundles libtrellis2 plus its libggml* set (the ggml
|
||||
# sonames collide across SIMD variants, so they can't share one directory).
|
||||
VARIANT=fallback
|
||||
|
||||
if [ "$(uname)" = "Darwin" ]; then
|
||||
LIBRARY="$CURDIR/variants/$VARIANT/libtrellis2.dylib"
|
||||
if [ ! -e "$LIBRARY" ]; then
|
||||
LIBRARY="$CURDIR/variants/$VARIANT/libtrellis2.so"
|
||||
fi
|
||||
export DYLD_LIBRARY_PATH="$CURDIR/variants/$VARIANT:$CURDIR/lib:$DYLD_LIBRARY_PATH"
|
||||
else
|
||||
if grep -q -e "\savx\s" /proc/cpuinfo ; then
|
||||
echo "CPU: AVX found OK"
|
||||
if [ -d "$CURDIR/variants/avx" ]; then
|
||||
VARIANT=avx
|
||||
fi
|
||||
fi
|
||||
|
||||
if grep -q -e "\savx2\s" /proc/cpuinfo ; then
|
||||
echo "CPU: AVX2 found OK"
|
||||
if [ -d "$CURDIR/variants/avx2" ]; then
|
||||
VARIANT=avx2
|
||||
fi
|
||||
fi
|
||||
|
||||
if grep -q -e "\savx512f\s" /proc/cpuinfo ; then
|
||||
echo "CPU: AVX512F found OK"
|
||||
if [ -d "$CURDIR/variants/avx512" ]; then
|
||||
VARIANT=avx512
|
||||
fi
|
||||
fi
|
||||
|
||||
LIBRARY="$CURDIR/variants/$VARIANT/libtrellis2.so"
|
||||
export LD_LIBRARY_PATH="$CURDIR/variants/$VARIANT:$CURDIR/lib:$LD_LIBRARY_PATH"
|
||||
fi
|
||||
|
||||
export TRELLIS2_LIBRARY=$LIBRARY
|
||||
|
||||
# If there is a lib/ld.so, use it
|
||||
if [ -f "$CURDIR"/lib/ld.so ]; then
|
||||
echo "Using lib/ld.so"
|
||||
echo "Using library: $LIBRARY"
|
||||
exec "$CURDIR"/lib/ld.so "$CURDIR"/trellis2cpp "$@"
|
||||
fi
|
||||
|
||||
echo "Using library: $LIBRARY"
|
||||
exec "$CURDIR"/trellis2cpp "$@"
|
||||
@@ -1,511 +0,0 @@
|
||||
package main
|
||||
|
||||
// trellis2.go — purego bindings to libtrellis2's flat C ABI (trellis2_capi.h)
|
||||
// plus the LocalAI backend implementation. Adapted from the upstream demo
|
||||
// server's engine.go; the t2_abi_version binding guards against header/library
|
||||
// drift.
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math/rand"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"unsafe"
|
||||
|
||||
"github.com/mudler/LocalAI/pkg/grpc/base"
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/LocalAI/pkg/utils"
|
||||
)
|
||||
|
||||
const abiVersion = 11
|
||||
|
||||
// Pipeline types (enum t2_pipeline_type) and background modes
|
||||
// (enum t2_background_mode).
|
||||
const (
|
||||
pipeAuto = 0
|
||||
pipeCoarse = 1
|
||||
pipe512 = 2
|
||||
pipe1024 = 3
|
||||
|
||||
backgroundAuto = 0
|
||||
backgroundKeep = 1
|
||||
backgroundBlack = 2
|
||||
backgroundWhite = 3
|
||||
)
|
||||
|
||||
// The bindings use typed pointers (*float32/*int32/*byte) rather than uintptr
|
||||
// for C-owned buffers so no uintptr->unsafe.Pointer conversions are needed;
|
||||
// only the opaque t2_pipeline / t2_mesh_result handles stay uintptr.
|
||||
var (
|
||||
t2AbiVersion func() int32
|
||||
t2PipelineLoad func(dino, ssFlow, ssDec, slatFlow, slatFlowHR, shapeDec,
|
||||
shapeEnc, texDec, texFlow, texFlowHR string,
|
||||
flags int32, err *byte, errLen int32) uintptr
|
||||
t2PipelineFree func(p uintptr)
|
||||
t2PipelineBackend func(p uintptr) string
|
||||
t2PipelineCaps func(p uintptr) int32
|
||||
t2Generate func(p uintptr, img *byte, imgLen int32,
|
||||
pipelineType, backgroundMode int32, seed uint64, steps int32,
|
||||
guidance float32, textureSteps int32,
|
||||
progress, user, preview, previewUser uintptr,
|
||||
err *byte, errLen int32) uintptr
|
||||
t2MeshNVerts func(r uintptr) int32
|
||||
t2MeshNTris func(r uintptr) int32
|
||||
t2MeshVerts func(r uintptr) *float32
|
||||
t2MeshTris func(r uintptr) *int32
|
||||
t2MeshHasPBR func(r uintptr) int32
|
||||
t2MeshPBR func(r uintptr) *float32
|
||||
t2MeshFree func(r uintptr)
|
||||
t2BakeGLB func(verts *float32, nv int32, tris *int32, nt int32,
|
||||
pbr *float32, texSize, componentFilter int32,
|
||||
outLen *int32, err *byte, errLen int32) *byte
|
||||
// CGAL Alpha Wrap print remeshing — availability is fixed at library build
|
||||
// time, so gate every use on t2_print_remesh_available.
|
||||
t2PrintRemeshAvailable func() int32
|
||||
t2PreparePrintMesh func(verts *float32, nv int32, tris *int32, nt int32,
|
||||
pbr *float32, componentFilter int32, alphaRatio, offsetRatio float32,
|
||||
err *byte, errLen int32) uintptr
|
||||
t2BakeProjectedGLB func(targetVerts *float32, targetNV int32,
|
||||
targetTris *int32, targetNT int32,
|
||||
sourceVerts *float32, sourceNV int32,
|
||||
sourceTris *int32, sourceNT int32,
|
||||
sourcePBR *float32, texSize, sourceComponentFilter int32,
|
||||
outLen *int32, err *byte, errLen int32) *byte
|
||||
t2FreeBuffer func(buf *byte)
|
||||
)
|
||||
|
||||
type libFunc struct {
|
||||
funcPtr any
|
||||
name string
|
||||
}
|
||||
|
||||
func registerLibFuncsWith(register func(fptr any, name string)) {
|
||||
for _, lf := range []libFunc{
|
||||
{&t2AbiVersion, "t2_abi_version"},
|
||||
{&t2PipelineLoad, "t2_pipeline_load"},
|
||||
{&t2PipelineFree, "t2_pipeline_free"},
|
||||
{&t2PipelineBackend, "t2_pipeline_backend"},
|
||||
{&t2PipelineCaps, "t2_pipeline_caps"},
|
||||
{&t2Generate, "t2_generate"},
|
||||
{&t2MeshNVerts, "t2_mesh_n_verts"},
|
||||
{&t2MeshNTris, "t2_mesh_n_tris"},
|
||||
{&t2MeshVerts, "t2_mesh_verts"},
|
||||
{&t2MeshTris, "t2_mesh_tris"},
|
||||
{&t2MeshHasPBR, "t2_mesh_has_pbr"},
|
||||
{&t2MeshPBR, "t2_mesh_pbr"},
|
||||
{&t2MeshFree, "t2_mesh_free"},
|
||||
{&t2BakeGLB, "t2_bake_glb"},
|
||||
{&t2PrintRemeshAvailable, "t2_print_remesh_available"},
|
||||
{&t2PreparePrintMesh, "t2_prepare_print_mesh"},
|
||||
{&t2BakeProjectedGLB, "t2_bake_projected_glb"},
|
||||
{&t2FreeBuffer, "t2_free_buffer"},
|
||||
} {
|
||||
register(lf.funcPtr, lf.name)
|
||||
}
|
||||
}
|
||||
|
||||
// modelSet holds the resolved path for every pipeline role; optional roles are
|
||||
// "" when disabled (the C side treats NULL/"" as "omit").
|
||||
type modelSet struct {
|
||||
dino, ssFlow, ssDec string
|
||||
slatFlow, slatFlow1024 string
|
||||
shapeDec string
|
||||
shapeEnc, texDec string
|
||||
texSlatFlow512, texSlatFlow1024 string
|
||||
}
|
||||
|
||||
// role → (option key, default filename) in t2_pipeline_load argument order.
|
||||
// The option keys follow the sd-ggml `*_path` convention; the default
|
||||
// filenames are the ones the upstream converters emit and the demo server
|
||||
// looks up, so a gallery install needs no options at all.
|
||||
type modelRole struct {
|
||||
key string
|
||||
filename string
|
||||
required bool
|
||||
assign func(*modelSet, string)
|
||||
}
|
||||
|
||||
var modelRoles = []modelRole{
|
||||
{"dino_path", "dino_f16.gguf", true, func(s *modelSet, p string) { s.dino = p }},
|
||||
{"ss_flow_path", "ss_flow_f16.gguf", true, func(s *modelSet, p string) { s.ssFlow = p }},
|
||||
{"ss_dec_path", "ss_dec_f16.gguf", true, func(s *modelSet, p string) { s.ssDec = p }},
|
||||
{"slat_flow_path", "slat_flow_f16.gguf", false, func(s *modelSet, p string) { s.slatFlow = p }},
|
||||
{"slat_flow_1024_path", "slat_flow_1024_f16.gguf", false, func(s *modelSet, p string) { s.slatFlow1024 = p }},
|
||||
{"shape_dec_path", "shape_dec_f16.gguf", false, func(s *modelSet, p string) { s.shapeDec = p }},
|
||||
{"shape_enc_path", "shape_enc_f16.gguf", false, func(s *modelSet, p string) { s.shapeEnc = p }},
|
||||
{"tex_dec_path", "tex_dec_f16.gguf", false, func(s *modelSet, p string) { s.texDec = p }},
|
||||
{"tex_slat_flow_512_path", "tex_slat_flow_512_f16.gguf", false, func(s *modelSet, p string) { s.texSlatFlow512 = p }},
|
||||
{"tex_slat_flow_1024_path", "tex_slat_flow_1024_f16.gguf", false, func(s *modelSet, p string) { s.texSlatFlow1024 = p }},
|
||||
}
|
||||
|
||||
// resolveModels maps LocalAI's model file + options onto the ten pipeline
|
||||
// roles. The model file only anchors the GGUF directory; each role resolves
|
||||
// to an explicit `<role>_path` option when given, else to its default
|
||||
// filename in that directory. Missing required files refuse the load (a
|
||||
// backend must not capture arbitrary GGUFs — see issue #9287); missing
|
||||
// optional files degrade capabilities the same way the upstream demo does.
|
||||
func resolveModels(modelFile, modelPath string, options []string) (modelSet, error) {
|
||||
base := modelFile
|
||||
if !filepath.IsAbs(base) {
|
||||
base = filepath.Join(modelPath, base)
|
||||
}
|
||||
ggufDir := filepath.Dir(base)
|
||||
|
||||
overrides := map[string]string{}
|
||||
for _, op := range options {
|
||||
key, value, found := strings.Cut(op, ":")
|
||||
if !found || !strings.HasSuffix(key, "_path") {
|
||||
continue
|
||||
}
|
||||
if !filepath.IsAbs(value) {
|
||||
value = filepath.Join(modelPath, value)
|
||||
if err := utils.VerifyPath(value, modelPath); err != nil {
|
||||
return modelSet{}, fmt.Errorf("option %s: %w", key, err)
|
||||
}
|
||||
}
|
||||
overrides[key] = value
|
||||
}
|
||||
|
||||
var set modelSet
|
||||
var missingRequired []string
|
||||
for _, role := range modelRoles {
|
||||
path, explicit := overrides[role.key]
|
||||
if !explicit {
|
||||
path = filepath.Join(ggufDir, role.filename)
|
||||
}
|
||||
if _, err := os.Stat(path); err != nil {
|
||||
if explicit {
|
||||
return modelSet{}, fmt.Errorf("option %s points at a missing file: %s", role.key, path)
|
||||
}
|
||||
if role.required {
|
||||
missingRequired = append(missingRequired, role.filename)
|
||||
}
|
||||
path = ""
|
||||
}
|
||||
role.assign(&set, path)
|
||||
}
|
||||
if len(missingRequired) > 0 {
|
||||
return modelSet{}, fmt.Errorf("not a trellis2 model set: missing required %s in %s", strings.Join(missingRequired, ", "), ggufDir)
|
||||
}
|
||||
|
||||
// Degradation mirrors the upstream demo: the 512 pair enables everything
|
||||
// finer than coarse; texturing needs its three-model set; a textured 1024
|
||||
// cascade additionally needs the HR texture flow.
|
||||
if set.slatFlow == "" || set.shapeDec == "" {
|
||||
set.slatFlow, set.shapeDec = "", ""
|
||||
set.slatFlow1024 = ""
|
||||
set.shapeEnc, set.texDec, set.texSlatFlow512, set.texSlatFlow1024 = "", "", "", ""
|
||||
return set, nil
|
||||
}
|
||||
if set.shapeEnc == "" || set.texDec == "" || set.texSlatFlow512 == "" {
|
||||
set.shapeEnc, set.texDec, set.texSlatFlow512, set.texSlatFlow1024 = "", "", "", ""
|
||||
} else if set.texSlatFlow1024 == "" {
|
||||
set.slatFlow1024 = ""
|
||||
}
|
||||
return set, nil
|
||||
}
|
||||
|
||||
func pipelineForQuality(quality string) int32 {
|
||||
switch quality {
|
||||
case "coarse":
|
||||
return pipeCoarse
|
||||
case "512":
|
||||
return pipe512
|
||||
case "1024":
|
||||
return pipe1024
|
||||
default:
|
||||
return pipeAuto
|
||||
}
|
||||
}
|
||||
|
||||
func backgroundForMode(background string) int32 {
|
||||
switch background {
|
||||
case "keep":
|
||||
return backgroundKeep
|
||||
case "black":
|
||||
return backgroundBlack
|
||||
case "white":
|
||||
return backgroundWhite
|
||||
default:
|
||||
return backgroundAuto
|
||||
}
|
||||
}
|
||||
|
||||
func componentFilterFor(components string) int32 {
|
||||
switch components {
|
||||
case "tiny":
|
||||
return 0 // remove only tiny islands
|
||||
case "largest":
|
||||
return 1 // keep the largest connected component
|
||||
default:
|
||||
return 2 // preserve every connected component (demo default)
|
||||
}
|
||||
}
|
||||
|
||||
func atoiOr(s string, fallback int32) int32 {
|
||||
if s == "" {
|
||||
return fallback
|
||||
}
|
||||
n, err := strconv.Atoi(s)
|
||||
if err != nil {
|
||||
return fallback
|
||||
}
|
||||
return int32(n)
|
||||
}
|
||||
|
||||
func boolParam(s string) bool {
|
||||
return s == "1" || strings.EqualFold(s, "true")
|
||||
}
|
||||
|
||||
// ratioOr parses a fraction-of-bounding-box-diagonal parameter. Out-of-range
|
||||
// or unparseable values fall back rather than error, mirroring atoiOr; the
|
||||
// accepted range matches what the upstream demo clamps to.
|
||||
func ratioOr(s string, fallback float32) float32 {
|
||||
if s == "" {
|
||||
return fallback
|
||||
}
|
||||
f, err := strconv.ParseFloat(s, 32)
|
||||
if err != nil || f < 0.00001 || f > 0.5 {
|
||||
return fallback
|
||||
}
|
||||
return float32(f)
|
||||
}
|
||||
|
||||
type Trellis2 struct {
|
||||
base.SingleThread
|
||||
// t2_generate is not thread-safe per pipeline. The gRPC server already
|
||||
// serializes calls via Locking(), but keep a local mutex too so the
|
||||
// invariant doesn't depend on the transport.
|
||||
mu sync.Mutex
|
||||
pipeline uintptr
|
||||
}
|
||||
|
||||
func (t *Trellis2) Load(opts *pb.ModelOptions) error {
|
||||
set, err := resolveModels(opts.ModelFile, opts.ModelPath, opts.Options)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
errBuf := make([]byte, 512)
|
||||
p := t2PipelineLoad(set.dino, set.ssFlow, set.ssDec,
|
||||
set.slatFlow, set.slatFlow1024, set.shapeDec,
|
||||
set.shapeEnc, set.texDec, set.texSlatFlow512, set.texSlatFlow1024,
|
||||
0 /*flags*/, &errBuf[0], int32(len(errBuf)))
|
||||
if p == 0 {
|
||||
return fmt.Errorf("trellis2 pipeline load: %s", cstr(errBuf))
|
||||
}
|
||||
|
||||
t.mu.Lock()
|
||||
if t.pipeline != 0 {
|
||||
t2PipelineFree(t.pipeline)
|
||||
}
|
||||
t.pipeline = p
|
||||
t.mu.Unlock()
|
||||
|
||||
fmt.Fprintf(os.Stderr, "trellis2 pipeline loaded: backend=%s caps=%#x\n",
|
||||
t2PipelineBackend(p), t2PipelineCaps(p))
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *Trellis2) Free() error {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
if t.pipeline != 0 {
|
||||
t2PipelineFree(t.pipeline)
|
||||
t.pipeline = 0
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *Trellis2) Generate3D(opts *pb.Generate3DRequest) error {
|
||||
if opts.Dst == "" {
|
||||
return fmt.Errorf("dst is empty")
|
||||
}
|
||||
if opts.GetParams()["operation"] == "print_remesh" {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
return remeshGLB(opts)
|
||||
}
|
||||
img, err := os.ReadFile(opts.Src)
|
||||
if err != nil {
|
||||
return fmt.Errorf("reading conditioning image: %w", err)
|
||||
}
|
||||
if len(img) == 0 {
|
||||
return fmt.Errorf("conditioning image is empty")
|
||||
}
|
||||
|
||||
seed := uint64(opts.Seed)
|
||||
if opts.Seed <= 0 {
|
||||
seed = rand.Uint64()
|
||||
}
|
||||
guidance := opts.CfgScale
|
||||
if guidance <= 0 {
|
||||
guidance = -1 // <0 selects the pipeline default (7.5)
|
||||
}
|
||||
texSize := atoiOr(opts.GetParams()["texture_size"], 0) // <=0 selects the bake default
|
||||
componentFilter := componentFilterFor(opts.GetParams()["components"])
|
||||
|
||||
// Optional CGAL Alpha Wrap: wrap the generated mesh into a watertight,
|
||||
// intersection-free 2-manifold for 3D printing. Ratios are fractions of
|
||||
// the bounding-box diagonal; offset defaults to alpha/30 per the CGAL
|
||||
// guideline the upstream demo uses. Offset is deliberately not an
|
||||
// independent parameter: looser values produce puffy or degenerate wraps.
|
||||
printRemesh := boolParam(opts.GetParams()["print_remesh"])
|
||||
alphaRatio := ratioOr(opts.GetParams()["alpha_ratio"], 0.005)
|
||||
offsetRatio := alphaRatio / 30
|
||||
if printRemesh && t2PrintRemeshAvailable() == 0 {
|
||||
return fmt.Errorf("print_remesh requested but libtrellis2 was built without CGAL Alpha Wrap")
|
||||
}
|
||||
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
if t.pipeline == 0 {
|
||||
return fmt.Errorf("model not loaded")
|
||||
}
|
||||
|
||||
errBuf := make([]byte, 512)
|
||||
r := t2Generate(t.pipeline, &img[0], int32(len(img)),
|
||||
pipelineForQuality(opts.Quality), backgroundForMode(opts.Background),
|
||||
seed, opts.Step, guidance, opts.TextureSteps,
|
||||
0, 0, 0, 0, // no progress/preview callbacks
|
||||
&errBuf[0], int32(len(errBuf)))
|
||||
if r == 0 {
|
||||
return fmt.Errorf("trellis2 generate: %s", cstr(errBuf))
|
||||
}
|
||||
defer t2MeshFree(r)
|
||||
|
||||
nv := t2MeshNVerts(r)
|
||||
nt := t2MeshNTris(r)
|
||||
if nv == 0 || nt == 0 {
|
||||
return fmt.Errorf("empty mesh")
|
||||
}
|
||||
var pbr *float32
|
||||
if t2MeshHasPBR(r) != 0 {
|
||||
pbr = t2MeshPBR(r)
|
||||
}
|
||||
|
||||
// Bake straight from the mesh accessor buffers — they stay valid until
|
||||
// t2_mesh_free, so no copies are needed.
|
||||
var outLen int32
|
||||
var glb *byte
|
||||
if printRemesh {
|
||||
wrap := t2PreparePrintMesh(t2MeshVerts(r), nv, t2MeshTris(r), nt, pbr,
|
||||
componentFilter, alphaRatio, offsetRatio,
|
||||
&errBuf[0], int32(len(errBuf)))
|
||||
if wrap == 0 {
|
||||
return fmt.Errorf("trellis2 print remesh: %s", cstr(errBuf))
|
||||
}
|
||||
defer t2MeshFree(wrap)
|
||||
wnv, wnt := t2MeshNVerts(wrap), t2MeshNTris(wrap)
|
||||
if wnv == 0 || wnt == 0 {
|
||||
return fmt.Errorf("empty print mesh")
|
||||
}
|
||||
if pbr != nil {
|
||||
// Wrapping creates new vertices, so the source material is
|
||||
// reprojected per texel onto the wrap's UV atlas (demo handleGLB).
|
||||
glb = t2BakeProjectedGLB(t2MeshVerts(wrap), wnv, t2MeshTris(wrap), wnt,
|
||||
t2MeshVerts(r), nv, t2MeshTris(r), nt, pbr,
|
||||
texSize, componentFilter,
|
||||
&outLen, &errBuf[0], int32(len(errBuf)))
|
||||
} else {
|
||||
glb = t2BakeGLB(t2MeshVerts(wrap), wnv, t2MeshTris(wrap), wnt, nil,
|
||||
texSize, 2, // the wrap output is already component-filtered
|
||||
&outLen, &errBuf[0], int32(len(errBuf)))
|
||||
}
|
||||
} else {
|
||||
glb = t2BakeGLB(t2MeshVerts(r), nv, t2MeshTris(r), nt, pbr,
|
||||
texSize, componentFilter,
|
||||
&outLen, &errBuf[0], int32(len(errBuf)))
|
||||
}
|
||||
if glb == nil {
|
||||
return fmt.Errorf("trellis2 GLB bake: %s", cstr(errBuf))
|
||||
}
|
||||
defer t2FreeBuffer(glb)
|
||||
|
||||
out := make([]byte, int(outLen))
|
||||
copy(out, unsafe.Slice(glb, int(outLen)))
|
||||
return os.WriteFile(opts.Dst, out, 0600)
|
||||
}
|
||||
|
||||
// remeshGLB applies the demo's post-generation print workflow to an existing
|
||||
// dense vertex-PBR GLB. It does not touch the inference pipeline: CGAL wrapping,
|
||||
// UV unwrapping, and PBR projection are CPU-only post-processing operations.
|
||||
func remeshGLB(opts *pb.Generate3DRequest) error {
|
||||
if opts.Src == "" {
|
||||
return fmt.Errorf("src is empty")
|
||||
}
|
||||
if t2PrintRemeshAvailable() == 0 {
|
||||
return fmt.Errorf("print remeshing is unavailable (libtrellis2 was built without CGAL Alpha Wrap)")
|
||||
}
|
||||
data, err := os.ReadFile(opts.Src)
|
||||
if err != nil {
|
||||
return fmt.Errorf("reading source GLB: %w", err)
|
||||
}
|
||||
mesh, err := parseVertexGLB(data)
|
||||
if err != nil {
|
||||
return fmt.Errorf("reading source GLB: %w", err)
|
||||
}
|
||||
|
||||
params := opts.GetParams()
|
||||
alphaRatio := ratioOr(params["alpha_ratio"], 0.005)
|
||||
offsetRatio := alphaRatio / 30
|
||||
componentFilter := componentFilterFor(params["components"])
|
||||
textureSize := atoiOr(params["texture_size"], 2048)
|
||||
var sourcePBR *float32
|
||||
if len(mesh.pbr) != 0 {
|
||||
sourcePBR = &mesh.pbr[0]
|
||||
}
|
||||
errBuf := make([]byte, 512)
|
||||
wrap := t2PreparePrintMesh(
|
||||
&mesh.verts[0], int32(len(mesh.verts)/3),
|
||||
&mesh.tris[0], int32(len(mesh.tris)/3),
|
||||
sourcePBR, componentFilter, alphaRatio, offsetRatio,
|
||||
&errBuf[0], int32(len(errBuf)),
|
||||
)
|
||||
if wrap == 0 {
|
||||
return fmt.Errorf("trellis2 print remesh: %s", cstr(errBuf))
|
||||
}
|
||||
defer t2MeshFree(wrap)
|
||||
|
||||
wrappedVerts, wrappedTris := t2MeshNVerts(wrap), t2MeshNTris(wrap)
|
||||
if wrappedVerts == 0 || wrappedTris == 0 {
|
||||
return fmt.Errorf("empty print mesh")
|
||||
}
|
||||
var outLen int32
|
||||
var glb *byte
|
||||
if sourcePBR != nil {
|
||||
glb = t2BakeProjectedGLB(
|
||||
t2MeshVerts(wrap), wrappedVerts, t2MeshTris(wrap), wrappedTris,
|
||||
&mesh.verts[0], int32(len(mesh.verts)/3),
|
||||
&mesh.tris[0], int32(len(mesh.tris)/3), sourcePBR,
|
||||
int32(textureSize), componentFilter,
|
||||
&outLen, &errBuf[0], int32(len(errBuf)),
|
||||
)
|
||||
} else {
|
||||
glb = t2BakeGLB(
|
||||
t2MeshVerts(wrap), wrappedVerts, t2MeshTris(wrap), wrappedTris,
|
||||
nil, int32(textureSize), 2,
|
||||
&outLen, &errBuf[0], int32(len(errBuf)),
|
||||
)
|
||||
}
|
||||
if glb == nil || outLen <= 0 {
|
||||
return fmt.Errorf("trellis2 GLB bake: %s", cstr(errBuf))
|
||||
}
|
||||
defer t2FreeBuffer(glb)
|
||||
|
||||
out := make([]byte, int(outLen))
|
||||
copy(out, unsafe.Slice(glb, int(outLen)))
|
||||
return os.WriteFile(opts.Dst, out, 0o600)
|
||||
}
|
||||
|
||||
func cstr(b []byte) string {
|
||||
for i, c := range b {
|
||||
if c == 0 {
|
||||
return string(b[:i])
|
||||
}
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
@@ -1,247 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/credentials/insecure"
|
||||
)
|
||||
|
||||
func TestTrellis2Cpp(t *testing.T) {
|
||||
RegisterFailHandler(Fail)
|
||||
RunSpecs(t, "trellis2cpp backend suite")
|
||||
}
|
||||
|
||||
// touch creates empty files — resolveModels only checks existence, so the
|
||||
// tests never need real GGUF weights.
|
||||
func touch(dir string, names ...string) {
|
||||
for _, name := range names {
|
||||
Expect(os.WriteFile(filepath.Join(dir, name), nil, 0o600)).To(Succeed())
|
||||
}
|
||||
}
|
||||
|
||||
var requiredFiles = []string{"dino_f16.gguf", "ss_flow_f16.gguf", "ss_dec_f16.gguf"}
|
||||
|
||||
var fullSet = append(append([]string{}, requiredFiles...),
|
||||
"slat_flow_f16.gguf", "slat_flow_1024_f16.gguf", "shape_dec_f16.gguf",
|
||||
"shape_enc_f16.gguf", "tex_dec_f16.gguf",
|
||||
"tex_slat_flow_512_f16.gguf", "tex_slat_flow_1024_f16.gguf")
|
||||
|
||||
var _ = Describe("resolveModels", func() {
|
||||
var dir string
|
||||
|
||||
BeforeEach(func() {
|
||||
dir = GinkgoT().TempDir()
|
||||
})
|
||||
|
||||
It("refuses a directory without the trellis2 component files", func() {
|
||||
touch(dir, "some-llm.gguf")
|
||||
|
||||
_, err := resolveModels("some-llm.gguf", dir, nil)
|
||||
Expect(err).To(MatchError(ContainSubstring("not a trellis2 model set")))
|
||||
})
|
||||
|
||||
It("resolves every role from the full default-named set", func() {
|
||||
touch(dir, fullSet...)
|
||||
|
||||
set, err := resolveModels("ss_flow_f16.gguf", dir, nil)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
for name, path := range map[string]string{
|
||||
"dino": set.dino,
|
||||
"ss_flow": set.ssFlow,
|
||||
"ss_dec": set.ssDec,
|
||||
"slat_flow": set.slatFlow,
|
||||
"slat_flow_1024": set.slatFlow1024,
|
||||
"shape_dec": set.shapeDec,
|
||||
"shape_enc": set.shapeEnc,
|
||||
"tex_dec": set.texDec,
|
||||
"tex_slat_flow_512": set.texSlatFlow512,
|
||||
"tex_slat_flow_1024": set.texSlatFlow1024,
|
||||
} {
|
||||
Expect(path).NotTo(BeEmpty(), "role %s", name)
|
||||
}
|
||||
})
|
||||
|
||||
It("degrades to coarse-only without the 512 pair, even when texture files exist", func() {
|
||||
touch(dir, requiredFiles...)
|
||||
touch(dir, "shape_enc_f16.gguf", "tex_dec_f16.gguf", "tex_slat_flow_512_f16.gguf", "slat_flow_1024_f16.gguf")
|
||||
|
||||
set, err := resolveModels("ss_flow_f16.gguf", dir, nil)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(set.slatFlow).To(BeEmpty())
|
||||
Expect(set.shapeDec).To(BeEmpty())
|
||||
Expect(set.slatFlow1024).To(BeEmpty())
|
||||
Expect(set.shapeEnc).To(BeEmpty())
|
||||
Expect(set.texDec).To(BeEmpty())
|
||||
Expect(set.texSlatFlow512).To(BeEmpty())
|
||||
Expect(set.texSlatFlow1024).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("disables texturing but keeps fine geometry when the texture set is incomplete", func() {
|
||||
touch(dir, requiredFiles...)
|
||||
touch(dir, "slat_flow_f16.gguf", "slat_flow_1024_f16.gguf", "shape_dec_f16.gguf", "tex_dec_f16.gguf")
|
||||
|
||||
set, err := resolveModels("ss_flow_f16.gguf", dir, nil)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(set.slatFlow).NotTo(BeEmpty())
|
||||
Expect(set.shapeDec).NotTo(BeEmpty())
|
||||
Expect(set.slatFlow1024).NotTo(BeEmpty())
|
||||
Expect(set.shapeEnc).To(BeEmpty())
|
||||
Expect(set.texDec).To(BeEmpty())
|
||||
Expect(set.texSlatFlow512).To(BeEmpty())
|
||||
Expect(set.texSlatFlow1024).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("drops the 1024 cascade when texturing lacks the HR texture flow", func() {
|
||||
touch(dir, fullSet...)
|
||||
Expect(os.Remove(filepath.Join(dir, "tex_slat_flow_1024_f16.gguf"))).To(Succeed())
|
||||
|
||||
set, err := resolveModels("ss_flow_f16.gguf", dir, nil)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(set.slatFlow1024).To(BeEmpty())
|
||||
Expect(set.shapeEnc).NotTo(BeEmpty())
|
||||
Expect(set.texDec).NotTo(BeEmpty())
|
||||
Expect(set.texSlatFlow512).NotTo(BeEmpty())
|
||||
})
|
||||
|
||||
It("honors explicit *_path option overrides", func() {
|
||||
touch(dir, fullSet...)
|
||||
custom := filepath.Join(dir, "custom")
|
||||
Expect(os.Mkdir(custom, 0o750)).To(Succeed())
|
||||
touch(custom, "my-dino.gguf")
|
||||
|
||||
set, err := resolveModels("ss_flow_f16.gguf", dir, []string{"dino_path:custom/my-dino.gguf"})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(set.dino).To(Equal(filepath.Join(custom, "my-dino.gguf")))
|
||||
})
|
||||
|
||||
It("fails when an explicitly configured file is missing", func() {
|
||||
touch(dir, fullSet...)
|
||||
|
||||
_, err := resolveModels("ss_flow_f16.gguf", dir, []string{"tex_dec_path:nope.gguf"})
|
||||
Expect(err).To(MatchError(ContainSubstring("missing file")))
|
||||
})
|
||||
|
||||
It("rejects option paths escaping the model directory", func() {
|
||||
touch(dir, fullSet...)
|
||||
|
||||
_, err := resolveModels("ss_flow_f16.gguf", dir, []string{"dino_path:../outside.gguf"})
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = DescribeTable("request parameter mapping",
|
||||
func(got, want int32) {
|
||||
Expect(got).To(Equal(want))
|
||||
},
|
||||
Entry("quality empty", pipelineForQuality(""), int32(pipeAuto)),
|
||||
Entry("quality auto", pipelineForQuality("auto"), int32(pipeAuto)),
|
||||
Entry("quality coarse", pipelineForQuality("coarse"), int32(pipeCoarse)),
|
||||
Entry("quality 512", pipelineForQuality("512"), int32(pipe512)),
|
||||
Entry("quality 1024", pipelineForQuality("1024"), int32(pipe1024)),
|
||||
Entry("background empty", backgroundForMode(""), int32(backgroundAuto)),
|
||||
Entry("background auto", backgroundForMode("auto"), int32(backgroundAuto)),
|
||||
Entry("background keep", backgroundForMode("keep"), int32(backgroundKeep)),
|
||||
Entry("background black", backgroundForMode("black"), int32(backgroundBlack)),
|
||||
Entry("background white", backgroundForMode("white"), int32(backgroundWhite)),
|
||||
Entry("components default", componentFilterFor(""), int32(2)),
|
||||
Entry("components all", componentFilterFor("all"), int32(2)),
|
||||
Entry("components largest", componentFilterFor("largest"), int32(1)),
|
||||
Entry("components tiny", componentFilterFor("tiny"), int32(0)),
|
||||
Entry("atoi empty", atoiOr("", 0), int32(0)),
|
||||
Entry("atoi value", atoiOr("2048", 0), int32(2048)),
|
||||
Entry("atoi junk", atoiOr("junk", 7), int32(7)),
|
||||
)
|
||||
|
||||
var _ = Describe("print remesh parameters", func() {
|
||||
It("parses the print_remesh toggle", func() {
|
||||
Expect(boolParam("1")).To(BeTrue())
|
||||
Expect(boolParam("true")).To(BeTrue())
|
||||
Expect(boolParam("TRUE")).To(BeTrue())
|
||||
Expect(boolParam("")).To(BeFalse())
|
||||
Expect(boolParam("0")).To(BeFalse())
|
||||
Expect(boolParam("no")).To(BeFalse())
|
||||
})
|
||||
|
||||
It("parses ratios and clamps junk to the fallback", func() {
|
||||
Expect(ratioOr("", 0.005)).To(BeNumerically("~", 0.005, 1e-6))
|
||||
Expect(ratioOr("0.01", 0.005)).To(BeNumerically("~", 0.01, 1e-6))
|
||||
Expect(ratioOr("junk", 0.005)).To(BeNumerically("~", 0.005, 1e-6))
|
||||
Expect(ratioOr("-1", 0.005)).To(BeNumerically("~", 0.005, 1e-6))
|
||||
Expect(ratioOr("0.9", 0.005)).To(BeNumerically("~", 0.005, 1e-6), "above the demo's 50% cap")
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("packaged backend", func() {
|
||||
It("starts and answers Health without loading model weights", func() {
|
||||
runScript := os.Getenv("TRELLIS2CPP_SMOKE_RUN")
|
||||
if runScript == "" {
|
||||
runScript = filepath.Join("package", "run.sh")
|
||||
}
|
||||
if _, err := os.Stat(runScript); os.IsNotExist(err) {
|
||||
Skip("packaged backend is not present; run make before the smoke test")
|
||||
}
|
||||
|
||||
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
addr := listener.Addr().String()
|
||||
Expect(listener.Close()).To(Succeed())
|
||||
|
||||
cmd := exec.Command("bash", runScript, "--addr="+addr)
|
||||
cmd.Stdout = GinkgoWriter
|
||||
cmd.Stderr = GinkgoWriter
|
||||
Expect(cmd.Start()).To(Succeed())
|
||||
processDone := make(chan error, 1)
|
||||
go func() { processDone <- cmd.Wait() }()
|
||||
processExited := false
|
||||
DeferCleanup(func() {
|
||||
if cmd.Process != nil && !processExited {
|
||||
_ = cmd.Process.Kill()
|
||||
<-processDone
|
||||
}
|
||||
})
|
||||
|
||||
Eventually(func() error {
|
||||
select {
|
||||
case err := <-processDone:
|
||||
processExited = true
|
||||
if err != nil {
|
||||
return StopTrying("backend exited before Health succeeded").Wrap(err)
|
||||
}
|
||||
return StopTrying("backend exited before Health succeeded")
|
||||
default:
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
|
||||
conn, err := grpc.DialContext(ctx, addr,
|
||||
grpc.WithTransportCredentials(insecure.NewCredentials()),
|
||||
grpc.WithBlock(),
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = conn.Close() }()
|
||||
|
||||
reply, err := pb.NewBackendClient(conn).Health(ctx, &pb.HealthMessage{})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if string(reply.GetMessage()) != "OK" {
|
||||
return fmt.Errorf("unexpected Health reply %q", reply.GetMessage())
|
||||
}
|
||||
return nil
|
||||
}, 30*time.Second, 200*time.Millisecond).Should(Succeed())
|
||||
})
|
||||
})
|
||||
@@ -1,17 +0,0 @@
|
||||
GOCMD=go
|
||||
|
||||
valkey-store:
|
||||
CGO_ENABLED=0 $(GOCMD) build -ldflags "$(LD_FLAGS)" -tags "$(GO_TAGS)" -o valkey-store ./
|
||||
|
||||
package:
|
||||
bash package.sh
|
||||
|
||||
build: valkey-store package
|
||||
|
||||
## Runs the backend's Ginkgo suite. The unit (mock) specs run without a
|
||||
## container; the integration specs skip automatically unless VALKEY_ADDR is set.
|
||||
test:
|
||||
$(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --flake-attempts $(TEST_FLAKES) -v -r ./
|
||||
|
||||
clean:
|
||||
rm -f valkey-store
|
||||
@@ -1,261 +0,0 @@
|
||||
package main
|
||||
|
||||
// Connection + index configuration for the Valkey-backed vector store.
|
||||
//
|
||||
// Configuration is read from the model config `options:` list (a repeated
|
||||
// `key:value` string carried over gRPC in ModelOptions.Options) rather than
|
||||
// from process-wide environment variables. Driving it from the model config is
|
||||
// the LocalAI convention and, crucially, lets multiple stores each have their
|
||||
// own Valkey config (a face registry on one server, a router cache on another)
|
||||
// within a single LocalAI process — something a single VALKEY_* env surface
|
||||
// could never express. Every default lives as a named constant below — no
|
||||
// magic literals sprinkled through the store logic — so the defaults can be
|
||||
// audited in one place and referenced by the unit tests.
|
||||
//
|
||||
// Example model YAML:
|
||||
//
|
||||
// name: my-vector-store
|
||||
// backend: valkey-store
|
||||
// options:
|
||||
// - addr:valkey.internal:6379
|
||||
// - index_algo:HNSW
|
||||
// - distance_metric:COSINE
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
const (
|
||||
// _defaultAddr is the single-node Valkey address used when the `addr`
|
||||
// option is unset. Matches the port Phase 0 reserved for integration tests.
|
||||
_defaultAddr = "localhost:6379"
|
||||
|
||||
// _defaultClientName is mandatory: every connection identifies itself with
|
||||
// this name so operators can spot LocalAI's traffic via CLIENT LIST. It is
|
||||
// always set on the client, even if the operator clears the client_name option.
|
||||
_defaultClientName = "localai-valkey-store"
|
||||
|
||||
// _defaultIndexAlgo is FLAT (exact brute-force KNN) to preserve parity with
|
||||
// local-store's linear scan and keep the exact-cosine test expectations.
|
||||
_defaultIndexAlgo = indexAlgoFlat
|
||||
|
||||
// _defaultDistanceMetric is COSINE so similarities match local-store
|
||||
// (sim = 1 - cosine_distance). L2/IP are opt-in.
|
||||
_defaultDistanceMetric = distanceCosine
|
||||
|
||||
// HNSW graph defaults (only used when index_algo=HNSW). Values follow
|
||||
// the Valkey Search documented defaults.
|
||||
_defaultHNSWM = 16
|
||||
_defaultHNSWEFConstruction = 200
|
||||
_defaultHNSWEFRuntime = 10
|
||||
|
||||
// _defaultRequestTimeoutMS bounds every command. We deliberately do NOT rely
|
||||
// on the client's built-in write timeout: index back-fill or a slow KNN can
|
||||
// exceed a short default, so we thread this explicit deadline into every
|
||||
// command context.
|
||||
_defaultRequestTimeoutMS = 5000
|
||||
|
||||
// Valkey Search index algorithms.
|
||||
indexAlgoFlat = "FLAT"
|
||||
indexAlgoHNSW = "HNSW"
|
||||
|
||||
// Supported distance metrics.
|
||||
distanceCosine = "COSINE"
|
||||
distanceL2 = "L2"
|
||||
distanceIP = "IP"
|
||||
|
||||
// Option keys recognised in the model config `options:` list. They mirror
|
||||
// the previous VALKEY_* env var names without the prefix and lower-cased, so
|
||||
// operators migrating a config have an obvious 1:1 mapping.
|
||||
optAddr = "addr"
|
||||
optUsername = "username"
|
||||
optPassword = "password"
|
||||
optUsernameEnv = "username_env"
|
||||
optPasswordEnv = "password_env"
|
||||
optTLS = "tls"
|
||||
optTLSSkipVerify = "tls_skip_verify"
|
||||
optTLSCACert = "tls_ca_cert"
|
||||
optClientName = "client_name"
|
||||
optDB = "db"
|
||||
optIndexAlgo = "index_algo"
|
||||
optDistanceMetric = "distance_metric"
|
||||
optHNSWM = "hnsw_m"
|
||||
optHNSWEFConstruction = "hnsw_ef_construction"
|
||||
optHNSWEFRuntime = "hnsw_ef_runtime"
|
||||
optRequestTimeoutMS = "request_timeout_ms"
|
||||
)
|
||||
|
||||
// hnswParams holds the HNSW-only tuning knobs. They are ignored unless
|
||||
// IndexAlgo == indexAlgoHNSW.
|
||||
type hnswParams struct {
|
||||
M int
|
||||
EFConstruction int
|
||||
EFRuntime int
|
||||
}
|
||||
|
||||
// Config is the fully-resolved store configuration produced by loadConfig().
|
||||
type Config struct {
|
||||
Addr string
|
||||
Username string
|
||||
Password string
|
||||
UseTLS bool
|
||||
TLSSkipVerify bool
|
||||
TLSCACert string
|
||||
ClientName string
|
||||
DB int
|
||||
IndexAlgo string
|
||||
DistanceMetric string
|
||||
HNSW hnswParams
|
||||
RequestTimeout time.Duration
|
||||
}
|
||||
|
||||
// parseOptions turns the repeated `key:value` ModelOptions.Options list into a
|
||||
// lookup map. The split is on the FIRST ':' via strings.Cut, so values that
|
||||
// themselves contain a colon (e.g. `addr:host:6379`) are preserved intact. A
|
||||
// malformed entry with no ':' is warned about and skipped rather than silently
|
||||
// dropped, so an operator typo is visible in the logs.
|
||||
func parseOptions(opts *pb.ModelOptions) map[string]string {
|
||||
m := make(map[string]string)
|
||||
if opts == nil {
|
||||
return m
|
||||
}
|
||||
for _, o := range opts.GetOptions() {
|
||||
k, v, ok := strings.Cut(o, ":")
|
||||
if !ok {
|
||||
xlog.Warn("valkey-store: ignoring malformed option (want key:value)", "option", o)
|
||||
continue
|
||||
}
|
||||
m[strings.ToLower(strings.TrimSpace(k))] = strings.TrimSpace(v)
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// loadConfig resolves the store configuration from the model config options and
|
||||
// returns a validated Config. It fails fast on an unknown index algorithm or
|
||||
// distance metric (and on a malformed integer) so a misconfiguration surfaces
|
||||
// at Load() rather than silently degrading search.
|
||||
func loadConfig(opts *pb.ModelOptions) (Config, error) {
|
||||
o := parseOptions(opts)
|
||||
|
||||
// intOr parses an integer option, failing fast on a malformed value the same
|
||||
// way an invalid index algo or distance metric does. A typo like
|
||||
// `hnsw_m:1x6` must surface at Load() rather than silently degrading to the
|
||||
// default and producing subtly wrong (and hard-to-diagnose) index
|
||||
// behaviour. The first parse error wins and is returned below.
|
||||
var parseErr error
|
||||
intOr := func(key string, fallback int) int {
|
||||
v, ok := o[key]
|
||||
if !ok || v == "" {
|
||||
return fallback
|
||||
}
|
||||
n, err := strconv.Atoi(v)
|
||||
if err != nil {
|
||||
if parseErr == nil {
|
||||
parseErr = fmt.Errorf("valkey-store: invalid option %s %q: %w", key, v, err)
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
cfg := Config{
|
||||
Addr: strOr(o, optAddr, _defaultAddr),
|
||||
Username: resolveCredential(o, optUsername, optUsernameEnv),
|
||||
Password: resolveCredential(o, optPassword, optPasswordEnv),
|
||||
UseTLS: boolOr(o, optTLS, false),
|
||||
TLSSkipVerify: boolOr(o, optTLSSkipVerify, false),
|
||||
TLSCACert: o[optTLSCACert],
|
||||
ClientName: strOr(o, optClientName, _defaultClientName),
|
||||
DB: intOr(optDB, 0),
|
||||
IndexAlgo: strings.ToUpper(strOr(o, optIndexAlgo, _defaultIndexAlgo)),
|
||||
DistanceMetric: strings.ToUpper(strOr(o, optDistanceMetric, _defaultDistanceMetric)),
|
||||
HNSW: hnswParams{
|
||||
M: intOr(optHNSWM, _defaultHNSWM),
|
||||
EFConstruction: intOr(optHNSWEFConstruction, _defaultHNSWEFConstruction),
|
||||
EFRuntime: intOr(optHNSWEFRuntime, _defaultHNSWEFRuntime),
|
||||
},
|
||||
RequestTimeout: time.Duration(intOr(optRequestTimeoutMS, _defaultRequestTimeoutMS)) * time.Millisecond,
|
||||
}
|
||||
if parseErr != nil {
|
||||
return Config{}, parseErr
|
||||
}
|
||||
|
||||
// ClientName is mandatory. Restore the default if the operator blanked it,
|
||||
// so the connection is always identifiable.
|
||||
if cfg.ClientName == "" {
|
||||
cfg.ClientName = _defaultClientName
|
||||
}
|
||||
|
||||
if cfg.DB < 0 {
|
||||
return Config{}, fmt.Errorf("valkey-store: invalid option %s %d (must be >= 0)", optDB, cfg.DB)
|
||||
}
|
||||
|
||||
switch cfg.IndexAlgo {
|
||||
case indexAlgoFlat, indexAlgoHNSW:
|
||||
default:
|
||||
return Config{}, fmt.Errorf("valkey-store: invalid option %s %q (want FLAT or HNSW)", optIndexAlgo, cfg.IndexAlgo)
|
||||
}
|
||||
|
||||
switch cfg.DistanceMetric {
|
||||
case distanceCosine, distanceL2, distanceIP:
|
||||
default:
|
||||
return Config{}, fmt.Errorf("valkey-store: invalid option %s %q (want COSINE, L2 or IP)", optDistanceMetric, cfg.DistanceMetric)
|
||||
}
|
||||
|
||||
if cfg.RequestTimeout <= 0 {
|
||||
cfg.RequestTimeout = time.Duration(_defaultRequestTimeoutMS) * time.Millisecond
|
||||
}
|
||||
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
// strOr returns the option value for key, or fallback when it is unset/empty.
|
||||
func strOr(o map[string]string, key, fallback string) string {
|
||||
if v, ok := o[key]; ok && v != "" {
|
||||
return v
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
// boolOr parses a boolean option, falling back to the default on an unset or
|
||||
// unparseable value. A typo is surfaced via a warning (like the previous env
|
||||
// behaviour) rather than failing Load for a coarse on/off switch.
|
||||
func boolOr(o map[string]string, key string, fallback bool) bool {
|
||||
v, ok := o[key]
|
||||
if !ok || v == "" {
|
||||
return fallback
|
||||
}
|
||||
b, err := strconv.ParseBool(v)
|
||||
if err != nil {
|
||||
xlog.Warn("valkey-store: ignoring unparseable option, using default", "key", key, "value", v, "default", fallback)
|
||||
return fallback
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// resolveCredential resolves a credential value with the following priority:
|
||||
// 1. Direct value from the model config option (e.g. `username:admin`)
|
||||
// 2. Env-indirection: if `username_env` names an env var, read the credential
|
||||
// from that variable (e.g. `username_env:MY_VALKEY_USER` → os.Getenv("MY_VALKEY_USER"))
|
||||
//
|
||||
// The env-indirection pattern (same as cloud-proxy's api_key_env) avoids putting
|
||||
// secrets directly in model YAML: distinct store configs can each reference a
|
||||
// different credential env var without any plaintext passwords in the config.
|
||||
func resolveCredential(o map[string]string, directKey, envKey string) string {
|
||||
// Direct value takes precedence (backward compatible).
|
||||
if v := o[directKey]; v != "" {
|
||||
return v
|
||||
}
|
||||
// Env indirection: the option names an env var that holds the credential.
|
||||
if envVar := o[envKey]; envVar != "" {
|
||||
return os.Getenv(envVar)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -1,80 +0,0 @@
|
||||
package main
|
||||
|
||||
// Vector⇄key encoding: the "vector IS the key" resolution.
|
||||
//
|
||||
// local-store keys entries *by* the vector itself (a []float32). Valkey hashes
|
||||
// are keyed by strings, so we synthesise a deterministic, lossless key:
|
||||
//
|
||||
// key = prefix + hex(little-endian float32 bytes of the vector)
|
||||
//
|
||||
// The same vector always produces the same bytes, so HSET is an upsert and
|
||||
// HGET/DEL are exact matches — and the encoding is reversible, so we can hand
|
||||
// the original []float32 back on Get/Find.
|
||||
//
|
||||
// Divergence from local-store (documented and tested): local-store compares
|
||||
// keys with slices.Compare, which treats -0.0 == +0.0 and orders NaN, so those
|
||||
// collapse to the same logical key. Byte-encoding makes -0.0 and +0.0 (and any
|
||||
// distinct NaN bit-pattern) *distinct* keys. We accept this on purpose: a
|
||||
// lossless, deterministic, exact round-trip is more valuable for a persistent
|
||||
// store than reproducing local-store's float-equality quirk, and callers never
|
||||
// rely on -0.0/+0.0 aliasing.
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"math"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// _float32Bytes is the wire width of a single FLOAT32 component.
|
||||
const _float32Bytes = 4
|
||||
|
||||
// vecToBytes encodes a vector as little-endian float32 bytes. This is byte-for
|
||||
// -byte identical to valkey.VectorString32, so the value we store in the hash
|
||||
// `vec` field and the bytes we hash into the key share one encoding.
|
||||
func vecToBytes(v []float32) []byte {
|
||||
b := make([]byte, len(v)*_float32Bytes)
|
||||
for i, e := range v {
|
||||
off := i * _float32Bytes
|
||||
binary.LittleEndian.PutUint32(b[off:off+_float32Bytes], math.Float32bits(e))
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// bytesToVec reverses vecToBytes. It rejects a payload whose length is not a
|
||||
// multiple of the float32 width, which would indicate a corrupted/foreign value.
|
||||
func bytesToVec(b []byte) ([]float32, error) {
|
||||
if len(b)%_float32Bytes != 0 {
|
||||
return nil, fmt.Errorf("valkey-store: vector byte length %d is not a multiple of %d", len(b), _float32Bytes)
|
||||
}
|
||||
v := make([]float32, len(b)/_float32Bytes)
|
||||
for i := range v {
|
||||
off := i * _float32Bytes
|
||||
v[i] = math.Float32frombits(binary.LittleEndian.Uint32(b[off : off+_float32Bytes]))
|
||||
}
|
||||
return v, nil
|
||||
}
|
||||
|
||||
// encodeKey builds the Valkey hash key for a vector: prefix + hex(bytes).
|
||||
// Hex keeps the key printable (so it is safe in FT.CREATE PREFIX and in logs)
|
||||
// while staying lossless.
|
||||
func encodeKey(prefix string, v []float32) string {
|
||||
return prefix + hex.EncodeToString(vecToBytes(v))
|
||||
}
|
||||
|
||||
// decodeKey reverses encodeKey. It is intentionally retained as the tested,
|
||||
// symmetric inverse of encodeKey — it is NOT on the hot Find path (StoresFind
|
||||
// decodes the returned `vec` bytes via bytesToVec directly), but keeping the
|
||||
// key↔vector mapping provably invertible guards the encoding contract and is
|
||||
// exercised by the round-trip unit tests.
|
||||
func decodeKey(prefix, key string) ([]float32, error) {
|
||||
if !strings.HasPrefix(key, prefix) {
|
||||
return nil, fmt.Errorf("valkey-store: key %q does not have expected prefix %q", key, prefix)
|
||||
}
|
||||
b, err := hex.DecodeString(strings.TrimPrefix(key, prefix))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("valkey-store: decode key hex: %w", err)
|
||||
}
|
||||
return bytesToVec(b)
|
||||
}
|
||||
@@ -1,77 +0,0 @@
|
||||
package main
|
||||
|
||||
// Unit tests for the vector⇄key encoding. These need no Valkey server: they
|
||||
// exercise the pure lossless-encoding contract that the whole store relies on,
|
||||
// including the documented edge cases (-0.0/+0.0 and NaN) where this encoding
|
||||
// intentionally diverges from local-store's slices.Compare float equality.
|
||||
|
||||
import (
|
||||
"math"
|
||||
"math/rand/v2"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
valkey "github.com/valkey-io/valkey-go"
|
||||
)
|
||||
|
||||
var _ = Describe("vector⇄bytes encoding", func() {
|
||||
It("round-trips vectors of varying dimensions", func() {
|
||||
r := rand.New(rand.NewPCG(1, 2))
|
||||
for _, dim := range []int{1, 3, 4, 16, 128, 768} {
|
||||
v := make([]float32, dim)
|
||||
for i := range v {
|
||||
v[i] = float32(r.NormFloat64())
|
||||
}
|
||||
got, err := bytesToVec(vecToBytes(v))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(got).To(Equal(v))
|
||||
}
|
||||
})
|
||||
|
||||
It("matches valkey.VectorString32 byte-for-byte", func() {
|
||||
// The stored `vec` field uses valkey.VectorString32; the key uses
|
||||
// vecToBytes. They must be the same encoding or Get/Find break.
|
||||
v := []float32{0.1, -0.2, 3.5, 0}
|
||||
Expect(valkey.BinaryString(vecToBytes(v))).To(Equal(valkey.VectorString32(v)))
|
||||
})
|
||||
|
||||
It("rejects a byte payload that is not a multiple of 4", func() {
|
||||
_, err := bytesToVec([]byte{1, 2, 3})
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("key encoding", func() {
|
||||
const prefix = "vs:test:"
|
||||
|
||||
It("round-trips key encode/decode", func() {
|
||||
v := []float32{0.5, 0.5, 0.5}
|
||||
key := encodeKey(prefix, v)
|
||||
Expect(key).To(HavePrefix(prefix))
|
||||
got, err := decodeKey(prefix, key)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(got).To(Equal(v))
|
||||
})
|
||||
|
||||
It("produces distinct keys for -0.0 and +0.0 (documented divergence)", func() {
|
||||
negZero := float32(math.Copysign(0, -1))
|
||||
posZero := float32(0)
|
||||
Expect(math.Signbit(float64(negZero))).To(BeTrue())
|
||||
Expect(encodeKey(prefix, []float32{negZero})).NotTo(Equal(encodeKey(prefix, []float32{posZero})))
|
||||
})
|
||||
|
||||
It("produces a stable, distinct key for a NaN component", func() {
|
||||
nan := float32(math.NaN())
|
||||
k1 := encodeKey(prefix, []float32{nan})
|
||||
k2 := encodeKey(prefix, []float32{nan})
|
||||
// Deterministic: same NaN bit-pattern → same key.
|
||||
Expect(k1).To(Equal(k2))
|
||||
// Distinct from a normal value.
|
||||
Expect(k1).NotTo(Equal(encodeKey(prefix, []float32{0})))
|
||||
})
|
||||
|
||||
It("rejects a key without the expected prefix", func() {
|
||||
_, err := decodeKey(prefix, "wrong:deadbeef")
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
})
|
||||
@@ -1,25 +0,0 @@
|
||||
package main
|
||||
|
||||
// Note: this is started internally by LocalAI and a server is allocated for each store
|
||||
|
||||
import (
|
||||
"flag"
|
||||
"os"
|
||||
|
||||
grpc "github.com/mudler/LocalAI/pkg/grpc"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
var (
|
||||
addr = flag.String("addr", "localhost:50051", "the address to connect to")
|
||||
)
|
||||
|
||||
func main() {
|
||||
xlog.SetLogger(xlog.NewLogger(xlog.LogLevel(os.Getenv("LOCALAI_LOG_LEVEL")), os.Getenv("LOCALAI_LOG_FORMAT")))
|
||||
|
||||
flag.Parse()
|
||||
|
||||
if err := grpc.StartServer(*addr, NewValkeyStore()); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
@@ -1,12 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Script to copy the appropriate libraries based on architecture
|
||||
# This script is used in the final stage of the Dockerfile
|
||||
|
||||
set -e
|
||||
|
||||
CURDIR=$(dirname "$(realpath $0)")
|
||||
|
||||
mkdir -p $CURDIR/package
|
||||
cp -avf $CURDIR/valkey-store $CURDIR/package/
|
||||
cp -rfv $CURDIR/run.sh $CURDIR/package/
|
||||
@@ -1,6 +0,0 @@
|
||||
#!/bin/bash
|
||||
set -ex
|
||||
|
||||
CURDIR=$(dirname "$(realpath "$0")")
|
||||
|
||||
exec "$CURDIR"/valkey-store "$@"
|
||||
@@ -1,692 +0,0 @@
|
||||
package main
|
||||
|
||||
// Valkey-backed vector store, exposed as a gRPC backend. It mirrors the public
|
||||
// contract of backend/go/local-store (the four Stores* RPCs + Load) but swaps
|
||||
// the in-memory sorted slices for Valkey Search (FT.*) so the data persists
|
||||
// across restarts and can scale beyond an O(N) scan (opt-in HNSW).
|
||||
//
|
||||
// Data model — each entry is a Valkey HASH keyed by
|
||||
//
|
||||
// prefix + hex(little-endian float32 bytes of the vector)
|
||||
//
|
||||
// with two fields: `vec` (the raw float32 bytes, indexed by a lazily-created
|
||||
// FT VECTOR index of the discovered dimension) and `val` (the opaque value
|
||||
// bytes). The vector-IS-the-key encoding (see encoding.go) makes Set an
|
||||
// HSET upsert, Get an HGET, Delete a DEL, and Find an FT.SEARCH KNN.
|
||||
//
|
||||
// Similarity — Valkey returns cosine *distance* (0 = identical, 2 = opposite),
|
||||
// while local-store returns cosine *similarity* (1 = identical, -1 = opposite).
|
||||
// We convert sim = 1 - distance for COSINE so the values match local-store's
|
||||
// integration expectations exactly. For L2/IP the raw score is passed through.
|
||||
//
|
||||
// Concurrency — base.SingleThread serialises gRPC calls, so the store's
|
||||
// scalar bookkeeping (keyLen, indexCreated) needs no extra locking. All Valkey
|
||||
// commands are synchronous via client.Do and bounded by an explicit
|
||||
// per-request deadline (cfg.RequestTimeout); there is no background event loop.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/mudler/LocalAI/pkg/grpc/base"
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/LocalAI/pkg/store"
|
||||
"github.com/mudler/xlog"
|
||||
valkey "github.com/valkey-io/valkey-go"
|
||||
)
|
||||
|
||||
const (
|
||||
// Hash field names. `vec` is the indexed vector; `val` is the opaque value.
|
||||
_vecField = "vec"
|
||||
_valField = "val"
|
||||
// _scoreField is the KNN distance alias produced by the query and returned
|
||||
// by FT.SEARCH. Double-underscore avoids colliding with a stored field.
|
||||
_scoreField = "__score"
|
||||
|
||||
// _keyPrefixPrefix / _indexPrefix namespace the keys and index per model so
|
||||
// two namespaces (e.g. a 512-d face store and a 192-d voice store) sharing
|
||||
// one Valkey server never collide.
|
||||
_keyPrefixPrefix = "vs:"
|
||||
_indexPrefix = "idx:"
|
||||
|
||||
// _maxTopK bounds a Find so an accidental or abusive huge TopK cannot force
|
||||
// an unbounded server-side LIMIT / allocation. local-store has no cap, but
|
||||
// it is in-memory; a networked backend wants a guard. Callers asking for
|
||||
// more than this get the top _maxTopK results.
|
||||
_maxTopK = 10000
|
||||
|
||||
// _maxNsTokenLen bounds the human-readable portion of a namespace token so
|
||||
// a very long model name cannot produce an unbounded key prefix / index
|
||||
// name. The appended short hash keeps distinct namespaces collision-free
|
||||
// even when their sanitized prefixes are truncated to the same value.
|
||||
_maxNsTokenLen = 64
|
||||
)
|
||||
|
||||
// ValkeyStore implements the gRPC store Backend against Valkey Search.
|
||||
type ValkeyStore struct {
|
||||
base.SingleThread
|
||||
|
||||
client valkey.Client
|
||||
cfg Config
|
||||
|
||||
// prefix is the per-namespace key prefix; indexName is the FT index name.
|
||||
prefix string
|
||||
indexName string
|
||||
|
||||
// keyLen is the vector dimension, learned from the first Set. -1 means
|
||||
// "no keys yet" — mirrors local-store so dimension-mismatch errors are
|
||||
// identical. indexCreated tracks whether FT.CREATE has run (lazy creation).
|
||||
keyLen int
|
||||
indexCreated bool
|
||||
}
|
||||
|
||||
// NewValkeyStore returns a store with an open dimension and no index yet. The
|
||||
// Valkey client is established in Load once the connection config is known.
|
||||
func NewValkeyStore() *ValkeyStore {
|
||||
return &ValkeyStore{keyLen: -1}
|
||||
}
|
||||
|
||||
// newWithClient builds a store around an already-constructed client for a given
|
||||
// namespace. It exists so unit tests can inject a mock client without a real
|
||||
// Valkey server; Load is the production path.
|
||||
func newWithClient(client valkey.Client, cfg Config, namespace string) *ValkeyStore {
|
||||
return &ValkeyStore{
|
||||
client: client,
|
||||
cfg: cfg,
|
||||
prefix: keyPrefix(namespace),
|
||||
indexName: indexName(namespace),
|
||||
keyLen: -1,
|
||||
}
|
||||
}
|
||||
|
||||
// Load reads the store config from the model config options, connects, and
|
||||
// verifies the connection. The mandatory ClientName is always set so the
|
||||
// connection is identifiable via CLIENT LIST. opts.Model is the namespace
|
||||
// identifier (one process per (backend, model) tuple upstream), so we derive
|
||||
// an isolated key prefix and index name from it, and opts.Options carries the
|
||||
// per-store connection/index configuration.
|
||||
//
|
||||
// The NamespacePrefix gate mirrors local-store: core's StoreBackend always
|
||||
// sends the model name with store.NamespacePrefix; anything else is the model
|
||||
// loader's greedy autoload probing with a real model name, which must be
|
||||
// refused or the LLM binds to the vector store (the #9287 failure mode).
|
||||
func (s *ValkeyStore) Load(opts *pb.ModelOptions) error {
|
||||
if opts == nil {
|
||||
return fmt.Errorf("valkey-store: refusing to load: nil model options (expected %q prefix)", store.NamespacePrefix)
|
||||
}
|
||||
if !strings.HasPrefix(opts.GetModel(), store.NamespacePrefix) {
|
||||
return fmt.Errorf("valkey-store: refusing to load %q: not a store namespace (expected %q prefix)", opts.GetModel(), store.NamespacePrefix)
|
||||
}
|
||||
|
||||
cfg, err := loadConfig(opts)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.cfg = cfg
|
||||
|
||||
namespace := opts.Model
|
||||
s.prefix = keyPrefix(namespace)
|
||||
s.indexName = indexName(namespace)
|
||||
|
||||
clientOpt := valkey.ClientOption{
|
||||
InitAddress: []string{cfg.Addr},
|
||||
Username: cfg.Username,
|
||||
Password: cfg.Password,
|
||||
ClientName: cfg.ClientName,
|
||||
// SelectDB picks a logical Valkey DB (SELECT n) for deployments that use
|
||||
// numbered DBs for isolation. Defaults to 0; namespace prefixing already
|
||||
// isolates keyspaces on a shared DB.
|
||||
SelectDB: cfg.DB,
|
||||
// Disable client-side caching: values are opaque blobs written once and
|
||||
// read rarely, so tracking invalidations would only add overhead.
|
||||
DisableCache: true,
|
||||
}
|
||||
if cfg.UseTLS {
|
||||
tlsCfg, err := buildTLSConfig(cfg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
clientOpt.TLSConfig = tlsCfg
|
||||
}
|
||||
|
||||
// Close any client from a previous Load so a re-entrant Load does not leak
|
||||
// the old connection. Not reachable in the one-process-per-namespace model
|
||||
// today, but keeps Load idempotent.
|
||||
if s.client != nil {
|
||||
s.client.Close()
|
||||
s.client = nil
|
||||
}
|
||||
|
||||
client, err := valkey.NewClient(clientOpt)
|
||||
if err != nil {
|
||||
return fmt.Errorf("valkey-store: connect to %s: %w", cfg.Addr, err)
|
||||
}
|
||||
s.client = client
|
||||
|
||||
// Fail fast if the server is unreachable, mirroring how a real vector DB
|
||||
// backend would refuse to load against a dead endpoint.
|
||||
ctx, cancel := s.ctx()
|
||||
defer cancel()
|
||||
if err := s.client.Do(ctx, s.client.B().Ping().Build()).Error(); err != nil {
|
||||
s.client.Close()
|
||||
s.client = nil
|
||||
return fmt.Errorf("valkey-store: ping %s: %w", cfg.Addr, err)
|
||||
}
|
||||
|
||||
// A durable Valkey may already hold this namespace's index from a previous
|
||||
// run (this is the persistence capability local-store lacks). Recover both
|
||||
// its existence AND its vector dimension so Find works before this fresh
|
||||
// process issues its first Set, and — critically — so a post-restart Set
|
||||
// validates the incoming dimension against the real persisted DIM instead
|
||||
// of silently re-learning a wrong one and dropping mismatched vectors from
|
||||
// the index (which would return success while making the entry unsearchable).
|
||||
s.loadIndexState(ctx)
|
||||
|
||||
// Log the sanitized index name (which identifies the namespace) rather than
|
||||
// the raw model-derived namespace, which could carry control characters.
|
||||
xlog.Info("valkey-store loaded", "addr", cfg.Addr, "index", s.indexName, "algo", cfg.IndexAlgo, "metric", cfg.DistanceMetric, "indexExists", s.indexCreated, "keyLen", s.keyLen)
|
||||
return nil
|
||||
}
|
||||
|
||||
// loadIndexState issues one FT.INFO at Load to recover the persisted index
|
||||
// state. FT.INFO returns an error for an unknown index, so a successful reply
|
||||
// means the index exists (indexCreated=true). We then recover the vector
|
||||
// dimension from the reply and seed keyLen with it: without this, keyLen would
|
||||
// stay -1 after a restart and the next Set would blindly re-learn whatever
|
||||
// dimension the caller happened to send, accepting a mismatched vector that
|
||||
// FT never indexes (silent search-side data loss). If the dimension can't be
|
||||
// parsed (e.g. an unexpected FT.INFO layout on some server version), keyLen
|
||||
// is left at -1 and validation degrades to the pre-restart lazy behaviour.
|
||||
func (s *ValkeyStore) loadIndexState(ctx context.Context) {
|
||||
msg, err := s.client.Do(ctx, s.client.B().FtInfo().Index(s.indexName).Build()).ToMessage()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
s.indexCreated = true
|
||||
if dim, ok := findDimensions(msg); ok && dim > 0 {
|
||||
s.keyLen = dim
|
||||
}
|
||||
}
|
||||
|
||||
// findDimensions walks an FT.INFO reply for the vector field's dimension. In
|
||||
// Valkey Search the VECTOR attribute nests its parameters under an `index`
|
||||
// array whose `dimensions` key holds the DIM the index was created with. The
|
||||
// reply is a nested array (RESP2) or map (RESP3), so we search recursively for
|
||||
// a `dimensions` key/token and read the value that follows it, tolerating both
|
||||
// integer and string-encoded values.
|
||||
func findDimensions(m valkey.ValkeyMessage) (int, bool) {
|
||||
if m.IsMap() {
|
||||
mp, err := m.AsMap()
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
for k, v := range mp {
|
||||
if strings.EqualFold(k, "dimensions") {
|
||||
if n, ok := msgToInt(v); ok {
|
||||
return n, true
|
||||
}
|
||||
}
|
||||
if d, ok := findDimensions(v); ok {
|
||||
return d, true
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
if m.IsArray() {
|
||||
arr, err := m.ToArray()
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
for i := range arr {
|
||||
if s, err := arr[i].ToString(); err == nil && strings.EqualFold(s, "dimensions") && i+1 < len(arr) {
|
||||
if n, ok := msgToInt(arr[i+1]); ok {
|
||||
return n, true
|
||||
}
|
||||
}
|
||||
if d, ok := findDimensions(arr[i]); ok {
|
||||
return d, true
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
// msgToInt reads an integer from a ValkeyMessage that may be an integer reply
|
||||
// or a string-encoded integer (FT.INFO mixes both across fields/versions).
|
||||
func msgToInt(m valkey.ValkeyMessage) (int, bool) {
|
||||
if n, err := m.ToInt64(); err == nil {
|
||||
return int(n), true
|
||||
}
|
||||
if s, err := m.ToString(); err == nil {
|
||||
if n, err := strconv.Atoi(s); err == nil {
|
||||
return n, true
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
// Free closes the Valkey client. Called by the gRPC server on shutdown.
|
||||
func (s *ValkeyStore) Free() error {
|
||||
if s.client != nil {
|
||||
s.client.Close()
|
||||
s.client = nil
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// buildTLSConfig assembles the tls.Config for a tls=true connection.
|
||||
// Go only auto-derives ServerName (SNI) from the dial address for hostnames;
|
||||
// for an IP-addressed endpoint (e.g. 10.0.0.5:6379) SNI is left empty and the
|
||||
// certificate's SANs won't match the raw IP, so verification fails. We set it
|
||||
// explicitly from the configured host so both hostname and IP endpoints verify.
|
||||
// A custom CA bundle (tls_ca_cert) and an explicit insecure-skip escape hatch
|
||||
// (tls_skip_verify) are supported for enterprise/self-signed setups.
|
||||
func buildTLSConfig(cfg Config) (*tls.Config, error) {
|
||||
tlsCfg := &tls.Config{}
|
||||
if host, _, err := net.SplitHostPort(cfg.Addr); err == nil && host != "" {
|
||||
tlsCfg.ServerName = host
|
||||
}
|
||||
if cfg.TLSSkipVerify {
|
||||
tlsCfg.InsecureSkipVerify = true
|
||||
}
|
||||
if cfg.TLSCACert != "" {
|
||||
pem, err := os.ReadFile(cfg.TLSCACert)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("valkey-store: read tls_ca_cert %q: %w", cfg.TLSCACert, err)
|
||||
}
|
||||
pool := x509.NewCertPool()
|
||||
if !pool.AppendCertsFromPEM(pem) {
|
||||
return nil, fmt.Errorf("valkey-store: tls_ca_cert %q: no valid certificate found", cfg.TLSCACert)
|
||||
}
|
||||
tlsCfg.RootCAs = pool
|
||||
}
|
||||
return tlsCfg, nil
|
||||
}
|
||||
|
||||
// ctx returns a request-scoped context bounded by the configured timeout. We
|
||||
// never rely on the client's built-in write timeout because index back-fill
|
||||
// and large KNN queries can legitimately exceed a short default.
|
||||
func (s *ValkeyStore) ctx() (context.Context, context.CancelFunc) {
|
||||
return context.WithTimeout(context.Background(), s.cfg.RequestTimeout)
|
||||
}
|
||||
|
||||
func (s *ValkeyStore) StoresSet(opts *pb.StoresSetOptions) error {
|
||||
keys := store.UnwrapKeys(opts.Keys)
|
||||
values := store.UnwrapValues(opts.Values)
|
||||
if len(keys) == 0 {
|
||||
return fmt.Errorf("valkey-store: Set: no keys to add")
|
||||
}
|
||||
if len(keys) != len(values) {
|
||||
return fmt.Errorf("valkey-store: Set: len(keys) = %d, len(values) = %d", len(keys), len(values))
|
||||
}
|
||||
|
||||
// Learn the dimension from the first key ever set (mirrors local-store's
|
||||
// keyLen == -1 sentinel), then reject anything that disagrees. checkDims is
|
||||
// the single source of truth for the per-key length check (shared with
|
||||
// Get/Delete/Find) so the four RPCs cannot drift apart.
|
||||
if s.keyLen == -1 {
|
||||
s.keyLen = len(keys[0])
|
||||
}
|
||||
if err := s.checkDims("Set", keys); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// The index needs the dimension up front, but local-store learns it from
|
||||
// the first Set — so we create it lazily here, once, before writing.
|
||||
if err := s.ensureIndex(s.keyLen); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Write each entry with an individual round-trip rather than pipelining the
|
||||
// whole batch via DoMulti. Valkey Search indexes every HSET into the
|
||||
// FLAT/HNSW index synchronously on the server's main thread; a large
|
||||
// pipeline of indexed writes can fill the socket buffers while that
|
||||
// indexing keeps the server from draining them, deadlocking the connection
|
||||
// (observed as an i/o timeout on high-dimension batches — a single 768-d
|
||||
// DoMulti of ~20 vectors hangs, while the same writes issued sequentially
|
||||
// complete in milliseconds). Sequential writes keep each command fully
|
||||
// round-tripped and stay fast (hundreds of 768-d vectors in a few hundred
|
||||
// ms). A single failure fails the whole Set — partial writes are surfaced,
|
||||
// not swallowed.
|
||||
//
|
||||
// The request timeout is applied PER command, not once across the whole
|
||||
// loop: an unbounded SetCols against a remote Valkey would otherwise exhaust
|
||||
// a single aggregate deadline mid-batch and leave a partial, non-atomic
|
||||
// write.
|
||||
for i, k := range keys {
|
||||
cmd := s.client.B().Hset().Key(encodeKey(s.prefix, k)).
|
||||
FieldValue().
|
||||
FieldValue(_vecField, valkey.BinaryString(vecToBytes(k))).
|
||||
FieldValue(_valField, valkey.BinaryString(values[i])).
|
||||
Build()
|
||||
ctx, cancel := s.ctx()
|
||||
err := s.client.Do(ctx, cmd).Error()
|
||||
cancel()
|
||||
if err != nil {
|
||||
return fmt.Errorf("valkey-store: Set: HSET key %d: %w", i, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// StoresGet fetches values for the given keys. Missing keys are omitted from
|
||||
// the result (not errored), matching local-store; returned slices are aligned.
|
||||
func (s *ValkeyStore) StoresGet(opts *pb.StoresGetOptions) (pb.StoresGetResult, error) {
|
||||
keys := store.UnwrapKeys(opts.Keys)
|
||||
if len(keys) == 0 {
|
||||
return pb.StoresGetResult{}, nil
|
||||
}
|
||||
if err := s.checkDims("Get", keys); err != nil {
|
||||
return pb.StoresGetResult{}, err
|
||||
}
|
||||
|
||||
// Reads pipeline the whole batch via DoMulti under ONE aggregate deadline,
|
||||
// unlike Set/Delete which use a per-command timeout. That asymmetry is
|
||||
// deliberate: HGET is non-mutating, so exhausting the deadline mid-batch
|
||||
// only truncates the result (surfaced as an error) — it can never leave a
|
||||
// partial write behind, which is the specific hazard the per-command timeout
|
||||
// guards against for Set/Delete. Pipelining is also safe here because these
|
||||
// are non-indexed reads (the indexed-DoMulti deadlock only affects writes).
|
||||
ctx, cancel := s.ctx()
|
||||
defer cancel()
|
||||
|
||||
cmds := make([]valkey.Completed, len(keys))
|
||||
for i, k := range keys {
|
||||
cmds[i] = s.client.B().Hget().Key(encodeKey(s.prefix, k)).Field(_valField).Build()
|
||||
}
|
||||
|
||||
var foundKeys [][]float32
|
||||
var foundValues [][]byte
|
||||
for i, res := range s.client.DoMulti(ctx, cmds...) {
|
||||
v, err := res.ToString()
|
||||
if err != nil {
|
||||
// A nil reply means the key/field is absent — omit it, don't error.
|
||||
if valkey.IsValkeyNil(err) {
|
||||
continue
|
||||
}
|
||||
return pb.StoresGetResult{}, fmt.Errorf("valkey-store: Get: HGET key %d: %w", i, err)
|
||||
}
|
||||
// The request vector is exact, so we return it verbatim as the key.
|
||||
foundKeys = append(foundKeys, keys[i])
|
||||
foundValues = append(foundValues, []byte(v))
|
||||
}
|
||||
|
||||
return pb.StoresGetResult{
|
||||
Keys: store.WrapKeys(foundKeys),
|
||||
Values: store.WrapValues(foundValues),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// StoresDelete removes entries by exact vector. Missing keys are tolerated
|
||||
// (DEL returns 0), matching local-store.
|
||||
func (s *ValkeyStore) StoresDelete(opts *pb.StoresDeleteOptions) error {
|
||||
keys := store.UnwrapKeys(opts.Keys)
|
||||
if len(keys) == 0 {
|
||||
return fmt.Errorf("valkey-store: Delete: no keys to delete")
|
||||
}
|
||||
if err := s.checkDims("Delete", keys); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Sequential DELs for the same reason StoresSet avoids DoMulti: a DEL of an
|
||||
// indexed key mutates the search index on the server's main thread, and a
|
||||
// large pipeline of such mutations can deadlock the connection. Missing
|
||||
// keys (DEL returns 0) are tolerated, matching local-store. As in Set, the
|
||||
// timeout is per command so a large DeleteCols cannot exhaust one aggregate
|
||||
// deadline mid-batch.
|
||||
for i, k := range keys {
|
||||
ctx, cancel := s.ctx()
|
||||
err := s.client.Do(ctx, s.client.B().Del().Key(encodeKey(s.prefix, k)).Build()).Error()
|
||||
cancel()
|
||||
if err != nil {
|
||||
return fmt.Errorf("valkey-store: Delete: DEL key %d: %w", i, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// StoresFind returns the topK nearest entries by the configured distance
|
||||
// metric, ordered most-similar first. An empty/uncreated index returns empty
|
||||
// slices and no error, matching local-store's empty-store behaviour.
|
||||
func (s *ValkeyStore) StoresFind(opts *pb.StoresFindOptions) (pb.StoresFindResult, error) {
|
||||
// Guard against a malformed gRPC request with a nil/empty Key before
|
||||
// dereferencing it — a nil opts.Key would otherwise panic the backend.
|
||||
if opts.Key == nil || len(opts.Key.Floats) == 0 {
|
||||
return pb.StoresFindResult{}, fmt.Errorf("valkey-store: Find: query key is empty")
|
||||
}
|
||||
query := opts.Key.Floats
|
||||
topK := int(opts.TopK)
|
||||
if topK < 1 {
|
||||
return pb.StoresFindResult{}, fmt.Errorf("valkey-store: Find: topK = %d, must be >= 1", topK)
|
||||
}
|
||||
if topK > _maxTopK {
|
||||
xlog.Warn("valkey-store: Find topK clamped", "requested", topK, "max", _maxTopK)
|
||||
topK = _maxTopK
|
||||
}
|
||||
// No index yet means nothing has been Set (and none was found at Load) —
|
||||
// an empty result, not an error.
|
||||
if !s.indexCreated {
|
||||
return pb.StoresFindResult{}, nil
|
||||
}
|
||||
// Enforce the query dimension against the known keyLen — recovered from
|
||||
// FT.INFO at Load after a restart, or learned from the first Set — so a
|
||||
// wrong-dimension query gets the clean local-store-style error. keyLen is
|
||||
// only -1 in the degraded case where FT.INFO gave no parseable dimension;
|
||||
// then we let Valkey's own FT.SEARCH validate the query vector.
|
||||
if s.keyLen != -1 && len(query) != s.keyLen {
|
||||
return pb.StoresFindResult{}, fmt.Errorf("valkey-store: Find: query length %d does not match existing %d", len(query), s.keyLen)
|
||||
}
|
||||
|
||||
ctx, cancel := s.ctx()
|
||||
defer cancel()
|
||||
|
||||
// KNN pre-filter query: match everything, rank by vector distance into the
|
||||
// __score alias. A pure KNN query already returns its topK results ordered
|
||||
// by distance ascending (nearest-first), so we do NOT add SORTBY __score:
|
||||
// Valkey Search rejects sorting on the KNN score alias ("Index field
|
||||
// `__score` does not exist" — it is a query-time computed field, not a
|
||||
// SORTABLE schema attribute). LIMIT 0 topK caps the result and DIALECT 2 is
|
||||
// required for the =>[KNN ...] vector syntax. The __score field is still
|
||||
// returned in each document and read back for the similarity conversion.
|
||||
//
|
||||
// Injection-safety: the only caller-controlled value interpolated here is
|
||||
// topK (an int, already bounded above). _vecField and _scoreField are
|
||||
// compile-time constants, so this Sprintf cannot be used to inject query
|
||||
// syntax. Do NOT make those fields operator-configurable without sanitizing
|
||||
// them first — the KNN query string is otherwise built only from constants.
|
||||
q := fmt.Sprintf("*=>[KNN %d @%s $q AS %s]", topK, _vecField, _scoreField)
|
||||
cmd := s.client.B().FtSearch().Index(s.indexName).Query(q).
|
||||
Return("3").Identifier(_vecField).Identifier(_valField).Identifier(_scoreField).
|
||||
Limit().OffsetNum(0, int64(topK)).
|
||||
Params().Nargs(2).NameValue().NameValue("q", valkey.VectorString32(query)).
|
||||
Dialect(2).
|
||||
Build()
|
||||
|
||||
_, docs, err := s.client.Do(ctx, cmd).AsFtSearch()
|
||||
if err != nil {
|
||||
// The cached indexCreated flag can go stale: an operator runs
|
||||
// FT.DROPINDEX out of band, or two processes race on a fresh namespace.
|
||||
// If the index is gone, mirror local-store's empty-store behaviour
|
||||
// (empty result, no error) and clear the flag so a later Set recreates
|
||||
// it, rather than surfacing a hard error for what looks like an empty
|
||||
// store to the caller.
|
||||
if isNoSuchIndexErr(err) {
|
||||
s.indexCreated = false
|
||||
return pb.StoresFindResult{}, nil
|
||||
}
|
||||
return pb.StoresFindResult{}, fmt.Errorf("valkey-store: Find: FT.SEARCH: %w", err)
|
||||
}
|
||||
|
||||
keys := make([][]float32, 0, len(docs))
|
||||
values := make([][]byte, 0, len(docs))
|
||||
sims := make([]float32, 0, len(docs))
|
||||
for _, doc := range docs {
|
||||
// Decode the key from the returned `vec` bytes rather than the Valkey
|
||||
// key string: this guarantees the exact original float ordering/values
|
||||
// without a hex round-trip.
|
||||
vecBytes := []byte(doc.Doc[_vecField])
|
||||
k, err := bytesToVec(vecBytes)
|
||||
if err != nil {
|
||||
return pb.StoresFindResult{}, fmt.Errorf("valkey-store: Find: decode vec: %w", err)
|
||||
}
|
||||
dist, err := strconv.ParseFloat(doc.Doc[_scoreField], 64)
|
||||
if err != nil {
|
||||
return pb.StoresFindResult{}, fmt.Errorf("valkey-store: Find: parse score %q: %w", doc.Doc[_scoreField], err)
|
||||
}
|
||||
keys = append(keys, k)
|
||||
values = append(values, []byte(doc.Doc[_valField]))
|
||||
sims = append(sims, distanceToSimilarity(s.cfg.DistanceMetric, dist))
|
||||
}
|
||||
|
||||
return pb.StoresFindResult{
|
||||
Keys: store.WrapKeys(keys),
|
||||
Values: store.WrapValues(values),
|
||||
Similarities: sims,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ensureIndex creates the FT vector index once, lazily, on the first Set. The
|
||||
// dimension is fixed at creation (a second guard on top of the Go-side keyLen
|
||||
// check). An "already exists" error is treated as success so a restart against
|
||||
// a persisted index is a no-op.
|
||||
func (s *ValkeyStore) ensureIndex(dim int) error {
|
||||
if s.indexCreated {
|
||||
return nil
|
||||
}
|
||||
|
||||
// VECTOR attribute tokens. The count that follows the algorithm name is the
|
||||
// number of these tokens, so we build the slice and derive the count from
|
||||
// it — no hand-maintained magic number that drifts when HNSW knobs change.
|
||||
attrs := []string{"TYPE", "FLOAT32", "DIM", strconv.Itoa(dim), "DISTANCE_METRIC", s.cfg.DistanceMetric}
|
||||
if s.cfg.IndexAlgo == indexAlgoHNSW {
|
||||
attrs = append(attrs,
|
||||
"M", strconv.Itoa(s.cfg.HNSW.M),
|
||||
"EF_CONSTRUCTION", strconv.Itoa(s.cfg.HNSW.EFConstruction),
|
||||
"EF_RUNTIME", strconv.Itoa(s.cfg.HNSW.EFRuntime),
|
||||
)
|
||||
}
|
||||
|
||||
args := []string{
|
||||
s.indexName,
|
||||
"ON", "HASH",
|
||||
"PREFIX", "1", s.prefix,
|
||||
"SCHEMA", _vecField, "VECTOR", s.cfg.IndexAlgo, strconv.Itoa(len(attrs)),
|
||||
}
|
||||
args = append(args, attrs...)
|
||||
|
||||
// FT.CREATE has no typed builder entry point, so we use the Arbitrary escape
|
||||
// hatch. All tokens are non-key args in standalone mode.
|
||||
ctx, cancel := s.ctx()
|
||||
defer cancel()
|
||||
err := s.client.Do(ctx, s.client.B().Arbitrary("FT.CREATE").Args(args...).Build()).Error()
|
||||
if err != nil && !isIndexExistsErr(err) {
|
||||
return fmt.Errorf("valkey-store: FT.CREATE %s: %w", s.indexName, err)
|
||||
}
|
||||
|
||||
s.indexCreated = true
|
||||
return nil
|
||||
}
|
||||
|
||||
// checkDims rejects any key whose dimension disagrees with the learned keyLen.
|
||||
// When keyLen is still open (-1, nothing set yet) there is nothing to check.
|
||||
func (s *ValkeyStore) checkDims(op string, keys [][]float32) error {
|
||||
if s.keyLen == -1 {
|
||||
return nil
|
||||
}
|
||||
for i, k := range keys {
|
||||
if len(k) != s.keyLen {
|
||||
return fmt.Errorf("valkey-store: %s: key %d length %d does not match existing %d", op, i, len(k), s.keyLen)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// distanceToSimilarity converts a Valkey distance into local-store's similarity
|
||||
// convention. Only COSINE has a defined [-1, 1] similarity (sim = 1 - dist);
|
||||
// for L2/IP the raw score is returned as the "similarity" with a documented
|
||||
// meaning (smaller L2 = closer; larger IP = closer).
|
||||
func distanceToSimilarity(metric string, dist float64) float32 {
|
||||
if metric == distanceCosine {
|
||||
return float32(1 - dist)
|
||||
}
|
||||
return float32(dist)
|
||||
}
|
||||
|
||||
// isIndexExistsErr reports whether an FT.CREATE error is the benign
|
||||
// "index already exists" case (e.g. after a restart against a persisted index).
|
||||
func isIndexExistsErr(err error) bool {
|
||||
return strings.Contains(strings.ToLower(err.Error()), "already exists")
|
||||
}
|
||||
|
||||
// isNoSuchIndexErr reports whether an FT.SEARCH error means the index no longer
|
||||
// exists (dropped out of band, or never really created despite a stale cached
|
||||
// flag). Valkey Search phrases this differently across versions, so we match
|
||||
// the common variants rather than one exact string.
|
||||
func isNoSuchIndexErr(err error) bool {
|
||||
msg := strings.ToLower(err.Error())
|
||||
if !strings.Contains(msg, "index") {
|
||||
return false
|
||||
}
|
||||
return strings.Contains(msg, "no such index") ||
|
||||
strings.Contains(msg, "not exist") ||
|
||||
strings.Contains(msg, "not found") ||
|
||||
strings.Contains(msg, "unknown index")
|
||||
}
|
||||
|
||||
// keyPrefix / indexName derive per-namespace identifiers from the model name so
|
||||
// entries and indexes never collide across namespaces on a shared server.
|
||||
func keyPrefix(namespace string) string {
|
||||
return _keyPrefixPrefix + nsToken(namespace) + ":"
|
||||
}
|
||||
|
||||
func indexName(namespace string) string {
|
||||
return _indexPrefix + nsToken(namespace)
|
||||
}
|
||||
|
||||
// nsToken maps a namespace to a collision-resistant, printable token. sanitize()
|
||||
// alone is lossy (many distinct characters all fold to '_'), so namespaces like
|
||||
// "a b", "a/b" and "a:b" would otherwise share one keyspace and FT index — a
|
||||
// silent data-isolation bug (one store reading/clobbering another). We append a
|
||||
// short hash of the ORIGINAL namespace so distinct names never collide, while
|
||||
// the sanitized part keeps the token human-readable. It is deterministic, so a
|
||||
// persisted index is found again after a restart.
|
||||
func nsToken(namespace string) string {
|
||||
sum := sha256.Sum256([]byte(namespace))
|
||||
// Cap the human-readable part so a pathologically long model name can't
|
||||
// produce an unbounded key prefix / index name (which would degrade Valkey
|
||||
// performance). The 8-char hash suffix below already guarantees collision
|
||||
// resistance regardless of truncation, so trimming the readable part is safe.
|
||||
readable := sanitize(namespace)
|
||||
if len(readable) > _maxNsTokenLen {
|
||||
readable = readable[:_maxNsTokenLen]
|
||||
}
|
||||
return readable + "-" + hex.EncodeToString(sum[:])[:8]
|
||||
}
|
||||
|
||||
// sanitize maps a namespace to a safe token for keys/index names: alphanumeric,
|
||||
// '_', '-' and '.' pass through; everything else becomes '_'. An empty
|
||||
// namespace becomes "default" so the key/index names stay well-formed.
|
||||
func sanitize(namespace string) string {
|
||||
if namespace == "" {
|
||||
return "default"
|
||||
}
|
||||
var b strings.Builder
|
||||
b.Grow(len(namespace))
|
||||
for _, r := range namespace {
|
||||
switch {
|
||||
case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9', r == '_', r == '-', r == '.':
|
||||
b.WriteRune(r)
|
||||
default:
|
||||
b.WriteRune('_')
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
@@ -1,13 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
func TestValkeyStore(t *testing.T) {
|
||||
RegisterFailHandler(Fail)
|
||||
RunSpecs(t, "valkey-store test suite")
|
||||
}
|
||||
@@ -1,598 +0,0 @@
|
||||
package main
|
||||
|
||||
// Unit tests for the Valkey store, using the valkey-go gomock client so they
|
||||
// run with no container. They assert the exact commands built for each RPC
|
||||
// (the wire contract) plus the local-store parity semantics: empty/len/dim
|
||||
// rejects, omit-missing Get, tolerate-missing Delete, topK<1 reject, the
|
||||
// sim = 1 - distance conversion, lazy FT.CREATE, and the HNSW arg-shape.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strconv"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/LocalAI/pkg/store"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
valkey "github.com/valkey-io/valkey-go"
|
||||
"github.com/valkey-io/valkey-go/mock"
|
||||
"go.uber.org/mock/gomock"
|
||||
)
|
||||
|
||||
const testNamespace = "test"
|
||||
|
||||
func testCfg() Config {
|
||||
cfg, err := loadConfig(nil) // reads defaults when no options are set
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
return cfg
|
||||
}
|
||||
|
||||
// opts builds a *pb.ModelOptions carrying the given key:value option strings,
|
||||
// mirroring how core threads a store's model-config `options:` list to the
|
||||
// backend's LoadModel.
|
||||
func opts(kv ...string) *pb.ModelOptions {
|
||||
return &pb.ModelOptions{Options: kv}
|
||||
}
|
||||
|
||||
func newMockStore(cfg Config) (*ValkeyStore, *mock.Client) {
|
||||
ctrl := gomock.NewController(GinkgoT())
|
||||
DeferCleanup(ctrl.Finish)
|
||||
c := mock.NewClient(ctrl)
|
||||
return newWithClient(c, cfg, testNamespace), c
|
||||
}
|
||||
|
||||
func wrapSet(keys [][]float32, values [][]byte) *pb.StoresSetOptions {
|
||||
return &pb.StoresSetOptions{Keys: store.WrapKeys(keys), Values: store.WrapValues(values)}
|
||||
}
|
||||
|
||||
var _ = Describe("loadConfig", func() {
|
||||
It("uses documented defaults", func() {
|
||||
cfg, err := loadConfig(nil)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(cfg.Addr).To(Equal("localhost:6379"))
|
||||
Expect(cfg.ClientName).To(Equal("localai-valkey-store"))
|
||||
Expect(cfg.IndexAlgo).To(Equal("FLAT"))
|
||||
Expect(cfg.DistanceMetric).To(Equal("COSINE"))
|
||||
Expect(cfg.RequestTimeout.Milliseconds()).To(Equal(int64(5000)))
|
||||
})
|
||||
|
||||
It("honours option overrides", func() {
|
||||
cfg, err := loadConfig(opts(
|
||||
"addr:valkey.example:6380",
|
||||
"index_algo:hnsw",
|
||||
"distance_metric:l2",
|
||||
"request_timeout_ms:1234",
|
||||
))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
// addr keeps its embedded colon: strings.Cut splits on the first ':'.
|
||||
Expect(cfg.Addr).To(Equal("valkey.example:6380"))
|
||||
Expect(cfg.IndexAlgo).To(Equal("HNSW"))
|
||||
Expect(cfg.DistanceMetric).To(Equal("L2"))
|
||||
Expect(cfg.RequestTimeout.Milliseconds()).To(Equal(int64(1234)))
|
||||
})
|
||||
|
||||
It("keeps the mandatory client name when blanked", func() {
|
||||
cfg, err := loadConfig(opts("client_name:"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(cfg.ClientName).To(Equal("localai-valkey-store"))
|
||||
})
|
||||
|
||||
It("ignores a malformed option without a colon", func() {
|
||||
cfg, err := loadConfig(opts("addr:valkey.example:6380", "not-a-kv-pair"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(cfg.Addr).To(Equal("valkey.example:6380"))
|
||||
})
|
||||
|
||||
It("rejects an invalid index algo", func() {
|
||||
_, err := loadConfig(opts("index_algo:bogus"))
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("rejects an invalid distance metric", func() {
|
||||
_, err := loadConfig(opts("distance_metric:bogus"))
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("fails fast on a malformed HNSW integer instead of silently defaulting", func() {
|
||||
_, err := loadConfig(opts("hnsw_m:1x6"))
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("hnsw_m"))
|
||||
})
|
||||
|
||||
It("honours a valid db override", func() {
|
||||
cfg, err := loadConfig(opts("db:3"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(cfg.DB).To(Equal(3))
|
||||
})
|
||||
|
||||
It("rejects a negative db", func() {
|
||||
_, err := loadConfig(opts("db:-1"))
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("resolves username from direct option", func() {
|
||||
cfg, err := loadConfig(opts("username:admin"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(cfg.Username).To(Equal("admin"))
|
||||
})
|
||||
|
||||
It("resolves password from env indirection via password_env", func() {
|
||||
GinkgoT().Setenv("TEST_VALKEY_PW_INDIRECT", "s3cret")
|
||||
cfg, err := loadConfig(opts("password_env:TEST_VALKEY_PW_INDIRECT"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(cfg.Password).To(Equal("s3cret"))
|
||||
})
|
||||
|
||||
It("resolves username from env indirection via username_env", func() {
|
||||
GinkgoT().Setenv("TEST_VALKEY_USER_INDIRECT", "myuser")
|
||||
cfg, err := loadConfig(opts("username_env:TEST_VALKEY_USER_INDIRECT"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(cfg.Username).To(Equal("myuser"))
|
||||
})
|
||||
|
||||
It("prefers the direct option over env indirection", func() {
|
||||
GinkgoT().Setenv("TEST_VALKEY_PW_CLASH", "from-env")
|
||||
cfg, err := loadConfig(opts("password:direct-value", "password_env:TEST_VALKEY_PW_CLASH"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(cfg.Password).To(Equal("direct-value"))
|
||||
})
|
||||
|
||||
It("returns empty when neither direct nor env indirection is set", func() {
|
||||
cfg, err := loadConfig(opts("addr:localhost:6379"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(cfg.Username).To(BeEmpty())
|
||||
Expect(cfg.Password).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("Load namespace gate", func() {
|
||||
It("accepts prefixed store namespaces", func() {
|
||||
s := NewValkeyStore()
|
||||
// Load will fail at the Valkey connect step (no server), but only after
|
||||
// passing the namespace gate. A network error is acceptable here — it
|
||||
// means the prefix check passed.
|
||||
err := s.Load(&pb.ModelOptions{Model: store.NamespacePrefix + "any-namespace", Options: []string{"addr:localhost:1"}})
|
||||
Expect(err).NotTo(MatchError(ContainSubstring("not a store namespace")))
|
||||
})
|
||||
|
||||
It("accepts the prefix alone (default store)", func() {
|
||||
s := NewValkeyStore()
|
||||
err := s.Load(&pb.ModelOptions{Model: store.NamespacePrefix, Options: []string{"addr:localhost:1"}})
|
||||
Expect(err).NotTo(MatchError(ContainSubstring("not a store namespace")))
|
||||
})
|
||||
|
||||
It("refuses model names without the namespace prefix", func() {
|
||||
s := NewValkeyStore()
|
||||
err := s.Load(&pb.ModelOptions{Model: "some-llm.gguf"})
|
||||
Expect(err).To(MatchError(ContainSubstring("not a store namespace")))
|
||||
})
|
||||
|
||||
It("refuses an empty model name", func() {
|
||||
s := NewValkeyStore()
|
||||
err := s.Load(&pb.ModelOptions{})
|
||||
Expect(err).To(MatchError(ContainSubstring("not a store namespace")))
|
||||
})
|
||||
|
||||
It("refuses nil opts", func() {
|
||||
s := NewValkeyStore()
|
||||
err := s.Load(nil)
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("StoresSet", func() {
|
||||
It("rejects empty input", func() {
|
||||
s, _ := newMockStore(testCfg())
|
||||
Expect(s.StoresSet(&pb.StoresSetOptions{})).NotTo(Succeed())
|
||||
})
|
||||
|
||||
It("rejects key/value length mismatch", func() {
|
||||
s, _ := newMockStore(testCfg())
|
||||
err := s.StoresSet(wrapSet([][]float32{{1, 0, 0}}, [][]byte{[]byte("a"), []byte("b")}))
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("rejects dimension mismatch on a later add", func() {
|
||||
s, c := newMockStore(testCfg())
|
||||
// First Set issues FT.CREATE then a sequential HSET (both via Do).
|
||||
c.EXPECT().Do(gomock.Any(), gomock.Any()).DoAndReturn(
|
||||
func(ctx context.Context, cmd valkey.Completed) valkey.ValkeyResult {
|
||||
if cmd.Commands()[0] == "FT.CREATE" {
|
||||
return assertFTCreate(3, "FLAT")(ctx, cmd)
|
||||
}
|
||||
return mock.Result(mock.ValkeyInt64(1)) // HSET
|
||||
}).AnyTimes()
|
||||
Expect(s.StoresSet(wrapSet([][]float32{{1, 0, 0}}, [][]byte{[]byte("3d")}))).To(Succeed())
|
||||
|
||||
err := s.StoresSet(wrapSet([][]float32{{1, 0}}, [][]byte{[]byte("2d")}))
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("rejects dimension mismatch within a batch", func() {
|
||||
s, _ := newMockStore(testCfg())
|
||||
err := s.StoresSet(wrapSet([][]float32{{1, 0, 0}, {1, 0}}, [][]byte{[]byte("3d"), []byte("2d")}))
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("creates the FLAT index once and HSETs each entry", func() {
|
||||
s, c := newMockStore(testCfg())
|
||||
// FT.CREATE must run exactly once (on the first Set); each entry is then
|
||||
// written with an individual sequential HSET (Do, not DoMulti — see the
|
||||
// pipeline-deadlock note in StoresSet).
|
||||
var ftCreateCount, hsetCount int
|
||||
c.EXPECT().Do(gomock.Any(), gomock.Any()).DoAndReturn(
|
||||
func(ctx context.Context, cmd valkey.Completed) valkey.ValkeyResult {
|
||||
toks := cmd.Commands()
|
||||
switch toks[0] {
|
||||
case "FT.CREATE":
|
||||
ftCreateCount++
|
||||
return assertFTCreate(3, "FLAT")(ctx, cmd)
|
||||
case "HSET":
|
||||
hsetCount++
|
||||
Expect(toks[1]).To(HavePrefix(s.prefix))
|
||||
Expect(toks).To(ContainElements("vec", "val"))
|
||||
return mock.Result(mock.ValkeyInt64(1))
|
||||
default:
|
||||
Fail("unexpected command: " + toks[0])
|
||||
return valkey.ValkeyResult{}
|
||||
}
|
||||
}).AnyTimes()
|
||||
Expect(s.StoresSet(wrapSet([][]float32{{1, 0, 0}}, [][]byte{[]byte("a")}))).To(Succeed())
|
||||
Expect(s.StoresSet(wrapSet([][]float32{{2, 0, 0}}, [][]byte{[]byte("b")}))).To(Succeed())
|
||||
|
||||
Expect(ftCreateCount).To(Equal(1))
|
||||
Expect(hsetCount).To(Equal(2))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("StoresGet", func() {
|
||||
It("round-trips values and omits missing keys", func() {
|
||||
s, c := newMockStore(testCfg())
|
||||
s.keyLen = 3
|
||||
s.indexCreated = true
|
||||
// First key present, second missing (nil reply).
|
||||
c.EXPECT().DoMulti(gomock.Any(), gomock.Any(), gomock.Any()).Return([]valkey.ValkeyResult{
|
||||
mock.Result(mock.ValkeyString("hello")),
|
||||
mock.Result(mock.ValkeyNil()),
|
||||
}).Times(1)
|
||||
|
||||
res, err := s.StoresGet(&pb.StoresGetOptions{
|
||||
Keys: store.WrapKeys([][]float32{{1, 0, 0}, {9, 0, 0}}),
|
||||
})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(res.Keys).To(HaveLen(1))
|
||||
Expect(res.Values).To(HaveLen(1))
|
||||
Expect(res.Values[0].Bytes).To(Equal([]byte("hello")))
|
||||
})
|
||||
|
||||
It("rejects dimension mismatch", func() {
|
||||
s, _ := newMockStore(testCfg())
|
||||
s.keyLen = 3
|
||||
_, err := s.StoresGet(&pb.StoresGetOptions{Keys: store.WrapKeys([][]float32{{1, 0}})})
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("StoresDelete", func() {
|
||||
It("issues DEL per key and tolerates missing", func() {
|
||||
s, c := newMockStore(testCfg())
|
||||
s.keyLen = 3
|
||||
// DEL of a missing key returns 0 — still a success. DELs are issued
|
||||
// sequentially (Do, not DoMulti — see the deadlock note in StoresSet).
|
||||
c.EXPECT().Do(gomock.Any(), gomock.Any()).DoAndReturn(
|
||||
func(_ context.Context, cmd valkey.Completed) valkey.ValkeyResult {
|
||||
Expect(cmd.Commands()[0]).To(Equal("DEL"))
|
||||
return mock.Result(mock.ValkeyInt64(0))
|
||||
}).Times(1)
|
||||
Expect(s.StoresDelete(&pb.StoresDeleteOptions{
|
||||
Keys: store.WrapKeys([][]float32{{9, 0, 0}}),
|
||||
})).To(Succeed())
|
||||
})
|
||||
|
||||
It("rejects dimension mismatch", func() {
|
||||
s, _ := newMockStore(testCfg())
|
||||
s.keyLen = 3
|
||||
err := s.StoresDelete(&pb.StoresDeleteOptions{Keys: store.WrapKeys([][]float32{{1, 0}})})
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("StoresFind", func() {
|
||||
It("builds the KNN query and converts distance to similarity nearest-first", func() {
|
||||
s, c := newMockStore(testCfg())
|
||||
s.keyLen = 3
|
||||
s.indexCreated = true
|
||||
|
||||
// Distances 0, 1, 2 must map to similarities 1, 0, -1 (COSINE).
|
||||
docs := []ftDoc{
|
||||
{vec: []float32{1, 0, 0}, val: "a", dist: 0},
|
||||
{vec: []float32{0, 1, 0}, val: "b", dist: 1},
|
||||
{vec: []float32{-1, 0, 0}, val: "c", dist: 2},
|
||||
}
|
||||
c.EXPECT().Do(gomock.Any(), gomock.Any()).DoAndReturn(
|
||||
func(_ context.Context, cmd valkey.Completed) valkey.ValkeyResult {
|
||||
toks := cmd.Commands()
|
||||
Expect(toks[0]).To(Equal("FT.SEARCH"))
|
||||
Expect(toks[1]).To(Equal(s.indexName))
|
||||
Expect(toks[2]).To(ContainSubstring("KNN 3 @vec $q AS __score"))
|
||||
Expect(toks).To(ContainElements("PARAMS", "2", "q", "DIALECT", "2"))
|
||||
return mock.Result(ftSearchReply(s.prefix, docs))
|
||||
}).Times(1)
|
||||
|
||||
keys, values, sims, err := findViaRPC(s, []float32{1, 0, 0}, 3)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(sims).To(Equal([]float32{1, 0, -1}))
|
||||
Expect(values[0]).To(Equal([]byte("a")))
|
||||
// Key decoded from returned vec bytes equals the original vector.
|
||||
Expect(keys[0]).To(Equal([]float32{1, 0, 0}))
|
||||
})
|
||||
|
||||
It("rejects topK < 1", func() {
|
||||
s, _ := newMockStore(testCfg())
|
||||
s.keyLen = 3
|
||||
s.indexCreated = true
|
||||
_, err := s.StoresFind(&pb.StoresFindOptions{Key: &pb.StoresKey{Floats: []float32{1, 0, 0}}, TopK: 0})
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("rejects a nil Key without panicking", func() {
|
||||
s, _ := newMockStore(testCfg())
|
||||
s.keyLen = 3
|
||||
s.indexCreated = true
|
||||
_, err := s.StoresFind(&pb.StoresFindOptions{TopK: 5})
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("rejects an empty query vector", func() {
|
||||
s, _ := newMockStore(testCfg())
|
||||
s.keyLen = 3
|
||||
s.indexCreated = true
|
||||
_, err := s.StoresFind(&pb.StoresFindOptions{Key: &pb.StoresKey{Floats: []float32{}}, TopK: 5})
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("rejects query dimension mismatch", func() {
|
||||
s, _ := newMockStore(testCfg())
|
||||
s.keyLen = 3
|
||||
s.indexCreated = true
|
||||
_, err := s.StoresFind(&pb.StoresFindOptions{Key: &pb.StoresKey{Floats: []float32{1, 0}}, TopK: 1})
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("returns empty (no error) when the index was never created", func() {
|
||||
s, _ := newMockStore(testCfg())
|
||||
res, err := s.StoresFind(&pb.StoresFindOptions{Key: &pb.StoresKey{Floats: []float32{1, 0, 0}}, TopK: 5})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(res.Keys).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("ensureIndex arg-shape", func() {
|
||||
It("emits HNSW tuning tokens when the algo is HNSW", func() {
|
||||
cfg := testCfg()
|
||||
cfg.IndexAlgo = indexAlgoHNSW
|
||||
cfg.HNSW = hnswParams{M: 16, EFConstruction: 200, EFRuntime: 10}
|
||||
s, c := newMockStore(cfg)
|
||||
c.EXPECT().Do(gomock.Any(), gomock.Any()).DoAndReturn(
|
||||
func(_ context.Context, cmd valkey.Completed) valkey.ValkeyResult {
|
||||
toks := cmd.Commands()
|
||||
Expect(toks).To(ContainElements("HNSW", "M", "16", "EF_CONSTRUCTION", "200", "EF_RUNTIME", "10"))
|
||||
return mock.Result(mock.ValkeyString("OK"))
|
||||
}).Times(1)
|
||||
Expect(s.ensureIndex(4)).To(Succeed())
|
||||
})
|
||||
|
||||
It("treats an already-exists error as success", func() {
|
||||
s, c := newMockStore(testCfg())
|
||||
c.EXPECT().Do(gomock.Any(), gomock.Any()).Return(
|
||||
mock.ErrorResult(fmt.Errorf("Index already exists"))).Times(1)
|
||||
Expect(s.ensureIndex(4)).To(Succeed())
|
||||
Expect(s.indexCreated).To(BeTrue())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("distanceToSimilarity", func() {
|
||||
It("converts cosine distance to similarity", func() {
|
||||
Expect(distanceToSimilarity(distanceCosine, 0)).To(Equal(float32(1)))
|
||||
Expect(distanceToSimilarity(distanceCosine, 1)).To(Equal(float32(0)))
|
||||
Expect(distanceToSimilarity(distanceCosine, 2)).To(Equal(float32(-1)))
|
||||
})
|
||||
|
||||
It("passes the raw score through for non-cosine metrics", func() {
|
||||
Expect(distanceToSimilarity(distanceL2, 0.42)).To(Equal(float32(0.42)))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("findDimensions", func() {
|
||||
It("recovers an integer dimension from a nested FT.INFO reply", func() {
|
||||
dim, ok := findDimensions(ftInfoReply(mock.ValkeyInt64(768)))
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(dim).To(Equal(768))
|
||||
})
|
||||
|
||||
It("recovers a string-encoded dimension", func() {
|
||||
dim, ok := findDimensions(ftInfoReply(mock.ValkeyString("384")))
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(dim).To(Equal(384))
|
||||
})
|
||||
|
||||
It("reports not-found when no dimension token is present", func() {
|
||||
reply := mock.ValkeyArray(
|
||||
mock.ValkeyString("index_name"), mock.ValkeyString("idx:test"),
|
||||
mock.ValkeyString("num_docs"), mock.ValkeyInt64(0),
|
||||
)
|
||||
_, ok := findDimensions(reply)
|
||||
Expect(ok).To(BeFalse())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("loadIndexState", func() {
|
||||
It("recovers indexCreated and keyLen from a persisted index", func() {
|
||||
s, c := newMockStore(testCfg())
|
||||
c.EXPECT().Do(gomock.Any(), gomock.Any()).DoAndReturn(
|
||||
func(_ context.Context, cmd valkey.Completed) valkey.ValkeyResult {
|
||||
Expect(cmd.Commands()[0]).To(Equal("FT.INFO"))
|
||||
return mock.Result(ftInfoReply(mock.ValkeyInt64(768)))
|
||||
}).Times(1)
|
||||
|
||||
ctx, cancel := s.ctx()
|
||||
defer cancel()
|
||||
s.loadIndexState(ctx)
|
||||
Expect(s.indexCreated).To(BeTrue())
|
||||
Expect(s.keyLen).To(Equal(768))
|
||||
})
|
||||
|
||||
It("leaves state untouched when the index does not exist", func() {
|
||||
s, c := newMockStore(testCfg())
|
||||
c.EXPECT().Do(gomock.Any(), gomock.Any()).Return(
|
||||
mock.ErrorResult(fmt.Errorf("Index with name 'idx:test' not found"))).Times(1)
|
||||
|
||||
ctx, cancel := s.ctx()
|
||||
defer cancel()
|
||||
s.loadIndexState(ctx)
|
||||
Expect(s.indexCreated).To(BeFalse())
|
||||
Expect(s.keyLen).To(Equal(-1))
|
||||
})
|
||||
|
||||
It("marks the index created but leaves keyLen open when the dim is unparseable", func() {
|
||||
s, c := newMockStore(testCfg())
|
||||
c.EXPECT().Do(gomock.Any(), gomock.Any()).Return(
|
||||
mock.Result(mock.ValkeyArray(mock.ValkeyString("index_name"), mock.ValkeyString("idx:test")))).Times(1)
|
||||
|
||||
ctx, cancel := s.ctx()
|
||||
defer cancel()
|
||||
s.loadIndexState(ctx)
|
||||
Expect(s.indexCreated).To(BeTrue())
|
||||
Expect(s.keyLen).To(Equal(-1))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("StoresFind on a dropped index", func() {
|
||||
It("returns empty (no error) and clears the stale flag when the index is gone", func() {
|
||||
s, c := newMockStore(testCfg())
|
||||
s.keyLen = 3
|
||||
s.indexCreated = true
|
||||
c.EXPECT().Do(gomock.Any(), gomock.Any()).Return(
|
||||
mock.ErrorResult(fmt.Errorf("Index with name 'idx:test' not found"))).Times(1)
|
||||
|
||||
res, err := s.StoresFind(&pb.StoresFindOptions{Key: &pb.StoresKey{Floats: []float32{1, 0, 0}}, TopK: 5})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(res.Keys).To(BeEmpty())
|
||||
Expect(s.indexCreated).To(BeFalse())
|
||||
})
|
||||
|
||||
It("still surfaces a genuine FT.SEARCH error", func() {
|
||||
s, c := newMockStore(testCfg())
|
||||
s.keyLen = 3
|
||||
s.indexCreated = true
|
||||
c.EXPECT().Do(gomock.Any(), gomock.Any()).Return(
|
||||
mock.ErrorResult(fmt.Errorf("timeout"))).Times(1)
|
||||
|
||||
_, err := s.StoresFind(&pb.StoresFindOptions{Key: &pb.StoresKey{Floats: []float32{1, 0, 0}}, TopK: 5})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(s.indexCreated).To(BeTrue())
|
||||
})
|
||||
})
|
||||
|
||||
// --- test helpers ---
|
||||
|
||||
type ftDoc struct {
|
||||
vec []float32
|
||||
val string
|
||||
dist float64
|
||||
}
|
||||
|
||||
func findViaRPC(s *ValkeyStore, query []float32, topK int) ([][]float32, [][]byte, []float32, error) {
|
||||
res, err := s.StoresFind(&pb.StoresFindOptions{Key: &pb.StoresKey{Floats: query}, TopK: int32(topK)})
|
||||
if err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
return store.UnwrapKeys(res.Keys), store.UnwrapValues(res.Values), res.Similarities, nil
|
||||
}
|
||||
|
||||
// assertFTCreate returns a DoAndReturn func that verifies the FT.CREATE command
|
||||
// carries the expected dimension and algorithm, then replies OK.
|
||||
func assertFTCreate(dim int, algo string) func(context.Context, valkey.Completed) valkey.ValkeyResult {
|
||||
return func(_ context.Context, cmd valkey.Completed) valkey.ValkeyResult {
|
||||
toks := cmd.Commands()
|
||||
Expect(toks[0]).To(Equal("FT.CREATE"))
|
||||
Expect(toks).To(ContainElements("VECTOR", algo, "TYPE", "FLOAT32", "DIM", strconv.Itoa(dim), "DISTANCE_METRIC", "COSINE"))
|
||||
return mock.Result(mock.ValkeyString("OK"))
|
||||
}
|
||||
}
|
||||
|
||||
// ftInfoReply builds a RESP2-shaped FT.INFO reply that mirrors Valkey Search's
|
||||
// nesting: the VECTOR attribute carries its params under an `index` array whose
|
||||
// `dimensions` key holds the DIM. dimValue is the message the parser must read
|
||||
// back (integer or string-encoded), so both wire shapes can be exercised.
|
||||
func ftInfoReply(dimValue valkey.ValkeyMessage) valkey.ValkeyMessage {
|
||||
vectorAttr := mock.ValkeyArray(
|
||||
mock.ValkeyString("identifier"), mock.ValkeyString(_vecField),
|
||||
mock.ValkeyString("attribute"), mock.ValkeyString(_vecField),
|
||||
mock.ValkeyString("type"), mock.ValkeyString("VECTOR"),
|
||||
mock.ValkeyString("index"), mock.ValkeyArray(
|
||||
mock.ValkeyString("capacity"), mock.ValkeyInt64(1000),
|
||||
mock.ValkeyString("dimensions"), dimValue,
|
||||
mock.ValkeyString("distance_metric"), mock.ValkeyString("COSINE"),
|
||||
mock.ValkeyString("data_type"), mock.ValkeyString("FLOAT32"),
|
||||
),
|
||||
)
|
||||
return mock.ValkeyArray(
|
||||
mock.ValkeyString("index_name"), mock.ValkeyString("idx:test"),
|
||||
mock.ValkeyString("attributes"), mock.ValkeyArray(vectorAttr),
|
||||
mock.ValkeyString("num_docs"), mock.ValkeyInt64(0),
|
||||
)
|
||||
}
|
||||
|
||||
// ftSearchReply builds a RESP2-shaped FT.SEARCH reply: [total, key, attrs, ...]
|
||||
// where attrs carries the returned vec/val/__score fields.
|
||||
func ftSearchReply(prefix string, docs []ftDoc) valkey.ValkeyMessage {
|
||||
arr := []valkey.ValkeyMessage{mock.ValkeyInt64(int64(len(docs)))}
|
||||
for _, d := range docs {
|
||||
arr = append(arr, mock.ValkeyString(encodeKey(prefix, d.vec)))
|
||||
attrs := mock.ValkeyArray(
|
||||
mock.ValkeyString(_vecField), mock.ValkeyString(valkey.BinaryString(vecToBytes(d.vec))),
|
||||
mock.ValkeyString(_valField), mock.ValkeyString(d.val),
|
||||
mock.ValkeyString(_scoreField), mock.ValkeyString(strconv.FormatFloat(d.dist, 'f', -1, 64)),
|
||||
)
|
||||
arr = append(arr, attrs)
|
||||
}
|
||||
return mock.ValkeyArray(arr...)
|
||||
}
|
||||
|
||||
var _ = Describe("namespace token", func() {
|
||||
It("is stable for the same namespace (so a persisted index is found again)", func() {
|
||||
Expect(nsToken("faces")).To(Equal(nsToken("faces")))
|
||||
})
|
||||
|
||||
It("does not collide for namespaces that sanitize to the same token", func() {
|
||||
// "a b", "a/b" and "a:b" all sanitize to "a_b"; the hash suffix must
|
||||
// keep them distinct so two logically-distinct stores never share one
|
||||
// keyspace/index (the data-isolation guarantee).
|
||||
Expect(sanitize("a b")).To(Equal(sanitize("a/b")))
|
||||
Expect(nsToken("a b")).NotTo(Equal(nsToken("a/b")))
|
||||
Expect(nsToken("a/b")).NotTo(Equal(nsToken("a:b")))
|
||||
})
|
||||
|
||||
It("keeps the sanitized part human-readable", func() {
|
||||
Expect(nsToken("faces")).To(HavePrefix("faces-"))
|
||||
})
|
||||
|
||||
It("maps an empty namespace to a stable default token", func() {
|
||||
Expect(nsToken("")).To(HavePrefix("default-"))
|
||||
Expect(nsToken("")).To(Equal(nsToken("")))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("sanitize", func() {
|
||||
It("passes through allowed runes and folds the rest to '_'", func() {
|
||||
Expect(sanitize("Ok_9.-")).To(Equal("Ok_9.-"))
|
||||
Expect(sanitize("a b/c:d")).To(Equal("a_b_c_d"))
|
||||
})
|
||||
|
||||
It("maps empty to 'default'", func() {
|
||||
Expect(sanitize("")).To(Equal("default"))
|
||||
})
|
||||
})
|
||||
@@ -29,6 +29,29 @@ ABI v2) through purego:
|
||||
LocalAI's Go-side grammar-constrained tool calling; JSON-schema / regex /
|
||||
choice constraints are also exposed by the ABI.
|
||||
|
||||
## Hardware coverage
|
||||
|
||||
The CUDA builds require the CUDA 13 toolchain and target Blackwell only:
|
||||
`sm_120a` + `sm_121a` on x86_64, `sm_121a` (GB10 / DGX Spark) on arm64. CUDA
|
||||
12.x nvcc cannot compile the Blackwell fp4 kernels, so no CUDA 12 variant is
|
||||
shipped and `backend/index.yaml` maps the `nvidia-cuda-12` /
|
||||
`nvidia-l4t-cuda-12` capabilities at the CPU build. Practically:
|
||||
|
||||
| Host | Installed build |
|
||||
|---|---|
|
||||
| x86_64 + CUDA 13 | `cuda13-vllm-cpp` |
|
||||
| DGX Spark / GB10 (JetPack 7, CUDA 13) | `nvidia-l4t-arm64-vllm-cpp` |
|
||||
| Jetson AGX Orin (sm_87, JetPack 6, CUDA 12) | `cpu-vllm-cpp` |
|
||||
| Apple Silicon | `metal-vllm-cpp` |
|
||||
| Anything else | `vulkan-vllm-cpp` or `cpu-vllm-cpp` |
|
||||
|
||||
The capability a host reports comes from `/run/localai/capability` inside the
|
||||
LocalAI container, which the image bakes in at build time (see `Dockerfile`).
|
||||
A DGX Spark running the CUDA 12 `-nvidia-l4t-arm64` image therefore reports
|
||||
`nvidia-l4t-cuda-12` and gets the CPU build; use the `-nvidia-l4t-arm64-cuda-13`
|
||||
image, or set `LOCALAI_FORCE_META_BACKEND_CAPABILITY=nvidia-l4t-cuda-13`, to
|
||||
get the GPU one.
|
||||
|
||||
Model config example:
|
||||
|
||||
```yaml
|
||||
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# whisper.cpp version
|
||||
WHISPER_REPO?=https://github.com/ggml-org/whisper.cpp
|
||||
WHISPER_CPP_VERSION?=97c56f1dc1d1100a9d859c865a20c82d22f823ed
|
||||
WHISPER_CPP_VERSION?=080bbbe85230f624f0b52127f1ae1218247989f9
|
||||
SO_TARGET?=libgowhisper.so
|
||||
|
||||
CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF
|
||||
|
||||
@@ -129,29 +129,6 @@
|
||||
nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-ds4"
|
||||
metal: "metal-ds4"
|
||||
metal-darwin-arm64: "metal-ds4"
|
||||
- &dllm
|
||||
name: "dllm"
|
||||
alias: "dllm"
|
||||
license: mit
|
||||
description: |
|
||||
mudler/dllm.cpp - DiffusionGemma block-diffusion LLM inference engine
|
||||
(C++/ggml, GGUF weights). Decodes whole token canvases per diffusion
|
||||
round instead of autoregressive sampling. Runs on CPU and NVIDIA CUDA 13
|
||||
(including Jetson/GB10 L4T targets).
|
||||
urls:
|
||||
- https://github.com/mudler/dllm.cpp
|
||||
tags:
|
||||
- text-to-text
|
||||
- LLM
|
||||
- gguf
|
||||
- diffusion
|
||||
- CPU
|
||||
- CUDA
|
||||
capabilities:
|
||||
default: "cpu-dllm"
|
||||
nvidia: "cuda13-dllm"
|
||||
nvidia-cuda-13: "cuda13-dllm"
|
||||
nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-dllm"
|
||||
- &whispercpp
|
||||
name: "whisper"
|
||||
alias: "whisper"
|
||||
@@ -190,6 +167,10 @@
|
||||
inference time. It loads Hugging Face safetensors and GGUF checkpoints, supports
|
||||
structured output (JSON schema / regex / choice / GBNF grammar) enforced in-engine,
|
||||
and runs on CPU, NVIDIA CUDA (Blackwell-family), Apple Metal and Vulkan.
|
||||
The CUDA builds require the CUDA 13 toolchain and target Blackwell only: sm_120a
|
||||
plus sm_121a on x86_64, and sm_121a (GB10 / DGX Spark) on arm64. Older NVIDIA
|
||||
hardware and CUDA 12 hosts - including Jetson AGX Orin (sm_87, JetPack 6) - run
|
||||
the CPU build instead.
|
||||
urls:
|
||||
- https://github.com/mudler/vllm.cpp
|
||||
tags:
|
||||
@@ -207,6 +188,12 @@
|
||||
nvidia-cuda-13: "cuda13-vllm-cpp"
|
||||
nvidia-l4t: "nvidia-l4t-arm64-vllm-cpp"
|
||||
nvidia-l4t-cuda-13: "nvidia-l4t-arm64-vllm-cpp"
|
||||
# No CUDA 12 variant exists: 12.x nvcc cannot compile the Blackwell fp4
|
||||
# kernels, so those hosts run the CPU build. Mapped explicitly rather than
|
||||
# left to the "default" catch-all so the fallback is visible here instead
|
||||
# of looking like an oversight.
|
||||
nvidia-cuda-12: "cpu-vllm-cpp"
|
||||
nvidia-l4t-cuda-12: "cpu-vllm-cpp"
|
||||
- !!merge <<: *vllm-cpp
|
||||
name: "vllm-cpp-development"
|
||||
capabilities:
|
||||
@@ -217,6 +204,8 @@
|
||||
nvidia-cuda-13: "cuda13-vllm-cpp-development"
|
||||
nvidia-l4t: "nvidia-l4t-arm64-vllm-cpp-development"
|
||||
nvidia-l4t-cuda-13: "nvidia-l4t-arm64-vllm-cpp-development"
|
||||
nvidia-cuda-12: "cpu-vllm-cpp-development"
|
||||
nvidia-l4t-cuda-12: "cpu-vllm-cpp-development"
|
||||
- &crispasr
|
||||
name: "crispasr"
|
||||
alias: "crispasr"
|
||||
@@ -448,32 +437,6 @@
|
||||
nvidia-cuda-12: "cuda12-stablediffusion-ggml"
|
||||
nvidia-l4t-cuda-12: "nvidia-l4t-arm64-stablediffusion-ggml"
|
||||
nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-stablediffusion-ggml"
|
||||
- &trellis2cpp
|
||||
name: "trellis2cpp"
|
||||
alias: "trellis2cpp"
|
||||
license: mit
|
||||
description: |
|
||||
TRELLIS.2 image-to-3D generation (GLB meshes with PBR textures) in C++/ggml
|
||||
urls:
|
||||
- https://github.com/localai-org/trellis2cpp
|
||||
- https://github.com/microsoft/TRELLIS.2
|
||||
tags:
|
||||
- image-to-3d
|
||||
- 3d-generation
|
||||
- CPU
|
||||
- GPU
|
||||
- CUDA
|
||||
- Metal
|
||||
capabilities:
|
||||
default: "cpu-trellis2cpp"
|
||||
nvidia: "cuda12-trellis2cpp"
|
||||
vulkan: "vulkan-trellis2cpp"
|
||||
nvidia-l4t: "nvidia-l4t-arm64-trellis2cpp"
|
||||
metal: "metal-trellis2cpp"
|
||||
nvidia-cuda-13: "cuda13-trellis2cpp"
|
||||
nvidia-cuda-12: "cuda12-trellis2cpp"
|
||||
nvidia-l4t-cuda-12: "nvidia-l4t-arm64-trellis2cpp"
|
||||
nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-trellis2cpp"
|
||||
- &rfdetr
|
||||
name: "rfdetr"
|
||||
alias: "rfdetr"
|
||||
@@ -1504,7 +1467,6 @@
|
||||
alias: "kokoro"
|
||||
name: "kokoro"
|
||||
capabilities:
|
||||
default: "cpu-kokoro"
|
||||
nvidia: "cuda12-kokoro"
|
||||
intel: "intel-kokoro"
|
||||
amd: "rocm-kokoro"
|
||||
@@ -1886,23 +1848,6 @@
|
||||
capabilities:
|
||||
default: "cpu-cloud-proxy"
|
||||
metal: "metal-cloud-proxy"
|
||||
- &valkey-store
|
||||
name: "valkey-store"
|
||||
urls:
|
||||
- https://github.com/mudler/LocalAI
|
||||
description: |
|
||||
Valkey Store is a Valkey Search (FT.*) backed vector store for LocalAI. It
|
||||
persists vectors across restarts and supports opt-in HNSW indexing. Requires
|
||||
a reachable Valkey Search server (valkey/valkey-bundle).
|
||||
tags:
|
||||
- vector-database
|
||||
- valkey
|
||||
- open-source
|
||||
- CPU
|
||||
license: MIT
|
||||
capabilities:
|
||||
default: "cpu-valkey-store"
|
||||
metal: "metal-valkey-store"
|
||||
- &kitten-tts
|
||||
name: "kitten-tts"
|
||||
urls:
|
||||
@@ -2024,13 +1969,6 @@
|
||||
nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-ds4-development"
|
||||
metal: "metal-ds4-development"
|
||||
metal-darwin-arm64: "metal-ds4-development"
|
||||
- !!merge <<: *dllm
|
||||
name: "dllm-development"
|
||||
capabilities:
|
||||
default: "cpu-dllm-development"
|
||||
nvidia: "cuda13-dllm-development"
|
||||
nvidia-cuda-13: "cuda13-dllm-development"
|
||||
nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-dllm-development"
|
||||
- !!merge <<: *stablediffusionggml
|
||||
name: "stablediffusion-ggml-development"
|
||||
capabilities:
|
||||
@@ -2045,18 +1983,6 @@
|
||||
nvidia-cuda-12: "cuda12-stablediffusion-ggml-development"
|
||||
nvidia-l4t-cuda-12: "nvidia-l4t-arm64-stablediffusion-ggml-development"
|
||||
nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-stablediffusion-ggml-development"
|
||||
- !!merge <<: *trellis2cpp
|
||||
name: "trellis2cpp-development"
|
||||
capabilities:
|
||||
default: "cpu-trellis2cpp-development"
|
||||
nvidia: "cuda12-trellis2cpp-development"
|
||||
vulkan: "vulkan-trellis2cpp-development"
|
||||
nvidia-l4t: "nvidia-l4t-arm64-trellis2cpp-development"
|
||||
metal: "metal-trellis2cpp-development"
|
||||
nvidia-cuda-13: "cuda13-trellis2cpp-development"
|
||||
nvidia-cuda-12: "cuda12-trellis2cpp-development"
|
||||
nvidia-l4t-cuda-12: "nvidia-l4t-arm64-trellis2cpp-development"
|
||||
nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-trellis2cpp-development"
|
||||
- !!merge <<: *neutts
|
||||
name: "cpu-neutts"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-neutts"
|
||||
@@ -2454,35 +2380,6 @@
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-cloud-proxy"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-metal-darwin-arm64-cloud-proxy
|
||||
- !!merge <<: *valkey-store
|
||||
name: "cpu-valkey-store"
|
||||
alias: "valkey-store"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-valkey-store"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-cpu-valkey-store
|
||||
- !!merge <<: *valkey-store
|
||||
name: "cpu-valkey-store-development"
|
||||
alias: "valkey-store"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-cpu-valkey-store"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-cpu-valkey-store
|
||||
- !!merge <<: *valkey-store
|
||||
name: "valkey-store-development"
|
||||
alias: "valkey-store"
|
||||
capabilities:
|
||||
default: "cpu-valkey-store-development"
|
||||
metal: "metal-valkey-store-development"
|
||||
- !!merge <<: *valkey-store
|
||||
name: "metal-valkey-store"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-metal-darwin-arm64-valkey-store"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-metal-darwin-arm64-valkey-store
|
||||
- !!merge <<: *valkey-store
|
||||
name: "metal-valkey-store-development"
|
||||
alias: "valkey-store"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-valkey-store"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-metal-darwin-arm64-valkey-store
|
||||
- !!merge <<: *opus
|
||||
name: "cpu-opus"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-opus"
|
||||
@@ -2888,37 +2785,6 @@
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-ds4"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-metal-darwin-arm64-ds4
|
||||
## dllm
|
||||
- !!merge <<: *dllm
|
||||
name: "cpu-dllm"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-dllm"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-cpu-dllm
|
||||
- !!merge <<: *dllm
|
||||
name: "cpu-dllm-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-cpu-dllm"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-cpu-dllm
|
||||
- !!merge <<: *dllm
|
||||
name: "cuda13-dllm"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-13-dllm"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-nvidia-cuda-13-dllm
|
||||
- !!merge <<: *dllm
|
||||
name: "cuda13-dllm-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-13-dllm"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-cuda-13-dllm
|
||||
- !!merge <<: *dllm
|
||||
name: "cuda13-nvidia-l4t-arm64-dllm"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-nvidia-l4t-cuda-13-arm64-dllm"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-nvidia-l4t-cuda-13-arm64-dllm
|
||||
- !!merge <<: *dllm
|
||||
name: "cuda13-nvidia-l4t-arm64-dllm-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-nvidia-l4t-cuda-13-arm64-dllm"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-nvidia-l4t-cuda-13-arm64-dllm
|
||||
## whisper
|
||||
- !!merge <<: *whispercpp
|
||||
name: "whisper-development"
|
||||
@@ -3833,77 +3699,6 @@
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-13-stablediffusion-ggml"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-cuda-13-stablediffusion-ggml
|
||||
## trellis2cpp
|
||||
- !!merge <<: *trellis2cpp
|
||||
name: "cpu-trellis2cpp"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-trellis2cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-cpu-trellis2cpp
|
||||
- !!merge <<: *trellis2cpp
|
||||
name: "cpu-trellis2cpp-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-cpu-trellis2cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-cpu-trellis2cpp
|
||||
- !!merge <<: *trellis2cpp
|
||||
name: "metal-trellis2cpp"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-metal-darwin-arm64-trellis2cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-metal-darwin-arm64-trellis2cpp
|
||||
- !!merge <<: *trellis2cpp
|
||||
name: "metal-trellis2cpp-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-trellis2cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-metal-darwin-arm64-trellis2cpp
|
||||
- !!merge <<: *trellis2cpp
|
||||
name: "vulkan-trellis2cpp"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-vulkan-trellis2cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-vulkan-trellis2cpp
|
||||
- !!merge <<: *trellis2cpp
|
||||
name: "vulkan-trellis2cpp-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-vulkan-trellis2cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-vulkan-trellis2cpp
|
||||
- !!merge <<: *trellis2cpp
|
||||
name: "cuda12-trellis2cpp"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-12-trellis2cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-nvidia-cuda-12-trellis2cpp
|
||||
- !!merge <<: *trellis2cpp
|
||||
name: "cuda12-trellis2cpp-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-12-trellis2cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-cuda-12-trellis2cpp
|
||||
- !!merge <<: *trellis2cpp
|
||||
name: "cuda13-trellis2cpp"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-13-trellis2cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-nvidia-cuda-13-trellis2cpp
|
||||
- !!merge <<: *trellis2cpp
|
||||
name: "cuda13-trellis2cpp-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-13-trellis2cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-cuda-13-trellis2cpp
|
||||
- !!merge <<: *trellis2cpp
|
||||
name: "nvidia-l4t-arm64-trellis2cpp"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-nvidia-l4t-arm64-trellis2cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-nvidia-l4t-arm64-trellis2cpp
|
||||
- !!merge <<: *trellis2cpp
|
||||
name: "nvidia-l4t-arm64-trellis2cpp-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-nvidia-l4t-arm64-trellis2cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-nvidia-l4t-arm64-trellis2cpp
|
||||
- !!merge <<: *trellis2cpp
|
||||
name: "cuda13-nvidia-l4t-arm64-trellis2cpp"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-nvidia-l4t-cuda-13-arm64-trellis2cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-nvidia-l4t-cuda-13-arm64-trellis2cpp
|
||||
- !!merge <<: *trellis2cpp
|
||||
name: "cuda13-nvidia-l4t-arm64-trellis2cpp-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-nvidia-l4t-cuda-13-arm64-trellis2cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-nvidia-l4t-cuda-13-arm64-trellis2cpp
|
||||
## privacy-filter
|
||||
- !!merge <<: *privacyfilter
|
||||
name: "cpu-privacy-filter"
|
||||
@@ -5403,22 +5198,11 @@
|
||||
- !!merge <<: *kokoro
|
||||
name: "kokoro-development"
|
||||
capabilities:
|
||||
default: "cpu-kokoro-development"
|
||||
nvidia: "cuda12-kokoro-development"
|
||||
intel: "intel-kokoro-development"
|
||||
amd: "rocm-kokoro-development"
|
||||
nvidia-l4t: "nvidia-l4t-kokoro-development"
|
||||
metal: "metal-kokoro-development"
|
||||
- !!merge <<: *kokoro
|
||||
name: "cpu-kokoro"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-kokoro"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-cpu-kokoro
|
||||
- !!merge <<: *kokoro
|
||||
name: "cpu-kokoro-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-cpu-kokoro"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-cpu-kokoro
|
||||
- !!merge <<: *kokoro
|
||||
name: "cuda12-kokoro-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-12-kokoro"
|
||||
|
||||
@@ -1,3 +1 @@
|
||||
git+https://github.com/Blaizzy/mlx-vlm@v0.4.4
|
||||
torch
|
||||
torchvision
|
||||
git+https://github.com/Blaizzy/mlx-vlm@v0.4.4
|
||||
@@ -14,28 +14,4 @@ if [ "x${BUILD_PROFILE}" == "xintel" ]; then
|
||||
EXTRA_PIP_INSTALL_FLAGS+=" --upgrade --index-strategy=unsafe-first-match"
|
||||
fi
|
||||
|
||||
# Darwin needs a newer interpreter than libbackend's 3.10 default. nemo_toolkit
|
||||
# pulls in text2num, a Rust extension built with maturin, and its macOS arm64
|
||||
# wheels start at cp311 (3.0.2 publishes cp311/cp312/cp313/cp314 and no cp310).
|
||||
# On 3.10 pip therefore falls back to the sdist and dies in the PEP 517 hook
|
||||
# with "No module named 'maturin'", since EXTRA_PIP_INSTALL_FLAGS carries
|
||||
# --no-build-isolation and nothing installs the build backend. Moving to 3.12
|
||||
# takes the prebuilt wheel and needs no Rust toolchain on the runner at all.
|
||||
#
|
||||
# Darwin only, deliberately: the Linux profiles resolve a cp310 manylinux wheel
|
||||
# for the same package and have no reason to move.
|
||||
if [ "x${BUILD_PROFILE}" == "xmps" ] || [ "x${BUILD_PROFILE}" == "xmetal" ]; then
|
||||
PYTHON_VERSION="3.12"
|
||||
# PYTHON_PATCH must move with it. libbackend builds the portable-Python URL
|
||||
# as cpython-${PYTHON_VERSION}.${PYTHON_PATCH}+${PY_STANDALONE_TAG}-..., and
|
||||
# the default patch is 18 for 3.10.18; leaving it alone asks for a 3.12.18
|
||||
# that was never released and the download 404s.
|
||||
#
|
||||
# 11, not the 12 that sglang/install.sh uses for l4t13: at the 20250818 tag
|
||||
# python-build-standalone published 3.12.12 for linux aarch64 but not for
|
||||
# aarch64-apple-darwin, where 3.12.11 is the newest. Verified against the
|
||||
# release assets rather than copied across.
|
||||
PYTHON_PATCH="11"
|
||||
fi
|
||||
|
||||
installRequirements
|
||||
|
||||
@@ -2,16 +2,3 @@
|
||||
# (FunctionCallParser, ReasoningParser) move between releases.
|
||||
# 0.5.11 is the floor for Gemma 4 support (PR sgl-project/sglang#21952).
|
||||
sglang[all]>=0.5.11
|
||||
|
||||
# Keep nvidia-modelopt on a stable release. sglang[all] pulls it in through its
|
||||
# `diffusion` extra with no version bound of its own, and install.sh passes a
|
||||
# GLOBAL --prerelease=allow (needed because flash-attn-4 only ships 4.0.0b*
|
||||
# wheels). Unbounded plus prereleases-allowed resolves to 0.46.0rc0, whose build
|
||||
# backend imports wheel_stub without declaring it as a build dependency; with
|
||||
# --no-build-isolation also in EXTRA_PIP_INSTALL_FLAGS nothing installs it, and
|
||||
# every cublas sglang image fails with "No module named 'wheel_stub'".
|
||||
#
|
||||
# Bounding this one package rather than dropping the global flag: the flag is
|
||||
# load-bearing for flash-attn-4, and this is the narrower change. Raise the
|
||||
# bound once 0.46.0 final ships.
|
||||
nvidia-modelopt<0.46
|
||||
|
||||
@@ -2,16 +2,3 @@
|
||||
# (FunctionCallParser, ReasoningParser) move between releases.
|
||||
# 0.5.11 is the floor for Gemma 4 support (PR sgl-project/sglang#21952).
|
||||
sglang[all]>=0.5.11
|
||||
|
||||
# Keep nvidia-modelopt on a stable release. sglang[all] pulls it in through its
|
||||
# `diffusion` extra with no version bound of its own, and install.sh passes a
|
||||
# GLOBAL --prerelease=allow (needed because flash-attn-4 only ships 4.0.0b*
|
||||
# wheels). Unbounded plus prereleases-allowed resolves to 0.46.0rc0, whose build
|
||||
# backend imports wheel_stub without declaring it as a build dependency; with
|
||||
# --no-build-isolation also in EXTRA_PIP_INSTALL_FLAGS nothing installs it, and
|
||||
# every cublas sglang image fails with "No module named 'wheel_stub'".
|
||||
#
|
||||
# Bounding this one package rather than dropping the global flag: the flag is
|
||||
# load-bearing for flash-attn-4, and this is the narrower change. Raise the
|
||||
# bound once 0.46.0 final ships.
|
||||
nvidia-modelopt<0.46
|
||||
|
||||
@@ -334,13 +334,6 @@ impl Backend for KokorosService {
|
||||
Err(Status::unimplemented("Not supported"))
|
||||
}
|
||||
|
||||
async fn generate3_d(
|
||||
&self,
|
||||
_: Request<backend::Generate3DRequest>,
|
||||
) -> Result<Response<backend::Result>, Status> {
|
||||
Err(Status::unimplemented("Not supported"))
|
||||
}
|
||||
|
||||
async fn audio_transcription(
|
||||
&self,
|
||||
_: Request<backend::TranscriptRequest>,
|
||||
|
||||
@@ -165,7 +165,7 @@ func newApplication(appConfig *config.ApplicationConfig) *Application {
|
||||
voiceStoreName = "localai-voice-biometrics"
|
||||
)
|
||||
faceStoreResolver := func(_ context.Context, storeName string) (pkggrpc.Backend, error) {
|
||||
return corebackend.StoreBackend(ml, appConfig, app.backendLoader, storeName, "")
|
||||
return corebackend.StoreBackend(ml, appConfig, storeName, "")
|
||||
}
|
||||
app.faceRegistry = facerecognition.NewStoreRegistry(faceStoreResolver, faceStoreName, faceEmbeddingDim)
|
||||
|
||||
@@ -173,7 +173,7 @@ func newApplication(appConfig *config.ApplicationConfig) *Application {
|
||||
// namespace so embedding spaces stay isolated (a face vector and a
|
||||
// speaker vector are not comparable and differ in dimensionality).
|
||||
voiceStoreResolver := func(_ context.Context, storeName string) (pkggrpc.Backend, error) {
|
||||
return corebackend.StoreBackend(ml, appConfig, app.backendLoader, storeName, "")
|
||||
return corebackend.StoreBackend(ml, appConfig, storeName, "")
|
||||
}
|
||||
app.voiceRegistry = voicerecognition.NewStoreRegistry(voiceStoreResolver, voiceStoreName, voiceEmbeddingDim)
|
||||
|
||||
|
||||
@@ -46,12 +46,12 @@ type lazyScorer struct {
|
||||
modelName string
|
||||
}
|
||||
|
||||
func (l *lazyScorer) Score(ctx context.Context, prompt string, stablePrefixLen int, candidates []string) ([]backend.CandidateScore, error) {
|
||||
func (l *lazyScorer) Score(ctx context.Context, prompt string, candidates []string) ([]backend.CandidateScore, error) {
|
||||
cfg := l.app.adapterConfig(l.modelName)
|
||||
if cfg == nil {
|
||||
return nil, fmt.Errorf("scorer: model %q no longer available", l.modelName)
|
||||
}
|
||||
return backend.NewScorer(l.app.modelLoader, *cfg, l.app.applicationConfig).Score(ctx, prompt, stablePrefixLen, candidates)
|
||||
return backend.NewScorer(l.app.modelLoader, *cfg, l.app.applicationConfig).Score(ctx, prompt, candidates)
|
||||
}
|
||||
|
||||
// TokenCounter returns a func so the middleware's literal field type accepts
|
||||
@@ -116,5 +116,5 @@ func (l *lazyEmbedder) Embed(ctx context.Context, text string) ([]float32, error
|
||||
// VectorStore takes a store name, not a model name — no adapterConfig, no
|
||||
// staleness to avoid.
|
||||
func (a *Application) VectorStore(storeName string) backend.VectorStore {
|
||||
return backend.NewVectorStore(a.modelLoader, a.applicationConfig, a.backendLoader, storeName)
|
||||
return backend.NewVectorStore(a.modelLoader, a.applicationConfig, storeName)
|
||||
}
|
||||
|
||||
@@ -109,7 +109,7 @@ var _ = Describe("router_factories lazy config resolution", func() {
|
||||
Expect(lazy.modelName).To(Equal("score-test"))
|
||||
|
||||
removeCfg("score-test")
|
||||
_, err := sc.Score(context.Background(), "prompt", 0, []string{"a"})
|
||||
_, err := sc.Score(context.Background(), "prompt", []string{"a"})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("no longer available"))
|
||||
})
|
||||
|
||||
@@ -1,104 +0,0 @@
|
||||
package backend
|
||||
|
||||
import (
|
||||
"maps"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/trace"
|
||||
"github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
model "github.com/mudler/LocalAI/pkg/model"
|
||||
)
|
||||
|
||||
// Model3DGenerationOptions is the backend-neutral request passed to 3D
|
||||
// generators. Image contains a staged local path by the time it reaches
|
||||
// this layer.
|
||||
type Model3DGenerationOptions struct {
|
||||
Image string
|
||||
Destination string
|
||||
Seed int32
|
||||
Step int32
|
||||
CFGScale float32
|
||||
TextureSteps int32
|
||||
Quality string
|
||||
Background string
|
||||
Params map[string]string
|
||||
}
|
||||
|
||||
func Model3DGeneration(options Model3DGenerationOptions, loader *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig) (func() error, error) {
|
||||
opts := ModelOptions(modelConfig, appConfig)
|
||||
inferenceModel, err := loader.Load(opts...)
|
||||
if err != nil {
|
||||
recordModelLoadFailure(appConfig, modelConfig.Name, modelConfig.Backend, err, nil)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
fn := func() error {
|
||||
_, err := inferenceModel.Generate3D(
|
||||
appConfig.Context,
|
||||
&proto.Generate3DRequest{
|
||||
Src: options.Image,
|
||||
Dst: options.Destination,
|
||||
Seed: options.Seed,
|
||||
Step: options.Step,
|
||||
CfgScale: options.CFGScale,
|
||||
TextureSteps: options.TextureSteps,
|
||||
Quality: options.Quality,
|
||||
Background: options.Background,
|
||||
Params: maps.Clone(options.Params),
|
||||
},
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
if appConfig.EnableTracing {
|
||||
trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes)
|
||||
|
||||
traceType := trace.BackendTrace3DGeneration
|
||||
traceSummary := "3d: " + options.Quality
|
||||
traceData := map[string]any{}
|
||||
if options.Params["operation"] == "print_remesh" {
|
||||
traceType = trace.BackendTrace3DRemesh
|
||||
traceSummary = "3d: remesh"
|
||||
traceData["detail_percent"] = options.Params["detail_percent"]
|
||||
traceData["has_mesh"] = options.Image != ""
|
||||
} else {
|
||||
traceData = map[string]any{
|
||||
"seed": options.Seed,
|
||||
"step": options.Step,
|
||||
"cfg_scale": options.CFGScale,
|
||||
"texture_steps": options.TextureSteps,
|
||||
"quality": options.Quality,
|
||||
"background": options.Background,
|
||||
"has_image": options.Image != "",
|
||||
}
|
||||
}
|
||||
|
||||
startTime := time.Now()
|
||||
originalFn := fn
|
||||
fn = func() error {
|
||||
err := originalFn()
|
||||
duration := time.Since(startTime)
|
||||
|
||||
errStr := ""
|
||||
if err != nil {
|
||||
errStr = err.Error()
|
||||
}
|
||||
|
||||
trace.RecordBackendTrace(trace.BackendTrace{
|
||||
Timestamp: startTime,
|
||||
Duration: duration,
|
||||
Type: traceType,
|
||||
ModelName: modelConfig.Name,
|
||||
Backend: modelConfig.Backend,
|
||||
Summary: trace.TruncateString(traceSummary, 200),
|
||||
Error: errStr,
|
||||
Data: traceData,
|
||||
})
|
||||
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return fn, nil
|
||||
}
|
||||
@@ -166,21 +166,6 @@ func estimateModelSizeBytes(c config.ModelConfig, modelsPath string) int64 {
|
||||
return int64(result.SizeBytes)
|
||||
}
|
||||
|
||||
// effectiveThreads resolves the thread count a backend is asked to use.
|
||||
// Per-model threads wins: SetDefaults already fills an unset per-model value
|
||||
// from the app-level --threads, so overriding a set value with the app value
|
||||
// here would make the YAML `threads:` knob dead config (it did, for years —
|
||||
// e.g. a tiny VAD model could never opt down from the global pool size).
|
||||
func effectiveThreads(c config.ModelConfig, appThreads int) int {
|
||||
if c.Threads != nil && *c.Threads > 0 {
|
||||
return *c.Threads
|
||||
}
|
||||
if appThreads > 0 {
|
||||
return appThreads
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func ModelOptions(c config.ModelConfig, so *config.ApplicationConfig, opts ...model.Option) []model.Option {
|
||||
defOpts := []model.Option{
|
||||
model.WithBackendString(c.Backend),
|
||||
@@ -193,7 +178,16 @@ func ModelOptions(c config.ModelConfig, so *config.ApplicationConfig, opts ...mo
|
||||
defOpts = append(defOpts, model.WithModelFile(c.ModelFileName()))
|
||||
}
|
||||
|
||||
threads := effectiveThreads(c, so.Threads)
|
||||
threads := 1
|
||||
|
||||
if c.Threads != nil {
|
||||
threads = *c.Threads
|
||||
}
|
||||
|
||||
if so.Threads != 0 {
|
||||
threads = so.Threads
|
||||
}
|
||||
|
||||
c.Threads = &threads
|
||||
|
||||
grpcOpts := grpcModelOpts(c, so.SystemState.Model.ModelsPath)
|
||||
@@ -422,7 +416,6 @@ func grpcModelOpts(c config.ModelConfig, modelPath string) *pb.ModelOptions {
|
||||
Options: withCompanionArtifactOptions(c.Options, c.Artifacts),
|
||||
Overrides: c.Overrides,
|
||||
EngineArgs: engineArgsJSON,
|
||||
EnableScore: c.HasUsecases(config.FLAG_SCORE),
|
||||
CLIPSkip: int32(c.Diffusers.ClipSkip),
|
||||
ControlNet: c.Diffusers.ControlNet,
|
||||
ContextSize: int32(ctxSize),
|
||||
@@ -482,7 +475,6 @@ func grpcModelOpts(c config.ModelConfig, modelPath string) *pb.ModelOptions {
|
||||
ApiKeyFile: c.Proxy.APIKeyFile,
|
||||
UpstreamModel: c.Proxy.UpstreamModel,
|
||||
RequestTimeoutSeconds: int32(c.Proxy.RequestTimeoutSeconds),
|
||||
CachePrompt: c.Proxy.CachePrompt,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -120,7 +120,6 @@ var _ = Describe("grpcModelOpts NBatch", func() {
|
||||
cfg := config.ModelConfig{Threads: &threads, LLMConfig: config.LLMConfig{ContextSize: &ctx}}
|
||||
opts := grpcModelOpts(cfg, "/tmp/models")
|
||||
Expect(opts.NBatch).To(BeEquivalentTo(512))
|
||||
Expect(opts.EnableScore).To(BeFalse())
|
||||
})
|
||||
|
||||
It("sizes the batch to the context window for score models", func() {
|
||||
@@ -129,14 +128,6 @@ var _ = Describe("grpcModelOpts NBatch", func() {
|
||||
cfg := config.ModelConfig{Threads: &threads, LLMConfig: config.LLMConfig{ContextSize: &ctx}, KnownUsecases: &scoreUsecase}
|
||||
opts := grpcModelOpts(cfg, "/tmp/models")
|
||||
Expect(opts.NBatch).To(BeEquivalentTo(4096))
|
||||
Expect(opts.EnableScore).To(BeTrue())
|
||||
})
|
||||
|
||||
It("enables score resources for a model with multiple usecases", func() {
|
||||
usecases := config.FLAG_CHAT | config.FLAG_SCORE
|
||||
cfg := config.ModelConfig{Threads: &threads, LLMConfig: config.LLMConfig{ContextSize: &ctx}, KnownUsecases: &usecases}
|
||||
opts := grpcModelOpts(cfg, "/tmp/models")
|
||||
Expect(opts.EnableScore).To(BeTrue())
|
||||
})
|
||||
|
||||
It("keeps an explicit batch over the score default", func() {
|
||||
@@ -364,23 +355,3 @@ var _ = Describe("gRPCPredictOpts model identity", func() {
|
||||
Expect(opts.ModelIdentity).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("effectiveThreads", func() {
|
||||
It("lets a per-model threads value override the app-level --threads", func() {
|
||||
one := 1
|
||||
cfg := config.ModelConfig{Threads: &one}
|
||||
Expect(effectiveThreads(cfg, 10)).To(Equal(1),
|
||||
"per-model threads is a real knob, not dead config under --threads")
|
||||
})
|
||||
|
||||
It("falls back to the app-level threads when the model sets none", func() {
|
||||
Expect(effectiveThreads(config.ModelConfig{}, 10)).To(Equal(10))
|
||||
zero := 0
|
||||
Expect(effectiveThreads(config.ModelConfig{Threads: &zero}, 10)).To(Equal(10),
|
||||
"an explicit threads: 0 means unset, not zero threads")
|
||||
})
|
||||
|
||||
It("never resolves to a non-positive thread count", func() {
|
||||
Expect(effectiveThreads(config.ModelConfig{}, 0)).To(Equal(1))
|
||||
})
|
||||
})
|
||||
|
||||
@@ -28,7 +28,7 @@ func PreloadModelByName(ctx context.Context, cl *config.ModelConfigLoader, ml *m
|
||||
return nil, err
|
||||
}
|
||||
|
||||
stages, err := pipelineStages(cl, &cfg.Pipeline, ml.ModelPath, appConfig.ToConfigLoaderOptions()...)
|
||||
stages, err := pipelineStages(cl, &cfg.Pipeline, ml.ModelPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -59,7 +59,7 @@ var loadStage = PreloadModel
|
||||
// pipeline itself uses. A stage that fails to resolve is a misconfiguration,
|
||||
// so it fails fast rather than being deferred to load. A pipeline with no
|
||||
// stages set returns nil, which callers treat as "not a pipeline".
|
||||
func pipelineStages(cl *config.ModelConfigLoader, p *config.Pipeline, modelPath string, opts ...config.ConfigLoaderOption) ([]PreloadStage, error) {
|
||||
func pipelineStages(cl *config.ModelConfigLoader, p *config.Pipeline, modelPath string) ([]PreloadStage, error) {
|
||||
voiceRec := ""
|
||||
if p.VoiceRecognition != nil {
|
||||
voiceRec = p.VoiceRecognition.Model
|
||||
@@ -76,7 +76,7 @@ func pipelineStages(cl *config.ModelConfigLoader, p *config.Pipeline, modelPath
|
||||
if s.name == "" {
|
||||
continue
|
||||
}
|
||||
cfg, err := cl.LoadResolvedModelConfig(s.name, modelPath, opts...)
|
||||
cfg, err := cl.LoadResolvedModelConfig(s.name, modelPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%s (%s): %w", s.role, s.name, err)
|
||||
}
|
||||
@@ -87,11 +87,9 @@ func pipelineStages(cl *config.ModelConfigLoader, p *config.Pipeline, modelPath
|
||||
|
||||
// PreloadStages loads every present stage at once and waits for all of them, so
|
||||
// a pipeline warms in the time of its slowest stage rather than the sum. Absent
|
||||
// stages are skipped. Some callers represent an unset optional stage with a
|
||||
// nil config, while others materialize a default config with an empty name. A
|
||||
// failed stage does not cancel the others — they all run to completion so the
|
||||
// joined error names every broken stage at once, alongside the names that did
|
||||
// load.
|
||||
// (nil-config) stages are skipped. A failed stage does not cancel the others —
|
||||
// they all run to completion so the joined error names every broken stage at
|
||||
// once, alongside the names that did load.
|
||||
func PreloadStages(ctx context.Context, ml *model.ModelLoader, appConfig *config.ApplicationConfig, stages []PreloadStage) ([]string, error) {
|
||||
var (
|
||||
wg sync.WaitGroup
|
||||
@@ -100,7 +98,7 @@ func PreloadStages(ctx context.Context, ml *model.ModelLoader, appConfig *config
|
||||
errs []error
|
||||
)
|
||||
for _, s := range stages {
|
||||
if s.Cfg == nil || s.Cfg.Name == "" {
|
||||
if s.Cfg == nil {
|
||||
continue
|
||||
}
|
||||
wg.Add(1)
|
||||
|
||||
@@ -103,13 +103,12 @@ var _ = Describe("PreloadStages", func() {
|
||||
return PreloadStage{Role: role, Cfg: &config.ModelConfig{Name: name}}
|
||||
}
|
||||
|
||||
It("loads every present stage, skips absent stages, and returns the loaded names", func() {
|
||||
It("loads every present stage, skips absent (nil-config) ones, and returns the loaded names", func() {
|
||||
stubLoader(nil)
|
||||
|
||||
loaded, err := PreloadStages(context.Background(), nil, nil, []PreloadStage{
|
||||
mkStage("vad", "vad-m"),
|
||||
{Role: "transcription"},
|
||||
mkStage("tts", ""),
|
||||
{Role: "transcription"}, // absent stage
|
||||
mkStage("llm", "llm-m"),
|
||||
})
|
||||
|
||||
|
||||
@@ -23,10 +23,6 @@ type ScoreOptions struct {
|
||||
// token count. Useful when comparing candidates of different
|
||||
// lengths — without it, longer candidates score lower by default.
|
||||
LengthNormalize bool
|
||||
// StablePrefixLen is the byte length of the prompt prefix that stays
|
||||
// identical across repeated scoring calls (0 = unknown); forwarded to
|
||||
// the backend as a state-reuse boundary hint.
|
||||
StablePrefixLen int
|
||||
}
|
||||
|
||||
// CandidateScore is the per-candidate result. Mirrors pb.CandidateScore
|
||||
@@ -46,13 +42,9 @@ type TokenLogProb struct {
|
||||
// Scorer evaluates a model's joint log-probability of each candidate
|
||||
// continuation given a shared prompt. Implemented by NewScorer over a
|
||||
// model-loaded backend; the router's score classifier consumes this
|
||||
// for multi-label policy selection. stablePrefixLen is the byte length
|
||||
// of the prompt prefix that stays identical across calls (0 = unknown)
|
||||
// — backends use it to place a state-reuse point at the boundary, which
|
||||
// is what keeps repeat scoring fast on models that cannot rewind
|
||||
// (hybrid/recurrent architectures).
|
||||
// for multi-label policy selection.
|
||||
type Scorer interface {
|
||||
Score(ctx context.Context, prompt string, stablePrefixLen int, candidates []string) ([]CandidateScore, error)
|
||||
Score(ctx context.Context, prompt string, candidates []string) ([]CandidateScore, error)
|
||||
}
|
||||
|
||||
// NewScorer binds (loader, modelConfig, appConfig) into a Scorer. The
|
||||
@@ -69,8 +61,8 @@ type modelScorer struct {
|
||||
appConfig *config.ApplicationConfig
|
||||
}
|
||||
|
||||
func (m *modelScorer) Score(ctx context.Context, prompt string, stablePrefixLen int, candidates []string) ([]CandidateScore, error) {
|
||||
fn, err := ModelScore(prompt, candidates, ScoreOptions{LengthNormalize: true, StablePrefixLen: stablePrefixLen}, m.loader, m.modelConfig, m.appConfig)
|
||||
func (m *modelScorer) Score(ctx context.Context, prompt string, candidates []string) ([]CandidateScore, error) {
|
||||
fn, err := ModelScore(prompt, candidates, ScoreOptions{LengthNormalize: true}, m.loader, m.modelConfig, m.appConfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -111,7 +103,6 @@ func ModelScore(prompt string, candidates []string, opts ScoreOptions, loader *m
|
||||
Candidates: candidates,
|
||||
IncludeTokenLogprobs: opts.IncludeTokenLogprobs,
|
||||
LengthNormalize: opts.LengthNormalize,
|
||||
StablePrefixLen: int32(opts.StablePrefixLen),
|
||||
})
|
||||
results := scoreResponseToCandidates(resp, opts.IncludeTokenLogprobs)
|
||||
if appConfig.EnableTracing {
|
||||
|
||||
@@ -9,7 +9,6 @@ import (
|
||||
"github.com/mudler/LocalAI/core/trace"
|
||||
|
||||
"github.com/mudler/LocalAI/pkg/grpc"
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/LocalAI/pkg/model"
|
||||
"github.com/mudler/LocalAI/pkg/store"
|
||||
)
|
||||
@@ -24,25 +23,21 @@ type VectorStore interface {
|
||||
|
||||
// NewVectorStore returns a VectorStore backed by the local-store
|
||||
// gRPC backend, namespaced by storeName so two routers don't collide.
|
||||
// cl resolves the per-store model config (backend + options); it may be nil,
|
||||
// in which case the store falls back to the default backend and its built-in
|
||||
// defaults.
|
||||
func NewVectorStore(loader *model.ModelLoader, appConfig *config.ApplicationConfig, cl *config.ModelConfigLoader, storeName string) VectorStore {
|
||||
func NewVectorStore(loader *model.ModelLoader, appConfig *config.ApplicationConfig, storeName string) VectorStore {
|
||||
if storeName == "" {
|
||||
return nil
|
||||
}
|
||||
return &localVectorStore{loader: loader, appConfig: appConfig, cl: cl, storeName: storeName}
|
||||
return &localVectorStore{loader: loader, appConfig: appConfig, storeName: storeName}
|
||||
}
|
||||
|
||||
type localVectorStore struct {
|
||||
loader *model.ModelLoader
|
||||
appConfig *config.ApplicationConfig
|
||||
cl *config.ModelConfigLoader
|
||||
storeName string
|
||||
}
|
||||
|
||||
func (s *localVectorStore) backend(_ context.Context) (grpc.Backend, error) {
|
||||
return StoreBackend(s.loader, s.appConfig, s.cl, s.storeName, "")
|
||||
return StoreBackend(s.loader, s.appConfig, s.storeName, "")
|
||||
}
|
||||
|
||||
func (s *localVectorStore) Search(ctx context.Context, vec []float32) (sim float64, payload []byte, ok bool, err error) {
|
||||
@@ -126,24 +121,7 @@ func (s *localVectorStore) recordTrace(start time.Time, op string, vecDim int, s
|
||||
})
|
||||
}
|
||||
|
||||
func StoreBackend(sl *model.ModelLoader, appConfig *config.ApplicationConfig, cl *config.ModelConfigLoader, storeName string, backend string) (grpc.Backend, error) {
|
||||
// Resolve the per-store model config (keyed by the store namespace, which
|
||||
// is the model ID for a store). This is the LocalAI-native config surface:
|
||||
// a store's backend selection and its backend-specific settings live in a
|
||||
// model YAML's `backend:` and `options:` fields, so different stores can
|
||||
// point at different servers/indexes. When no config exists for the store,
|
||||
// we fall back to the default backend and let the backend apply its own
|
||||
// built-in defaults — preserving the zero-config experience.
|
||||
var loadOpts []string
|
||||
if cl != nil {
|
||||
if cfg, ok := cl.GetModelConfig(storeName); ok {
|
||||
if backend == "" {
|
||||
backend = cfg.Backend
|
||||
}
|
||||
loadOpts = cfg.Options
|
||||
}
|
||||
}
|
||||
|
||||
func StoreBackend(sl *model.ModelLoader, appConfig *config.ApplicationConfig, storeName string, backend string) (grpc.Backend, error) {
|
||||
if backend == "" {
|
||||
backend = model.LocalStoreBackend
|
||||
}
|
||||
@@ -167,12 +145,5 @@ func StoreBackend(sl *model.ModelLoader, appConfig *config.ApplicationConfig, cl
|
||||
model.WithModel(store.NamespacePrefix + storeName),
|
||||
}
|
||||
|
||||
// Thread the store's configured options through to the backend's LoadModel
|
||||
// via ModelOptions.Options (field 62). The loader clones these opts and
|
||||
// overrides only Model/ModelFile, so the namespace set above is preserved.
|
||||
if len(loadOpts) > 0 {
|
||||
sc = append(sc, model.WithLoadGRPCLoadModelOpts(&pb.ModelOptions{Options: loadOpts}))
|
||||
}
|
||||
|
||||
return sl.Load(sc...)
|
||||
}
|
||||
|
||||
@@ -243,23 +243,6 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
activatedListeners, err := systemdActivatedListeners()
|
||||
if err != nil {
|
||||
return fmt.Errorf("loading systemd socket activation listeners: %w", err)
|
||||
}
|
||||
activatedListener, err := selectSystemdListener(activatedListeners)
|
||||
if err != nil {
|
||||
for _, listener := range activatedListeners {
|
||||
_ = listener.Close()
|
||||
}
|
||||
return err
|
||||
}
|
||||
if activatedListener != nil {
|
||||
defer func() {
|
||||
_ = activatedListener.Close()
|
||||
}()
|
||||
}
|
||||
|
||||
os.MkdirAll(r.BackendsPath, 0750)
|
||||
os.MkdirAll(r.ModelsPath, 0750)
|
||||
|
||||
@@ -749,13 +732,8 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
|
||||
// LAN, or VPN that's the historical "trusted network" deployment, but on
|
||||
// a public IP it makes every model, gallery install, settings change, and
|
||||
// admin endpoint reachable by anyone who can connect to the port.
|
||||
listenAddress := r.Address
|
||||
if activatedListener != nil {
|
||||
listenAddress = activatedListener.Addr().String()
|
||||
}
|
||||
|
||||
authConfigured := app.AuthDB() != nil || len(r.APIKeys) > 0
|
||||
if err := requireAuthOrTrustedBind(listenAddress, authConfigured, r.AllowInsecurePublicBind); err != nil {
|
||||
if err := requireAuthOrTrustedBind(r.Address, authConfigured, r.AllowInsecurePublicBind); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -765,11 +743,7 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
|
||||
return err
|
||||
}
|
||||
|
||||
if activatedListener != nil {
|
||||
appHTTP.Listener = activatedListener
|
||||
xlog.Info("Using systemd socket activation listener", "address", listenAddress)
|
||||
}
|
||||
xlog.Info("LocalAI is started and running", "address", listenAddress)
|
||||
xlog.Info("LocalAI is started and running", "address", r.Address)
|
||||
|
||||
// Start P2P if token was provided via CLI/env or loaded from runtime_settings.json
|
||||
if token != "" || app.ApplicationConfig().P2PToken != "" {
|
||||
@@ -788,11 +762,11 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
|
||||
// backends like PostgreSQL need to call the embeddings API during
|
||||
// collection initialization.
|
||||
go func() {
|
||||
waitForServerReady(listenAddress, app.ApplicationConfig().Context)
|
||||
waitForServerReady(r.Address, app.ApplicationConfig().Context)
|
||||
app.StartAgentPool()
|
||||
}()
|
||||
|
||||
return appHTTP.Start(listenAddress)
|
||||
return appHTTP.Start(r.Address)
|
||||
}
|
||||
|
||||
// waitForServerReady polls the given address until the HTTP server is
|
||||
|
||||
@@ -1,17 +0,0 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
)
|
||||
|
||||
func selectSystemdListener(listeners []net.Listener) (net.Listener, error) {
|
||||
switch len(listeners) {
|
||||
case 0:
|
||||
return nil, nil
|
||||
case 1:
|
||||
return listeners[0], nil
|
||||
default:
|
||||
return nil, fmt.Errorf("systemd socket activation requires exactly one stream listener, got %d", len(listeners))
|
||||
}
|
||||
}
|
||||
@@ -1,71 +0,0 @@
|
||||
//go:build linux
|
||||
|
||||
package cli
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
const systemdListenFDStart = 3
|
||||
|
||||
func systemdActivatedListeners() ([]net.Listener, error) {
|
||||
listenPID := os.Getenv("LISTEN_PID")
|
||||
listenFDs := os.Getenv("LISTEN_FDS")
|
||||
if listenPID == "" && listenFDs == "" {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
defer func() {
|
||||
for _, key := range []string{"LISTEN_PID", "LISTEN_FDS", "LISTEN_FDNAMES"} {
|
||||
_ = os.Unsetenv(key)
|
||||
}
|
||||
}()
|
||||
|
||||
pid, err := strconv.Atoi(listenPID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid LISTEN_PID %q: %w", listenPID, err)
|
||||
}
|
||||
count, err := strconv.Atoi(listenFDs)
|
||||
if err != nil || count < 0 {
|
||||
return nil, fmt.Errorf("invalid LISTEN_FDS %q", listenFDs)
|
||||
}
|
||||
if pid != os.Getpid() || count == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
return listenersFromSystemdFDs(systemdListenFDStart, count)
|
||||
}
|
||||
|
||||
func listenersFromSystemdFDs(start, count int) (_ []net.Listener, err error) {
|
||||
listeners := make([]net.Listener, 0, count)
|
||||
defer func() {
|
||||
if err != nil {
|
||||
for _, listener := range listeners {
|
||||
_ = listener.Close()
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
for offset := range count {
|
||||
fd := uintptr(start + offset)
|
||||
file := os.NewFile(fd, fmt.Sprintf("LISTEN_FD_%d", fd))
|
||||
if file == nil {
|
||||
return nil, fmt.Errorf("opening systemd listener file descriptor %d", fd)
|
||||
}
|
||||
listener, listenerErr := net.FileListener(file)
|
||||
closeErr := file.Close()
|
||||
if listenerErr != nil {
|
||||
return nil, fmt.Errorf("using systemd file descriptor %d as a stream listener: %w", fd, listenerErr)
|
||||
}
|
||||
if closeErr != nil {
|
||||
_ = listener.Close()
|
||||
return nil, fmt.Errorf("closing inherited systemd file descriptor %d: %w", fd, closeErr)
|
||||
}
|
||||
listeners = append(listeners, listener)
|
||||
}
|
||||
|
||||
return listeners, nil
|
||||
}
|
||||
@@ -1,9 +0,0 @@
|
||||
//go:build !linux
|
||||
|
||||
package cli
|
||||
|
||||
import "net"
|
||||
|
||||
func systemdActivatedListeners() ([]net.Listener, error) {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -1,101 +0,0 @@
|
||||
//go:build linux
|
||||
|
||||
package cli
|
||||
|
||||
import (
|
||||
"net"
|
||||
"os"
|
||||
"strconv"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("selectSystemdListener", func() {
|
||||
It("keeps normal address binding when systemd passes no listener", func() {
|
||||
listener, err := selectSystemdListener(nil)
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(listener).To(BeNil())
|
||||
})
|
||||
|
||||
It("uses the single stream listener passed by systemd", func() {
|
||||
inherited, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(inherited.Close)
|
||||
|
||||
listener, err := selectSystemdListener([]net.Listener{inherited})
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(listener).To(BeIdenticalTo(inherited))
|
||||
})
|
||||
|
||||
It("rejects ambiguous activation with multiple stream listeners", func() {
|
||||
first, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(first.Close)
|
||||
second, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(second.Close)
|
||||
|
||||
listener, err := selectSystemdListener([]net.Listener{first, second})
|
||||
|
||||
Expect(err).To(MatchError(ContainSubstring("exactly one")))
|
||||
Expect(listener).To(BeNil())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("systemdActivatedListeners", func() {
|
||||
It("turns an inherited TCP file descriptor into a working listener", func() {
|
||||
original, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
file, err := original.(*net.TCPListener).File()
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(original.Close()).To(Succeed())
|
||||
|
||||
listeners, err := listenersFromSystemdFDs(int(file.Fd()), 1)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(listeners).To(HaveLen(1))
|
||||
DeferCleanup(listeners[0].Close)
|
||||
|
||||
client, err := net.Dial("tcp", listeners[0].Addr().String())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(client.Close)
|
||||
server, err := listeners[0].Accept()
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(server.Close()).To(Succeed())
|
||||
})
|
||||
|
||||
It("ignores descriptors intended for another process and clears the activation environment", func() {
|
||||
Expect(os.Setenv("LISTEN_PID", strconv.Itoa(os.Getpid()+1))).To(Succeed())
|
||||
Expect(os.Setenv("LISTEN_FDS", "1")).To(Succeed())
|
||||
Expect(os.Setenv("LISTEN_FDNAMES", "localai-http")).To(Succeed())
|
||||
DeferCleanup(func() {
|
||||
_ = os.Unsetenv("LISTEN_PID")
|
||||
_ = os.Unsetenv("LISTEN_FDS")
|
||||
_ = os.Unsetenv("LISTEN_FDNAMES")
|
||||
})
|
||||
|
||||
listeners, err := systemdActivatedListeners()
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(listeners).To(BeEmpty())
|
||||
Expect(os.Getenv("LISTEN_PID")).To(BeEmpty())
|
||||
Expect(os.Getenv("LISTEN_FDS")).To(BeEmpty())
|
||||
Expect(os.Getenv("LISTEN_FDNAMES")).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("reports malformed activation metadata instead of silently binding another socket", func() {
|
||||
Expect(os.Setenv("LISTEN_PID", strconv.Itoa(os.Getpid()))).To(Succeed())
|
||||
Expect(os.Setenv("LISTEN_FDS", "not-a-number")).To(Succeed())
|
||||
DeferCleanup(func() {
|
||||
_ = os.Unsetenv("LISTEN_PID")
|
||||
_ = os.Unsetenv("LISTEN_FDS")
|
||||
})
|
||||
|
||||
listeners, err := systemdActivatedListeners()
|
||||
|
||||
Expect(err).To(MatchError(ContainSubstring("LISTEN_FDS")))
|
||||
Expect(listeners).To(BeNil())
|
||||
})
|
||||
})
|
||||
@@ -16,7 +16,6 @@ const (
|
||||
UsecaseTokenize = "tokenize"
|
||||
UsecaseImage = "image"
|
||||
UsecaseVideo = "video"
|
||||
Usecase3D = "3d"
|
||||
UsecaseTranscript = "transcript"
|
||||
UsecaseTTS = "tts"
|
||||
UsecaseSoundGeneration = "sound_generation"
|
||||
@@ -31,7 +30,6 @@ const (
|
||||
UsecaseFaceRecognition = "face_recognition"
|
||||
UsecaseSpeakerRecognition = "speaker_recognition"
|
||||
UsecaseTokenClassify = "token_classify"
|
||||
UsecaseScore = "score"
|
||||
)
|
||||
|
||||
// GRPCMethod identifies a Backend service RPC from backend.proto.
|
||||
@@ -43,7 +41,6 @@ const (
|
||||
MethodEmbedding GRPCMethod = "Embedding"
|
||||
MethodGenerateImage GRPCMethod = "GenerateImage"
|
||||
MethodGenerateVideo GRPCMethod = "GenerateVideo"
|
||||
MethodGenerate3D GRPCMethod = "Generate3D"
|
||||
MethodAudioTranscription GRPCMethod = "AudioTranscription"
|
||||
MethodTTS GRPCMethod = "TTS"
|
||||
MethodTTSStream GRPCMethod = "TTSStream"
|
||||
@@ -63,7 +60,6 @@ const (
|
||||
MethodVoiceEmbed GRPCMethod = "VoiceEmbed"
|
||||
MethodVoiceAnalyze GRPCMethod = "VoiceAnalyze"
|
||||
MethodTokenClassify GRPCMethod = "TokenClassify"
|
||||
MethodScore GRPCMethod = "Score"
|
||||
)
|
||||
|
||||
// UsecaseInfo describes a single known_usecase value and how it maps
|
||||
@@ -126,11 +122,6 @@ var UsecaseInfoMap = map[string]UsecaseInfo{
|
||||
GRPCMethod: MethodGenerateVideo,
|
||||
Description: "Video generation via the GenerateVideo RPC, with optional image or audio conditioning when supported by the backend.",
|
||||
},
|
||||
Usecase3D: {
|
||||
Flag: FLAG_3D,
|
||||
GRPCMethod: MethodGenerate3D,
|
||||
Description: "Image-conditioned 3D asset generation via the Generate3D RPC — a binary glTF (GLB) mesh with optional PBR material (TRELLIS.2).",
|
||||
},
|
||||
UsecaseTranscript: {
|
||||
Flag: FLAG_TRANSCRIPT,
|
||||
GRPCMethod: MethodAudioTranscription,
|
||||
@@ -201,11 +192,6 @@ var UsecaseInfoMap = map[string]UsecaseInfo{
|
||||
GRPCMethod: MethodTokenClassify,
|
||||
Description: "Per-token classification (NER) via the TokenClassify RPC — the PII detector tier. Declared explicitly via known_usecases; never auto-guessed, since the token-classification head is not useful as general generation or embeddings.",
|
||||
},
|
||||
UsecaseScore: {
|
||||
Flag: FLAG_SCORE,
|
||||
GRPCMethod: MethodScore,
|
||||
Description: "Joint log-probability scoring of candidate continuations via the Score RPC. Declared explicitly via known_usecases and usable alongside generation usecases.",
|
||||
},
|
||||
}
|
||||
|
||||
// BackendCapability describes which gRPC methods and usecases a backend supports.
|
||||
@@ -255,8 +241,8 @@ func referenceVoiceCloning() *VoiceCloningCapability {
|
||||
var BackendCapabilities = map[string]BackendCapability{
|
||||
// --- LLM / text generation backends ---
|
||||
"llama-cpp": {
|
||||
GRPCMethods: []GRPCMethod{MethodPredict, MethodPredictStream, MethodEmbedding, MethodTokenizeString, MethodScore},
|
||||
PossibleUsecases: []string{UsecaseChat, UsecaseCompletion, UsecaseEdit, UsecaseEmbeddings, UsecaseTokenize, UsecaseVision, UsecaseScore},
|
||||
GRPCMethods: []GRPCMethod{MethodPredict, MethodPredictStream, MethodEmbedding, MethodTokenizeString},
|
||||
PossibleUsecases: []string{UsecaseChat, UsecaseCompletion, UsecaseEdit, UsecaseEmbeddings, UsecaseTokenize, UsecaseVision},
|
||||
DefaultUsecases: []string{UsecaseChat},
|
||||
AcceptsImages: true, // requires mmproj
|
||||
Description: "llama.cpp GGUF models — LLM inference with optional vision via mmproj",
|
||||
@@ -358,14 +344,6 @@ var BackendCapabilities = map[string]BackendCapability{
|
||||
Description: "Stable Diffusion via GGML quantized models",
|
||||
},
|
||||
|
||||
// --- 3D generation backends ---
|
||||
"trellis2cpp": {
|
||||
GRPCMethods: []GRPCMethod{MethodGenerate3D},
|
||||
PossibleUsecases: []string{Usecase3D},
|
||||
DefaultUsecases: []string{Usecase3D},
|
||||
Description: "trellis2.cpp — C++/GGML port of Microsoft TRELLIS.2: single-image to textured 3D mesh (GLB)",
|
||||
},
|
||||
|
||||
// --- Speech-to-text backends ---
|
||||
"whisper": {
|
||||
GRPCMethods: []GRPCMethod{MethodAudioTranscription, MethodVAD},
|
||||
|
||||
@@ -623,13 +623,6 @@ func DefaultRegistry() map[string]FieldMetaOverride {
|
||||
Component: "toggle",
|
||||
Order: 89,
|
||||
},
|
||||
"pipeline.turn_detection.vad_window_sec": {
|
||||
Section: "pipeline",
|
||||
Label: "VAD Window (s)",
|
||||
Description: "Widen the slice of recent audio the VAD rescans each turn-detection tick. Sized automatically from the commit silence threshold (server_vad silence window, or the semantic eagerness fallback) plus a warm-up margin — set only to widen it; values below the automatic floor are ignored.",
|
||||
Component: "number",
|
||||
Order: 90,
|
||||
},
|
||||
"pipeline.disable_warmup": {
|
||||
Section: "pipeline",
|
||||
Label: "Disable Warmup",
|
||||
@@ -637,99 +630,6 @@ func DefaultRegistry() map[string]FieldMetaOverride {
|
||||
Component: "toggle",
|
||||
Order: 90,
|
||||
},
|
||||
"pipeline.classifier.enabled": {
|
||||
Section: "pipeline",
|
||||
Label: "Classifier Mode",
|
||||
Description: "Replace autoregressive generation with prefill-only option selection: each user turn is scored against the option list via the Score primitive and the winning option's canned reply / tool call is emitted. Built for hardware that can afford prompt processing but not decode (e.g. a Raspberry Pi).",
|
||||
Component: "toggle",
|
||||
Order: 91,
|
||||
},
|
||||
"pipeline.classifier.options": {
|
||||
Section: "pipeline",
|
||||
Label: "Classifier Options",
|
||||
Description: "The intents the classifier scores each turn against. Each option has an id (also the scored route label — keep it short), a description of when it applies, an optional canned spoken reply, and an optional canned tool call {name, arguments}. A tool may also declare slots ([{name, type: number|enum|string, values, default, hint}]) whose \"{{name}}\" placeholders in arguments (and, optionally, the reply) are filled by a short grammar-constrained completion when the option wins — the hybrid between prefill-only classification and full generation (requires completion in the scoring model's known_usecases). Clients can replace the list per session via session.update localai_classifier.",
|
||||
Component: "json-editor",
|
||||
Order: 92,
|
||||
},
|
||||
"pipeline.classifier.threshold": {
|
||||
Section: "pipeline",
|
||||
Label: "Classifier Threshold",
|
||||
Description: "Softmax-probability floor the best option must clear; below it the fallback applies. 0 always picks the argmax.",
|
||||
Component: "slider",
|
||||
Min: f64(0),
|
||||
Max: f64(0.99),
|
||||
Step: f64(0.01),
|
||||
Order: 93,
|
||||
},
|
||||
"pipeline.classifier.fallback.mode": {
|
||||
Section: "pipeline",
|
||||
Label: "Classifier Fallback",
|
||||
Description: "What happens when no option clears the threshold: complete with no output, speak the canned fallback reply, or fall through to normal (slow) generation.",
|
||||
Component: "select",
|
||||
Options: []FieldOption{
|
||||
{Value: "none", Label: "none (empty response)"},
|
||||
{Value: "reply", Label: "canned reply"},
|
||||
{Value: "generate", Label: "generate"},
|
||||
},
|
||||
Order: 94,
|
||||
},
|
||||
"pipeline.classifier.fallback.reply": {
|
||||
Section: "pipeline",
|
||||
Label: "Classifier Fallback Reply",
|
||||
Description: "The canned reply spoken when the fallback mode is 'reply' and no option clears the threshold.",
|
||||
Component: "text",
|
||||
Order: 95,
|
||||
},
|
||||
"pipeline.classifier.normalization": {
|
||||
Section: "pipeline",
|
||||
Label: "Classifier Normalization",
|
||||
Description: "How option scores feed the softmax: 'raw' compares joint log-probs (default); 'mean' divides by token count, which is fairer when option ids have very different lengths.",
|
||||
Component: "select",
|
||||
Options: []FieldOption{
|
||||
{Value: "raw", Label: "raw (joint log-prob)"},
|
||||
{Value: "mean", Label: "mean (per-token)"},
|
||||
},
|
||||
Order: 96,
|
||||
},
|
||||
"pipeline.classifier.history_items": {
|
||||
Section: "pipeline",
|
||||
Label: "Classifier History Items",
|
||||
Description: "What gets scored: 0 or -1 (default) score only the latest user message; a positive N includes the trailing N conversation messages, role-labeled. Prior turns echo option names and can dominate small scoring models — only opt in with a larger scorer.",
|
||||
Component: "number",
|
||||
Order: 97,
|
||||
},
|
||||
"pipeline.classifier.model": {
|
||||
Section: "pipeline",
|
||||
Label: "Classifier Scoring Model",
|
||||
Description: "Optionally score on a different model config. Empty uses the pipeline LLM — scoring runs through the same llama.cpp slot as generation and shares its prompt cache, so a separate model is rarely needed.",
|
||||
Component: "model-select",
|
||||
AutocompleteProvider: ProviderModels,
|
||||
Order: 98,
|
||||
},
|
||||
"pipeline.classifier.address.names": {
|
||||
Section: "pipeline",
|
||||
Label: "Classifier Address Names",
|
||||
Description: "Wake-word gate: only act on turns that mention one of these names as a whole word ('Drone go up', not just 'go up'). Matching is deterministic on the transcript; unaddressed turns skip scoring entirely.",
|
||||
Component: "string-list",
|
||||
Order: 99,
|
||||
},
|
||||
"pipeline.classifier.address.mode": {
|
||||
Section: "pipeline",
|
||||
Label: "Classifier Address Mode",
|
||||
Description: "What to do with unaddressed turns: 'ignore' completes silently (right for ambient conversation), 'reply' speaks the address reply.",
|
||||
Component: "select",
|
||||
Options: []FieldOption{
|
||||
{Value: "ignore", Label: "ignore (stay silent)"},
|
||||
{Value: "reply", Label: "reply (speak the address reply)"},
|
||||
},
|
||||
Order: 100,
|
||||
},
|
||||
"pipeline.classifier.address.reply": {
|
||||
Section: "pipeline",
|
||||
Label: "Classifier Address Reply",
|
||||
Description: "Spoken when an unaddressed turn arrives in 'reply' mode.",
|
||||
Order: 101,
|
||||
},
|
||||
|
||||
// --- Functions ---
|
||||
"function.grammar.parallel_calls": {
|
||||
@@ -922,13 +822,6 @@ func DefaultRegistry() map[string]FieldMetaOverride {
|
||||
Min: f64(0),
|
||||
Order: 213,
|
||||
},
|
||||
"proxy.cache_prompt": {
|
||||
Section: "proxy",
|
||||
Label: "Proxy Anthropic Prompt Cache",
|
||||
Description: "Inject Anthropic prompt-cache breakpoints (cache_control: ephemeral) on the stable prefix (system, tools, last message) when mode is translate and provider is anthropic. Serves the repeated prefix at the cache-read rate on multi-turn/agentic calls. No effect otherwise.",
|
||||
Component: "checkbox",
|
||||
Order: 214,
|
||||
},
|
||||
|
||||
// --- MITM intercept hosts ---
|
||||
// Each host listed here is claimed by this model config; the
|
||||
|
||||
@@ -15,10 +15,9 @@ const (
|
||||
ModalityImage = "image"
|
||||
ModalityAudio = "audio"
|
||||
ModalityVideo = "video"
|
||||
Modality3D = "3d"
|
||||
)
|
||||
|
||||
var modalityOrder = []string{ModalityText, ModalityImage, ModalityAudio, ModalityVideo, Modality3D}
|
||||
var modalityOrder = []string{ModalityText, ModalityImage, ModalityAudio, ModalityVideo}
|
||||
|
||||
func declaredModalities(modalities []string) map[string]bool {
|
||||
declared := make(map[string]bool, len(modalities))
|
||||
@@ -155,7 +154,6 @@ func (c *ModelConfig) Capabilities() []string {
|
||||
add(c.HasUsecases(FLAG_SOUND_GENERATION), UsecaseSoundGeneration)
|
||||
add(c.HasUsecases(FLAG_IMAGE), UsecaseImage)
|
||||
add(c.HasUsecases(FLAG_VIDEO), UsecaseVideo)
|
||||
add(c.HasUsecases(FLAG_3D), Usecase3D)
|
||||
add(c.HasUsecases(FLAG_VAD), UsecaseVAD)
|
||||
add(c.HasUsecases(FLAG_DETECTION), UsecaseDetection)
|
||||
add(c.HasUsecases(FLAG_DEPTH), UsecaseDepth)
|
||||
@@ -183,10 +181,9 @@ func (c *ModelConfig) InputModalities() []string {
|
||||
c.HasUsecases(FLAG_TTS) || c.HasUsecases(FLAG_SOUND_GENERATION) || imageGen || videoGen
|
||||
|
||||
// Image input via a chat model requires vision (gated on chat, like the
|
||||
// Ollama surface); detection/depth/face/3D models consume images directly.
|
||||
// Ollama surface); detection/depth/face models consume images directly.
|
||||
imageIn := (chatish && c.VisionSupported()) || c.LimitMMPerPrompt.LimitImagePerPrompt > 0 ||
|
||||
c.HasUsecases(FLAG_DETECTION) || c.HasUsecases(FLAG_DEPTH) || c.HasUsecases(FLAG_FACE_RECOGNITION) ||
|
||||
c.HasUsecases(FLAG_3D)
|
||||
c.HasUsecases(FLAG_DETECTION) || c.HasUsecases(FLAG_DEPTH) || c.HasUsecases(FLAG_FACE_RECOGNITION)
|
||||
|
||||
audioIn := c.AudioInputSupported() || c.HasUsecases(FLAG_TRANSCRIPT) || c.HasUsecases(FLAG_AUDIO_TRANSFORM) ||
|
||||
c.HasUsecases(FLAG_REALTIME_AUDIO) || c.HasUsecases(FLAG_VAD) || c.HasUsecases(FLAG_DIARIZATION) ||
|
||||
@@ -211,12 +208,10 @@ func (c *ModelConfig) OutputModalities() []string {
|
||||
audioOut := c.HasUsecases(FLAG_TTS) || c.HasUsecases(FLAG_SOUND_GENERATION) ||
|
||||
c.HasUsecases(FLAG_AUDIO_TRANSFORM) || c.HasUsecases(FLAG_REALTIME_AUDIO)
|
||||
videoOut := c.HasUsecases(FLAG_VIDEO)
|
||||
threeDOut := c.HasUsecases(FLAG_3D)
|
||||
|
||||
modalities[ModalityText] = modalities[ModalityText] || textOut
|
||||
modalities[ModalityImage] = modalities[ModalityImage] || imageOut
|
||||
modalities[ModalityAudio] = modalities[ModalityAudio] || audioOut
|
||||
modalities[ModalityVideo] = modalities[ModalityVideo] || videoOut
|
||||
modalities[Modality3D] = modalities[Modality3D] || threeDOut
|
||||
return orderedModalities(modalities)
|
||||
}
|
||||
|
||||
@@ -109,25 +109,6 @@ var _ = Describe("Model capabilities derivation", func() {
|
||||
Expect(cfg.OutputModalities()).To(Equal([]string{"image"}))
|
||||
})
|
||||
|
||||
It("guesses the 3d usecase from the trellis2cpp backend and only that backend", func() {
|
||||
cfg := &ModelConfig{Backend: "trellis2cpp"}
|
||||
Expect(cfg.HasUsecases(FLAG_3D)).To(BeTrue())
|
||||
Expect(cfg.Capabilities()).To(ContainElement(Usecase3D))
|
||||
|
||||
other := &ModelConfig{Backend: "llama-cpp"}
|
||||
Expect(other.HasUsecases(FLAG_3D)).To(BeFalse())
|
||||
})
|
||||
|
||||
It("a 3D-generation model reads an image and writes a 3D asset", func() {
|
||||
// Pins the wire strings the UI depends on: capability "3d",
|
||||
// input modality "image" (no text prompt — TRELLIS.2 is
|
||||
// image-conditioned only), output modality "3d".
|
||||
cfg := &ModelConfig{KnownUsecases: usecaseBits(FLAG_3D), Backend: "trellis2cpp"}
|
||||
Expect(cfg.Capabilities()).To(Equal([]string{Usecase3D}))
|
||||
Expect(cfg.InputModalities()).To(Equal([]string{ModalityImage}))
|
||||
Expect(cfg.OutputModalities()).To(Equal([]string{Modality3D}))
|
||||
})
|
||||
|
||||
It("conditioned video uses declared modalities without backend-specific inference", func() {
|
||||
cfg := &ModelConfig{
|
||||
KnownUsecases: usecaseBits(FLAG_VIDEO),
|
||||
|
||||
@@ -232,15 +232,6 @@ type ProxyConfig struct {
|
||||
// means no per-request timeout (only the request context, which
|
||||
// is bound to the client connection, applies).
|
||||
RequestTimeoutSeconds int `yaml:"request_timeout_seconds,omitempty" json:"request_timeout_seconds,omitempty"`
|
||||
|
||||
// CachePrompt enables automatic Anthropic prompt-cache breakpoints
|
||||
// (cache_control: ephemeral) on the stable prefix — system prompt,
|
||||
// tools, and the last message block — when mode=translate and
|
||||
// provider=anthropic. Anthropic then serves the repeated prefix at
|
||||
// the cache-read rate (0.1x input), which sharply cuts cost on
|
||||
// agentic/multi-turn workloads that re-send a large stable prefix.
|
||||
// No effect for passthrough mode or non-Anthropic providers.
|
||||
CachePrompt bool `yaml:"cache_prompt,omitempty" json:"cache_prompt,omitempty"`
|
||||
}
|
||||
|
||||
// Proxy mode names. Validate() normalises an empty Mode to
|
||||
@@ -678,16 +669,6 @@ type Pipeline struct {
|
||||
// per session; retranscribe is server-side only. Unset keeps server_vad.
|
||||
TurnDetection PipelineTurnDetection `yaml:"turn_detection,omitempty" json:"turn_detection,omitempty"`
|
||||
|
||||
// Classifier switches realtime responses to prefill-only option
|
||||
// selection (LocalAI classifier mode): each user turn is scored
|
||||
// against a fixed option list via the Score primitive and the winning
|
||||
// option's canned reply / tool call is emitted, so weak hardware
|
||||
// never pays for autoregressive decode. Nil means disabled; clients
|
||||
// can still enable per session via session.update localai_classifier.
|
||||
// Validated (and rejected loudly) at realtime session setup, like the
|
||||
// pipeline model slots.
|
||||
Classifier *PipelineClassifier `yaml:"classifier,omitempty" json:"classifier,omitempty"`
|
||||
|
||||
// DisableWarmup turns off eager pre-loading of the pipeline's sub-models at
|
||||
// realtime session start. By default (false) LocalAI loads every configured
|
||||
// sub-model backend (VAD, transcription, LLM, TTS, sound detection, voice
|
||||
@@ -701,65 +682,6 @@ type Pipeline struct {
|
||||
DisableWarmup bool `yaml:"disable_warmup,omitempty" json:"disable_warmup,omitempty"`
|
||||
}
|
||||
|
||||
// PipelineClassifier is the YAML mirror of the realtime API's
|
||||
// localai_classifier extension (see
|
||||
// core/http/endpoints/openai/types/classifier.go, which documents the
|
||||
// field semantics and owns validation — the realtime session converts and
|
||||
// validates this block at setup).
|
||||
type PipelineClassifier struct {
|
||||
Enabled bool `yaml:"enabled,omitempty" json:"enabled,omitempty"`
|
||||
// Model optionally names a different config to score on. Empty uses
|
||||
// the pipeline's llm — with slot-based Score the same process serves
|
||||
// both scoring and generation and shares its prompt cache.
|
||||
Model string `yaml:"model,omitempty" json:"model,omitempty"`
|
||||
Threshold float64 `yaml:"threshold,omitempty" json:"threshold,omitempty"`
|
||||
Normalization string `yaml:"normalization,omitempty" json:"normalization,omitempty"`
|
||||
HistoryItems int `yaml:"history_items,omitempty" json:"history_items,omitempty"`
|
||||
Fallback *PipelineClassifierFallback `yaml:"fallback,omitempty" json:"fallback,omitempty"`
|
||||
Options []PipelineClassifierOption `yaml:"options,omitempty" json:"options,omitempty"`
|
||||
// Address gates every turn on the assistant being addressed by one of
|
||||
// these names (wake-word behavior); see types.ClassifierAddress.
|
||||
Address *PipelineClassifierAddress `yaml:"address,omitempty" json:"address,omitempty"`
|
||||
}
|
||||
|
||||
// PipelineClassifierAddress mirrors types.ClassifierAddress for YAML.
|
||||
type PipelineClassifierAddress struct {
|
||||
Names []string `yaml:"names,omitempty" json:"names,omitempty"`
|
||||
Mode string `yaml:"mode,omitempty" json:"mode,omitempty"`
|
||||
Reply string `yaml:"reply,omitempty" json:"reply,omitempty"`
|
||||
}
|
||||
|
||||
type PipelineClassifierOption struct {
|
||||
ID string `yaml:"id" json:"id"`
|
||||
Description string `yaml:"description" json:"description"`
|
||||
Reply string `yaml:"reply,omitempty" json:"reply,omitempty"`
|
||||
Tool *PipelineClassifierTool `yaml:"tool,omitempty" json:"tool,omitempty"`
|
||||
}
|
||||
|
||||
type PipelineClassifierTool struct {
|
||||
Name string `yaml:"name" json:"name"`
|
||||
// Arguments is a plain YAML map; the realtime session marshals it to
|
||||
// the JSON arguments string of the emitted function call. With Slots
|
||||
// it is a template: "{{name}}" values are filled by a constrained
|
||||
// completion when the option wins.
|
||||
Arguments map[string]any `yaml:"arguments,omitempty" json:"arguments,omitempty"`
|
||||
// Slots declares the inferred arguments; see types.ClassifierSlot.
|
||||
Slots []PipelineClassifierSlot `yaml:"slots,omitempty" json:"slots,omitempty"`
|
||||
}
|
||||
|
||||
type PipelineClassifierSlot struct {
|
||||
Name string `yaml:"name" json:"name"`
|
||||
Type string `yaml:"type" json:"type"` // number | enum | string
|
||||
Values []string `yaml:"values,omitempty" json:"values,omitempty"`
|
||||
Default string `yaml:"default,omitempty" json:"default,omitempty"`
|
||||
Hint string `yaml:"hint,omitempty" json:"hint,omitempty"`
|
||||
}
|
||||
|
||||
type PipelineClassifierFallback struct {
|
||||
Mode string `yaml:"mode,omitempty" json:"mode,omitempty"`
|
||||
Reply string `yaml:"reply,omitempty" json:"reply,omitempty"`
|
||||
}
|
||||
|
||||
// PipelineCompaction configures summarize-then-drop for a realtime pipeline.
|
||||
type PipelineCompaction struct {
|
||||
// Enabled turns summarize-then-drop on. Default false.
|
||||
@@ -1060,12 +982,6 @@ type PipelineTurnDetection struct {
|
||||
// are compared in the logs — a diagnostic for streaming/batch alignment
|
||||
// at the cost of one extra decode per turn.
|
||||
Retranscribe *bool `yaml:"retranscribe,omitempty" json:"retranscribe,omitempty"`
|
||||
// VadWindowSec widens the slice of recent audio the VAD rescans each
|
||||
// tick. The pipeline sizes it automatically from the commit silence
|
||||
// threshold (server_vad silence window, or the semantic eagerness
|
||||
// fallback) plus a warm-up margin; set this only to widen it further —
|
||||
// values below the automatic floor are ignored.
|
||||
VadWindowSec float64 `yaml:"vad_window_sec,omitempty" json:"vad_window_sec,omitempty"`
|
||||
}
|
||||
|
||||
// TurnDetectionSemantic reports whether this pipeline defaults sessions to
|
||||
@@ -1519,9 +1435,20 @@ func (c *ModelConfig) Validate() (bool, error) {
|
||||
ProxyProviderOpenAI, ProxyProviderAnthropic)
|
||||
}
|
||||
|
||||
// Score on llama-cpp runs through the slot loop (SERVER_TASK_TYPE_SCORE,
|
||||
// see backend/cpp/llama-cpp/patches/), so it is safe to combine with
|
||||
// chat/completion/embeddings on one config — no conflict check needed.
|
||||
// Score on llama-cpp bypasses the slot loop and races the
|
||||
// llama_context against concurrent generation/embedding traffic
|
||||
// (see backend/cpp/llama-cpp/grpc-server.cpp on Score). Reject the
|
||||
// combination here so operators are forced to split the model.
|
||||
// (token_classify is unaffected — it runs on the standalone
|
||||
// privacy-filter backend, not llama-cpp.)
|
||||
const scoreConflicts = FLAG_CHAT | FLAG_COMPLETION | FLAG_EMBEDDINGS
|
||||
if (c.Backend == "llama-cpp" || c.Backend == "llama") &&
|
||||
c.HasUsecases(FLAG_SCORE) && c.KnownUsecases != nil &&
|
||||
*c.KnownUsecases&scoreConflicts != 0 {
|
||||
return false, fmt.Errorf(
|
||||
"known_usecases conflict on llama-cpp: score is incompatible " +
|
||||
"with chat/completion/embeddings — split into separate model configs")
|
||||
}
|
||||
|
||||
// Pattern detector: validate built-in names and that each operator-defined
|
||||
// pattern is a well-formed, anchored, bounded restricted-regex. Reject at
|
||||
@@ -1649,10 +1576,9 @@ const (
|
||||
// Marks a model as wired for the Score gRPC primitive (joint
|
||||
// log-prob of candidate continuations under a shared prompt). Must
|
||||
// be declared explicitly via `known_usecases: [score]` — there's
|
||||
// no heuristic for it. On llama-cpp, Score runs through the slot
|
||||
// loop (SERVER_TASK_TYPE_SCORE), so it may combine freely with
|
||||
// chat/completion/embeddings on one config and shares the slot's
|
||||
// prompt cache with generation.
|
||||
// no heuristic for it. On llama-cpp, Score bypasses the slot loop
|
||||
// (direct llama_decode), so combining score with
|
||||
// chat/completion/embeddings in one config is rejected at validation.
|
||||
FLAG_SCORE ModelConfigUsecase = 0b10000000000000000000
|
||||
|
||||
// Marks a model as wired for the Depth gRPC primitive (per-pixel
|
||||
@@ -1673,11 +1599,6 @@ const (
|
||||
// labels via the SoundDetection RPC, e.g. ced).
|
||||
FLAG_SOUND_CLASSIFICATION ModelConfigUsecase = 0b10000000000000000000000
|
||||
|
||||
// Marks a model as wired for the Generate3D gRPC primitive
|
||||
// (image-conditioned 3D asset generation — a binary glTF mesh with
|
||||
// optional PBR material, e.g. trellis2cpp).
|
||||
FLAG_3D ModelConfigUsecase = 0b100000000000000000000000
|
||||
|
||||
// Common Subsets
|
||||
FLAG_LLM ModelConfigUsecase = FLAG_CHAT | FLAG_COMPLETION | FLAG_EDIT
|
||||
)
|
||||
@@ -1691,7 +1612,7 @@ var ModalityGroups = []ModelConfigUsecase{
|
||||
FLAG_TRANSCRIPT | FLAG_REALTIME_AUDIO | FLAG_SOUND_CLASSIFICATION, // audio input — realtime_audio is any-to-any, so it counts here too
|
||||
FLAG_TTS | FLAG_SOUND_GENERATION | FLAG_REALTIME_AUDIO, // audio output — and here, so a lone realtime_audio flag still reads as multimodal
|
||||
FLAG_AUDIO_TRANSFORM, // audio in/out transforms
|
||||
FLAG_IMAGE | FLAG_VIDEO | FLAG_3D, // visual generation
|
||||
FLAG_IMAGE | FLAG_VIDEO, // visual generation
|
||||
}
|
||||
|
||||
// IsMultimodal returns true if the given usecases span two or more orthogonal
|
||||
@@ -1738,7 +1659,6 @@ func GetAllModelConfigUsecases() map[string]ModelConfigUsecase {
|
||||
"FLAG_SCORE": FLAG_SCORE,
|
||||
"FLAG_DEPTH": FLAG_DEPTH,
|
||||
"FLAG_TOKEN_CLASSIFY": FLAG_TOKEN_CLASSIFY,
|
||||
"FLAG_3D": FLAG_3D,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1771,9 +1691,9 @@ func GetUsecasesFromYAML(input []string) *ModelConfigUsecase {
|
||||
// either, they reserved the model for an internal direct-decode primitive
|
||||
// (the router classifier, or the PII NER tier). Letting GuessUsecases
|
||||
// paint chat/completion/embeddings on top would surface it in pickers it
|
||||
// was deliberately kept out of. So a declared score or token_classify
|
||||
// list is authoritative; declare the generation usecases explicitly
|
||||
// alongside score to serve both from one config.
|
||||
// was deliberately kept out of, and (on llama-cpp) reintroduce the slot
|
||||
// contention the conflict check exists to prevent. So a declared score or
|
||||
// token_classify list is authoritative.
|
||||
func (c *ModelConfig) HasUsecases(u ModelConfigUsecase) bool {
|
||||
if c.KnownUsecases != nil {
|
||||
if (u & *c.KnownUsecases) == u {
|
||||
@@ -1890,13 +1810,6 @@ func (c *ModelConfig) GuessUsecases(u ModelConfigUsecase) bool {
|
||||
}
|
||||
}
|
||||
|
||||
if (u & FLAG_3D) == FLAG_3D {
|
||||
threeDBackends := []string{"trellis2cpp"}
|
||||
if !slices.Contains(threeDBackends, c.Backend) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
if (u & FLAG_FACE_RECOGNITION) == FLAG_FACE_RECOGNITION {
|
||||
faceBackends := []string{"insightface"}
|
||||
if !slices.Contains(faceBackends, c.Backend) {
|
||||
@@ -1970,8 +1883,8 @@ func (c *ModelConfig) GuessUsecases(u ModelConfigUsecase) bool {
|
||||
|
||||
if (u & FLAG_SCORE) == FLAG_SCORE {
|
||||
// No heuristic: Score-intent is a deliberate operator choice
|
||||
// (it keeps the model out of pickers it wasn't meant for), so
|
||||
// HasUsecases(FLAG_SCORE) is true only when KnownUsecases
|
||||
// (it reserves the model from generation traffic on llama-cpp),
|
||||
// so HasUsecases(FLAG_SCORE) is true only when KnownUsecases
|
||||
// declares it explicitly.
|
||||
return false
|
||||
}
|
||||
|
||||
@@ -201,8 +201,8 @@ func (bcl *ModelConfigLoader) LoadModelConfigFileByNameDefaultOptions(modelName
|
||||
// survives unresolved into model loading and fails downstream — notably in
|
||||
// distributed mode with "backend name is empty". Mirrors the top-level alias
|
||||
// resolution in core/http/middleware/request.go.
|
||||
func (bcl *ModelConfigLoader) LoadResolvedModelConfig(modelName, modelPath string, opts ...ConfigLoaderOption) (*ModelConfig, error) {
|
||||
cfg, err := bcl.LoadModelConfigFileByName(modelName, modelPath, opts...)
|
||||
func (bcl *ModelConfigLoader) LoadResolvedModelConfig(modelName, modelPath string) (*ModelConfig, error) {
|
||||
cfg, err := bcl.LoadModelConfigFileByName(modelName, modelPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -49,21 +49,4 @@ alias: real-llm
|
||||
Expect(direct.Backend).To(Equal("llama-cpp"))
|
||||
Expect(direct.Name).To(Equal("real-llm"))
|
||||
})
|
||||
|
||||
It("applies loader defaults while preserving explicit model threads", func() {
|
||||
tmpDir := GinkgoT().TempDir()
|
||||
Expect(os.WriteFile(filepath.Join(tmpDir, "defaulted.yaml"), []byte("name: defaulted\nbackend: llama-cpp\n"), 0644)).To(Succeed())
|
||||
Expect(os.WriteFile(filepath.Join(tmpDir, "explicit.yaml"), []byte("name: explicit\nbackend: llama-cpp\nthreads: 3\n"), 0644)).To(Succeed())
|
||||
|
||||
cl := config.NewModelConfigLoader(tmpDir)
|
||||
defaulted, err := cl.LoadResolvedModelConfig("defaulted", tmpDir, config.LoadOptionThreads(11))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(defaulted.Threads).NotTo(BeNil())
|
||||
Expect(*defaulted.Threads).To(Equal(11))
|
||||
|
||||
explicit, err := cl.LoadResolvedModelConfig("explicit", tmpDir, config.LoadOptionThreads(11))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(explicit.Threads).NotTo(BeNil())
|
||||
Expect(*explicit.Threads).To(Equal(3))
|
||||
})
|
||||
})
|
||||
|
||||
@@ -127,19 +127,21 @@ parameters:
|
||||
Expect(err).To(BeNil())
|
||||
Expect(valid).To(BeTrue())
|
||||
|
||||
// Score runs through the llama-cpp slot loop, so mixing the
|
||||
// score usecase with chat/completion/embeddings on one config
|
||||
// is valid — the slot scheduler serializes score against
|
||||
// generation and shares the prompt cache between them.
|
||||
// llama-cpp configs can't mix the score usecase with
|
||||
// chat/completion/embeddings — Score bypasses the slot loop
|
||||
// and would race the llama_context. (token_classify is exempt:
|
||||
// it runs on the privacy-filter backend, not llama-cpp, so the
|
||||
// token_classify combinations below stay valid.)
|
||||
scoreFlag := FLAG_SCORE | FLAG_CHAT
|
||||
scoringChat := ModelConfig{
|
||||
Name: "router-and-chat",
|
||||
conflicting := ModelConfig{
|
||||
Name: "router-but-also-chat",
|
||||
Backend: "llama-cpp",
|
||||
KnownUsecases: &scoreFlag,
|
||||
}
|
||||
valid, err = scoringChat.Validate()
|
||||
Expect(valid).To(BeTrue())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
valid, err = conflicting.Validate()
|
||||
Expect(valid).To(BeFalse())
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("score is incompatible"))
|
||||
|
||||
scoreOnly := FLAG_SCORE
|
||||
dedicated := ModelConfig{
|
||||
|
||||
83
core/gallery/backend_index_capabilities_test.go
Normal file
83
core/gallery/backend_index_capabilities_test.go
Normal file
@@ -0,0 +1,83 @@
|
||||
package gallery_test
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gopkg.in/yaml.v3"
|
||||
|
||||
"github.com/mudler/LocalAI/core/gallery"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
)
|
||||
|
||||
// loadBackendIndex parses backend/index.yaml once for the whole suite.
|
||||
var loadBackendIndex = sync.OnceValues(func() (gallery.GalleryElements[*gallery.GalleryBackend], error) {
|
||||
data, err := os.ReadFile(filepath.Join("..", "..", "backend", "index.yaml"))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var entries gallery.GalleryElements[*gallery.GalleryBackend]
|
||||
if err := yaml.Unmarshal(data, &entries); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return entries, nil
|
||||
})
|
||||
|
||||
var _ = Describe("backend/index.yaml capability maps", func() {
|
||||
var entries gallery.GalleryElements[*gallery.GalleryBackend]
|
||||
|
||||
BeforeEach(func() {
|
||||
var err error
|
||||
entries, err = loadBackendIndex()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(entries).ToNot(BeEmpty())
|
||||
})
|
||||
|
||||
// A capability pointing at a name that does not exist is invisible until a
|
||||
// host with exactly that capability tries to install: FindBestBackendFromMeta
|
||||
// returns nil and the install fails with "no backend found".
|
||||
It("resolves every capability reference to an entry in the index", func() {
|
||||
names := map[string]struct{}{}
|
||||
for _, e := range entries {
|
||||
names[e.Name] = struct{}{}
|
||||
}
|
||||
|
||||
dangling := []string{}
|
||||
for _, e := range entries {
|
||||
for capability, target := range e.CapabilitiesMap {
|
||||
if _, ok := names[target]; !ok {
|
||||
dangling = append(dangling, fmt.Sprintf(" %s -> %s: %q", e.Name, capability, target))
|
||||
}
|
||||
}
|
||||
}
|
||||
Expect(dangling).To(BeEmpty(), "capabilities naming a missing entry:\n%s", strings.Join(dangling, "\n"))
|
||||
})
|
||||
|
||||
// vllm.cpp's CUDA kernels need the CUDA 13 toolchain (12.x nvcc cannot
|
||||
// compile the Blackwell fp4 paths), so CUDA 12 hosts have no GPU build to
|
||||
// install and must land on the CPU one. Assert the fallback is explicit
|
||||
// rather than an accident of the "default" catch-all, so mapping these
|
||||
// capabilities at a CUDA image later is a test failure and not a host that
|
||||
// pulls kernels it cannot run.
|
||||
DescribeTable("routes vllm-cpp hosts to the build their toolchain supports",
|
||||
func(metaName, capability, expected string) {
|
||||
meta := entries.FindByName(metaName)
|
||||
Expect(meta).ToNot(BeNil())
|
||||
|
||||
resolved := meta.FindBestBackendFromMeta(system.NewCapabilityState(capability), entries)
|
||||
Expect(resolved).ToNot(BeNil())
|
||||
Expect(resolved.Name).To(Equal(expected))
|
||||
},
|
||||
Entry("CUDA 12 x86_64 gets the CPU build", "vllm-cpp", "nvidia-cuda-12", "cpu-vllm-cpp"),
|
||||
Entry("CUDA 12 Jetson (AGX Orin) gets the CPU build", "vllm-cpp", "nvidia-l4t-cuda-12", "cpu-vllm-cpp"),
|
||||
Entry("CUDA 13 Jetson (DGX Spark) gets the L4T build", "vllm-cpp", "nvidia-l4t-cuda-13", "nvidia-l4t-arm64-vllm-cpp"),
|
||||
Entry("CUDA 13 x86_64 gets the CUDA build", "vllm-cpp", "nvidia-cuda-13", "cuda13-vllm-cpp"),
|
||||
Entry("development CUDA 12 Jetson gets the CPU build", "vllm-cpp-development", "nvidia-l4t-cuda-12", "cpu-vllm-cpp-development"),
|
||||
Entry("development CUDA 13 Jetson gets the L4T build", "vllm-cpp-development", "nvidia-l4t-cuda-13", "nvidia-l4t-arm64-vllm-cpp-development"),
|
||||
)
|
||||
})
|
||||
@@ -32,42 +32,6 @@ var _ = Describe("Runtime capability-based backend selection", func() {
|
||||
os.RemoveAll(tempDir)
|
||||
})
|
||||
|
||||
It("keeps the Kokoro CPU fallback installable from the backend gallery", func() {
|
||||
backends, err := ReadConfigFile[[]*GalleryBackend](filepath.Join("..", "..", "backend", "index.yaml"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
byName := make(map[string]*GalleryBackend, len(*backends))
|
||||
for _, backend := range *backends {
|
||||
byName[backend.Name] = backend
|
||||
}
|
||||
|
||||
Expect(byName).To(HaveKey("kokoro"))
|
||||
Expect(byName["kokoro"].CapabilitiesMap).To(HaveKeyWithValue("default", "cpu-kokoro"))
|
||||
Expect(byName).To(HaveKey("cpu-kokoro"))
|
||||
Expect(byName["cpu-kokoro"].URI).To(Equal("quay.io/go-skynet/local-ai-backends:latest-cpu-kokoro"))
|
||||
|
||||
type matrixEntry struct {
|
||||
Backend string `yaml:"backend"`
|
||||
Platforms string `yaml:"platforms"`
|
||||
PlatformTag string `yaml:"platform-tag"`
|
||||
TagSuffix string `yaml:"tag-suffix"`
|
||||
}
|
||||
type backendMatrix struct {
|
||||
Include []matrixEntry `yaml:"include"`
|
||||
}
|
||||
|
||||
matrix, err := ReadConfigFile[backendMatrix](filepath.Join("..", "..", ".github", "backend-matrix.yml"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
var cpuArchitectures []string
|
||||
for _, entry := range matrix.Include {
|
||||
if entry.Backend == "kokoro" && entry.TagSuffix == "-cpu-kokoro" {
|
||||
cpuArchitectures = append(cpuArchitectures, entry.Platforms+"/"+entry.PlatformTag)
|
||||
}
|
||||
}
|
||||
Expect(cpuArchitectures).To(ConsistOf("linux/amd64/amd64", "linux/arm64/arm64"))
|
||||
})
|
||||
|
||||
It("ListSystemBackends prefers optimal alias candidate", func() {
|
||||
// Arrange two installed backends sharing the same alias
|
||||
must := func(err error) { Expect(err).NotTo(HaveOccurred()) }
|
||||
|
||||
@@ -143,11 +143,6 @@ var defaultImporters = []Importer{
|
||||
&CoquiImporter{},
|
||||
// Image/Video (Batch 3)
|
||||
&StableDiffusionGGMLImporter{},
|
||||
// Trellis2CppImporter (TRELLIS.2 image-to-3D, native C++/ggml port) must
|
||||
// run before LlamaCPPImporter so its GGUF sets aren't claimed by the
|
||||
// generic .gguf importer; matches only trellis-named URIs/repos or the
|
||||
// distinctive component filenames, so arbitrary GGUFs are never claimed.
|
||||
&Trellis2CppImporter{},
|
||||
&ACEStepImporter{},
|
||||
// LongCat repositories carry generic Diffusers metadata, so this exact
|
||||
// owner/repo matcher must run before DiffuserImporter.
|
||||
|
||||
@@ -1,170 +0,0 @@
|
||||
package importers
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/gallery"
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
"go.yaml.in/yaml/v2"
|
||||
)
|
||||
|
||||
var _ Importer = &Trellis2CppImporter{}
|
||||
|
||||
// trellis2File describes one component of the TRELLIS.2 GGUF set hosted on
|
||||
// the LocalAI-io HuggingFace org. The pipeline spans three source repos
|
||||
// (TRELLIS.2-4B, TRELLIS-image-large for the SS decoder, and a DINOv3
|
||||
// mirror), so a single import URI always expands to this full set — no one
|
||||
// repo can describe it alone. Filenames follow the trellis2cpp converter
|
||||
// defaults, which the backend resolves without any options.
|
||||
type trellis2File struct {
|
||||
filename string
|
||||
uri string
|
||||
sha256 string
|
||||
}
|
||||
|
||||
var trellis2Files = []trellis2File{
|
||||
{"dino_f16.gguf", "https://huggingface.co/LocalAI-io/dinov3-vitl16-pretrain-lvd1689m-GGUF/resolve/main/dino_f16.gguf", "385d8186a38a2328ec740fb2ac1f33f9194d8774efc7ccafd4aa2e51cf5f6450"},
|
||||
{"ss_flow_f16.gguf", "https://huggingface.co/LocalAI-io/TRELLIS.2-4B-GGUF/resolve/main/ss_flow_f16.gguf", "1dded5b74237d24e6876a642a26f90b43742e3554418573860f810e3bbe61e8c"},
|
||||
{"ss_dec_f16.gguf", "https://huggingface.co/LocalAI-io/TRELLIS-image-large-GGUF/resolve/main/ss_dec_f16.gguf", "9c2210b7ed830fdc8286961a8189878ff5bcfd3bfc83ab4eacee005d293d2185"},
|
||||
{"slat_flow_f16.gguf", "https://huggingface.co/LocalAI-io/TRELLIS.2-4B-GGUF/resolve/main/slat_flow_f16.gguf", "2f94bad7b1c524ad8c01943bc38fcc0c314e7d482ce896f3c6e96eb6e7cec15c"},
|
||||
{"slat_flow_1024_f16.gguf", "https://huggingface.co/LocalAI-io/TRELLIS.2-4B-GGUF/resolve/main/slat_flow_1024_f16.gguf", "b6a2270131e2e9235e9b6cb525193eb85ae132fa5af3274322aacd39e40a6bc5"},
|
||||
{"shape_dec_f16.gguf", "https://huggingface.co/LocalAI-io/TRELLIS.2-4B-GGUF/resolve/main/shape_dec_f16.gguf", "6fe53f1d7763dabf7c8d72bc38f4053d87fde6f65bf17a9d378d27edb39d3530"},
|
||||
{"shape_enc_f16.gguf", "https://huggingface.co/LocalAI-io/TRELLIS.2-4B-GGUF/resolve/main/shape_enc_f16.gguf", "3ec80ff580987fcdb9bc594fc8b6fda890d63101ca442eb2b26f5dc315e8696c"},
|
||||
{"tex_dec_f16.gguf", "https://huggingface.co/LocalAI-io/TRELLIS.2-4B-GGUF/resolve/main/tex_dec_f16.gguf", "afd304f4dfcb8c94df851b85519b415b99f04070f7d29de1320c50631b1be4e0"},
|
||||
{"tex_slat_flow_512_f16.gguf", "https://huggingface.co/LocalAI-io/TRELLIS.2-4B-GGUF/resolve/main/tex_slat_flow_512_f16.gguf", "89a081b7f5487a5b31f03d240e4d959a56db0cc2c46c327230097a2554da52ae"},
|
||||
{"tex_slat_flow_1024_f16.gguf", "https://huggingface.co/LocalAI-io/TRELLIS.2-4B-GGUF/resolve/main/tex_slat_flow_1024_f16.gguf", "bbb55b0910c7929aac5e0612a9bb15113837a2c674cafb9f0f170eda8b5558a8"},
|
||||
}
|
||||
|
||||
// trellis2ComponentNames are the distinctive default component filenames. A
|
||||
// raw .gguf URL with one of these basenames is a strong trellis2 signal.
|
||||
// dino_f16.gguf is deliberately absent — DINO checkpoints are common enough
|
||||
// that the bare name would over-claim.
|
||||
var trellis2ComponentNames = map[string]struct{}{
|
||||
"ss_flow_f16.gguf": {},
|
||||
"ss_dec_f16.gguf": {},
|
||||
"slat_flow_f16.gguf": {},
|
||||
"slat_flow_1024_f16.gguf": {},
|
||||
"shape_dec_f16.gguf": {},
|
||||
"shape_enc_f16.gguf": {},
|
||||
"tex_dec_f16.gguf": {},
|
||||
"tex_slat_flow_512_f16.gguf": {},
|
||||
"tex_slat_flow_1024_f16.gguf": {},
|
||||
}
|
||||
|
||||
// Trellis2CppImporter recognises Microsoft TRELLIS.2 image-to-3D GGUF sets
|
||||
// (the trellis2.cpp converter outputs hosted under LocalAI-io). It must be
|
||||
// registered BEFORE LlamaCPPImporter so llama-cpp does not steal the .gguf
|
||||
// match. preferences.backend="trellis2cpp" overrides detection.
|
||||
type Trellis2CppImporter struct{}
|
||||
|
||||
func (i *Trellis2CppImporter) Name() string { return "trellis2cpp" }
|
||||
func (i *Trellis2CppImporter) Modality() string { return "3d" }
|
||||
func (i *Trellis2CppImporter) AutoDetects() bool { return true }
|
||||
|
||||
// containsTrellisToken reports whether s (compared case-insensitively)
|
||||
// carries a TRELLIS marker ("trellis" covers TRELLIS.2 / trellis2 too).
|
||||
func containsTrellisToken(s string) bool {
|
||||
return strings.Contains(strings.ToLower(s), "trellis")
|
||||
}
|
||||
|
||||
func (i *Trellis2CppImporter) Match(details Details) bool {
|
||||
preferences, err := details.Preferences.MarshalJSON()
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
preferencesMap := make(map[string]any)
|
||||
if len(preferences) > 0 {
|
||||
if err := json.Unmarshal(preferences, &preferencesMap); err != nil {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
if b, ok := preferencesMap["backend"].(string); ok && b != "" {
|
||||
return b == "trellis2cpp"
|
||||
}
|
||||
|
||||
// Raw .gguf URL named after a distinctive pipeline component.
|
||||
if strings.HasSuffix(strings.ToLower(details.URI), ".gguf") {
|
||||
base := strings.ToLower(filepath.Base(details.URI))
|
||||
if _, ok := trellis2ComponentNames[base]; ok {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
// A trellis-named URI or HF repo carrying GGUFs.
|
||||
if containsTrellisToken(details.URI) {
|
||||
if strings.HasSuffix(strings.ToLower(details.URI), ".gguf") {
|
||||
return true
|
||||
}
|
||||
if details.HuggingFace != nil && hasGGUF(details.HuggingFace.Files) {
|
||||
return true
|
||||
}
|
||||
// HF details may be nil (tree-listing quirk) — decide from the
|
||||
// owner/repo alone.
|
||||
if _, repo, ok := HFOwnerRepoFromURI(details.URI); ok && containsTrellisToken(repo) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
func (i *Trellis2CppImporter) Import(details Details) (gallery.ModelConfig, error) {
|
||||
preferences, err := details.Preferences.MarshalJSON()
|
||||
if err != nil {
|
||||
return gallery.ModelConfig{}, err
|
||||
}
|
||||
preferencesMap := make(map[string]any)
|
||||
if len(preferences) > 0 {
|
||||
if err := json.Unmarshal(preferences, &preferencesMap); err != nil {
|
||||
return gallery.ModelConfig{}, err
|
||||
}
|
||||
}
|
||||
|
||||
name, ok := preferencesMap["name"].(string)
|
||||
if !ok {
|
||||
name = "trellis2-4b"
|
||||
}
|
||||
|
||||
description, ok := preferencesMap["description"].(string)
|
||||
if !ok {
|
||||
description = "TRELLIS.2 image-to-3D (GLB with PBR textures) — imported from " + details.URI
|
||||
}
|
||||
|
||||
cfg := gallery.ModelConfig{
|
||||
Name: name,
|
||||
Description: description,
|
||||
}
|
||||
// The full pipeline spans three HF repos, so any trellis URI imports the
|
||||
// complete known-good set rather than whatever single repo was pasted.
|
||||
for _, f := range trellis2Files {
|
||||
cfg.Files = append(cfg.Files, gallery.File{
|
||||
URI: f.uri,
|
||||
Filename: f.filename,
|
||||
SHA256: f.sha256,
|
||||
})
|
||||
}
|
||||
|
||||
modelConfig := config.ModelConfig{
|
||||
Name: name,
|
||||
Description: description,
|
||||
Backend: "trellis2cpp",
|
||||
KnownUsecaseStrings: []string{"FLAG_3D"},
|
||||
PredictionOptions: schema.PredictionOptions{
|
||||
// ss_flow anchors the GGUF directory; the backend resolves the
|
||||
// other components from their default filenames next to it.
|
||||
BasicModelRequest: schema.BasicModelRequest{Model: "ss_flow_f16.gguf"},
|
||||
},
|
||||
}
|
||||
|
||||
data, err := yaml.Marshal(modelConfig)
|
||||
if err != nil {
|
||||
return gallery.ModelConfig{}, err
|
||||
}
|
||||
|
||||
cfg.ConfigFile = string(data)
|
||||
return cfg, nil
|
||||
}
|
||||
@@ -1,105 +0,0 @@
|
||||
package importers_test
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"github.com/mudler/LocalAI/core/gallery/importers"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("Trellis2CppImporter", func() {
|
||||
Context("detection from HuggingFace", func() {
|
||||
// LocalAI-io/TRELLIS.2-4B-GGUF is the canonical GGUF conversion of
|
||||
// microsoft/TRELLIS.2-4B produced by the trellis2cpp converters.
|
||||
// Detection must route it to trellis2cpp (and NOT to llama-cpp,
|
||||
// which otherwise steals every .gguf repo).
|
||||
It("matches the TRELLIS.2 GGUF repo and imports the full component set", func() {
|
||||
uri := "https://huggingface.co/LocalAI-io/TRELLIS.2-4B-GGUF"
|
||||
preferences := json.RawMessage(`{}`)
|
||||
|
||||
modelConfig, err := importers.DiscoverModelConfig(uri, preferences)
|
||||
|
||||
Expect(err).ToNot(HaveOccurred(), fmt.Sprintf("Error: %v", err))
|
||||
Expect(modelConfig.ConfigFile).To(ContainSubstring("backend: trellis2cpp"), fmt.Sprintf("Model config: %+v", modelConfig))
|
||||
Expect(modelConfig.ConfigFile).To(ContainSubstring("known_usecases"))
|
||||
Expect(modelConfig.ConfigFile).To(ContainSubstring("FLAG_3D"))
|
||||
// The pipeline spans three repos; the import must carry the whole
|
||||
// set, anchored on ss_flow.
|
||||
Expect(modelConfig.Files).To(HaveLen(10))
|
||||
Expect(modelConfig.ConfigFile).To(ContainSubstring("model: ss_flow_f16.gguf"))
|
||||
})
|
||||
|
||||
It("matches a raw .gguf URL named after a distinctive pipeline component", func() {
|
||||
uri := "https://example.com/models/tex_slat_flow_512_f16.gguf"
|
||||
preferences := json.RawMessage(`{}`)
|
||||
|
||||
modelConfig, err := importers.DiscoverModelConfig(uri, preferences)
|
||||
|
||||
Expect(err).ToNot(HaveOccurred(), fmt.Sprintf("Error: %v", err))
|
||||
Expect(modelConfig.ConfigFile).To(ContainSubstring("backend: trellis2cpp"), fmt.Sprintf("Model config: %+v", modelConfig))
|
||||
})
|
||||
})
|
||||
|
||||
Context("preference override", func() {
|
||||
It("honours preferences.backend=trellis2cpp for arbitrary URIs", func() {
|
||||
uri := "https://example.com/some-unrelated-model"
|
||||
preferences := json.RawMessage(`{"backend": "trellis2cpp"}`)
|
||||
|
||||
modelConfig, err := importers.DiscoverModelConfig(uri, preferences)
|
||||
|
||||
Expect(err).ToNot(HaveOccurred(), fmt.Sprintf("Error: %v", err))
|
||||
Expect(modelConfig.ConfigFile).To(ContainSubstring("backend: trellis2cpp"), fmt.Sprintf("Model config: %+v", modelConfig))
|
||||
})
|
||||
|
||||
It("does not override a different explicit backend", func() {
|
||||
imp := &importers.Trellis2CppImporter{}
|
||||
match := imp.Match(importers.Details{
|
||||
URI: "https://example.com/models/tex_slat_flow_512_f16.gguf",
|
||||
Preferences: json.RawMessage(`{"backend": "llama-cpp"}`),
|
||||
})
|
||||
|
||||
Expect(match).To(BeFalse())
|
||||
})
|
||||
|
||||
It("still auto-detects when the backend preference is empty", func() {
|
||||
imp := &importers.Trellis2CppImporter{}
|
||||
match := imp.Match(importers.Details{
|
||||
URI: "https://example.com/models/tex_slat_flow_512_f16.gguf",
|
||||
Preferences: json.RawMessage(`{"backend": ""}`),
|
||||
})
|
||||
|
||||
Expect(match).To(BeTrue())
|
||||
})
|
||||
})
|
||||
|
||||
Context("negative detection", func() {
|
||||
It("does not claim an unrelated raw .gguf URL", func() {
|
||||
imp := &importers.Trellis2CppImporter{}
|
||||
match := imp.Match(importers.Details{
|
||||
URI: "https://example.com/models/llama-3-8b-Q4_K.gguf",
|
||||
Preferences: json.RawMessage(`{}`),
|
||||
})
|
||||
Expect(match).To(BeFalse())
|
||||
})
|
||||
|
||||
It("does not claim a bare dino_f16.gguf (too generic a name)", func() {
|
||||
imp := &importers.Trellis2CppImporter{}
|
||||
match := imp.Match(importers.Details{
|
||||
URI: "https://example.com/models/dino_f16.gguf",
|
||||
Preferences: json.RawMessage(`{}`),
|
||||
})
|
||||
Expect(match).To(BeFalse())
|
||||
})
|
||||
})
|
||||
|
||||
Context("Importer interface metadata", func() {
|
||||
It("exposes name/modality/autodetect", func() {
|
||||
imp := &importers.Trellis2CppImporter{}
|
||||
Expect(imp.Name()).To(Equal("trellis2cpp"))
|
||||
Expect(imp.Modality()).To(Equal("3d"))
|
||||
Expect(imp.AutoDetects()).To(BeTrue())
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -55,12 +55,6 @@ var quietPaths = []string{"/api/operations", "/api/resources", "/healthz", "/rea
|
||||
// conditional revalidation round-trip.
|
||||
const immutableAssetCacheControl = "public, max-age=31536000, immutable"
|
||||
|
||||
func defaultBodyLimitSkipper(c echo.Context) bool {
|
||||
// Remeshing accepts generated GLBs that routinely exceed the default
|
||||
// upload limit. The route has its own tighter, format-specific limit.
|
||||
return c.Request().Method == http.MethodPost && c.Path() == "/3d/remesh"
|
||||
}
|
||||
|
||||
// applyModelLoadCooldown maps a ModelLoadCooldownError anywhere in err's chain
|
||||
// to HTTP 503 with a Retry-After header (whole seconds, floor 1), so a client
|
||||
// polling a model whose load recently failed backs off instead of triggering a
|
||||
@@ -129,10 +123,7 @@ func API(application *application.Application) (*echo.Echo, error) {
|
||||
|
||||
// Set body limit
|
||||
if application.ApplicationConfig().UploadLimitMB > 0 {
|
||||
e.Use(middleware.BodyLimitWithConfig(middleware.BodyLimitConfig{
|
||||
Limit: fmt.Sprintf("%dM", application.ApplicationConfig().UploadLimitMB),
|
||||
Skipper: defaultBodyLimitSkipper,
|
||||
}))
|
||||
e.Use(middleware.BodyLimit(fmt.Sprintf("%dM", application.ApplicationConfig().UploadLimitMB)))
|
||||
}
|
||||
|
||||
// SPA fallback handler, set later when React UI is available
|
||||
@@ -314,22 +305,14 @@ func API(application *application.Application) (*echo.Echo, error) {
|
||||
audioPath := filepath.Join(application.ApplicationConfig().GeneratedContentDir, "audio")
|
||||
imagePath := filepath.Join(application.ApplicationConfig().GeneratedContentDir, "images")
|
||||
videoPath := filepath.Join(application.ApplicationConfig().GeneratedContentDir, "videos")
|
||||
threeDPath := filepath.Join(application.ApplicationConfig().GeneratedContentDir, "3d")
|
||||
|
||||
os.MkdirAll(audioPath, 0750)
|
||||
os.MkdirAll(imagePath, 0750)
|
||||
os.MkdirAll(videoPath, 0750)
|
||||
_ = os.MkdirAll(threeDPath, 0750)
|
||||
|
||||
// Go's built-in MIME table has no .glb entry and minimal containers
|
||||
// ship no /etc/mime.types, so generated GLBs would otherwise be
|
||||
// served as application/octet-stream.
|
||||
_ = mime.AddExtensionType(".glb", "model/gltf-binary")
|
||||
|
||||
e.Static("/generated-audio", audioPath)
|
||||
e.Static("/generated-images", imagePath)
|
||||
e.Static("/generated-videos", videoPath)
|
||||
e.Static("/generated-3d", threeDPath)
|
||||
}
|
||||
|
||||
// Usage recording is initialised in application/startup.go and
|
||||
|
||||
@@ -91,10 +91,6 @@ var RouteFeatureRegistry = []RouteFeature{
|
||||
// Video
|
||||
{"POST", "/video", FeatureVideo},
|
||||
|
||||
// 3D generation
|
||||
{"POST", "/3d/generations", Feature3D},
|
||||
{"POST", "/3d/remesh", Feature3D},
|
||||
|
||||
// Sound generation
|
||||
{"POST", "/v1/sound-generation", FeatureSound},
|
||||
|
||||
@@ -186,7 +182,6 @@ func APIFeatureMetas() []FeatureMeta {
|
||||
{FeatureVAD, "Voice Activity Detection", true},
|
||||
{FeatureDetection, "Detection", true},
|
||||
{FeatureVideo, "Video Generation", true},
|
||||
{Feature3D, "3D Generation", true},
|
||||
{FeatureEmbeddings, "Embeddings", true},
|
||||
{FeatureSound, "Sound Generation", true},
|
||||
{FeatureRealtime, "Realtime", true},
|
||||
|
||||
@@ -584,7 +584,6 @@ func isAPIPath(path string) bool {
|
||||
strings.HasPrefix(path, "/tts") ||
|
||||
strings.HasPrefix(path, "/vad") ||
|
||||
strings.HasPrefix(path, "/video") ||
|
||||
strings.HasPrefix(path, "/3d/") ||
|
||||
strings.HasPrefix(path, "/stores/") ||
|
||||
strings.HasPrefix(path, "/system") ||
|
||||
strings.HasPrefix(path, "/ws/") ||
|
||||
|
||||
@@ -156,16 +156,6 @@ var _ = Describe("Auth Middleware", func() {
|
||||
Expect(rec.Code).To(Equal(http.StatusUnauthorized))
|
||||
})
|
||||
|
||||
It("returns 401 for unauthenticated 3D generation requests", func() {
|
||||
rec := doRequest(app, http.MethodPost, "/3d/generations")
|
||||
Expect(rec.Code).To(Equal(http.StatusUnauthorized))
|
||||
})
|
||||
|
||||
It("returns 401 for unauthenticated 3D remesh requests", func() {
|
||||
rec := doRequest(app, http.MethodPost, "/3d/remesh")
|
||||
Expect(rec.Code).To(Equal(http.StatusUnauthorized))
|
||||
})
|
||||
|
||||
It("allows unauthenticated access to non-API paths when no legacy keys", func() {
|
||||
rec := doRequest(app, http.MethodGet, "/app")
|
||||
Expect(rec.Code).To(Equal(http.StatusOK))
|
||||
|
||||
@@ -47,7 +47,6 @@ const (
|
||||
FeatureVAD = "vad"
|
||||
FeatureDetection = "detection"
|
||||
FeatureVideo = "video"
|
||||
Feature3D = "3d"
|
||||
FeatureEmbeddings = "embeddings"
|
||||
FeatureSound = "sound"
|
||||
FeatureRealtime = "realtime"
|
||||
@@ -74,7 +73,7 @@ var GeneralFeatures = []string{FeatureFineTuning, FeatureQuantization}
|
||||
var APIFeatures = []string{
|
||||
FeatureChat, FeatureImages, FeatureAudioSpeech, FeatureAudioTranscription,
|
||||
FeatureAudioDiarization, FeatureAudioClassification,
|
||||
FeatureVAD, FeatureDetection, FeatureVideo, Feature3D, FeatureEmbeddings, FeatureSound,
|
||||
FeatureVAD, FeatureDetection, FeatureVideo, FeatureEmbeddings, FeatureSound,
|
||||
FeatureRealtime, FeatureRerank, FeatureTokenize, FeatureMCP, FeatureStores,
|
||||
FeatureFaceRecognition, FeatureVoiceRecognition, FeatureAudioTransform,
|
||||
FeaturePIIFilter,
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user