diff --git a/.agents/api-endpoints-and-auth.md b/.agents/api-endpoints-and-auth.md index 77f816849..fffee21c9 100644 --- a/.agents/api-endpoints-and-auth.md +++ b/.agents/api-endpoints-and-auth.md @@ -304,7 +304,9 @@ React pages that want to filter the ModelSelector by capability import this symb ### 4. `docs/content/` (user-facing documentation) -A new capability deserves its own page under `docs/content/features/`, plus cross-links from related features and an entry in `docs/content/whats-new.md`. See the pattern used by `face-recognition.md` / `object-detection.md`. +A new capability deserves its own page under `docs/content/features/`, plus cross-links from related features. See the pattern used by `face-recognition.md` / `object-detection.md`. + +Announcing it is the release's job, not this page's: the capability gets covered in the release blog post under `website/content/blog/`. See [preparing-a-release.md](preparing-a-release.md). `docs/content/whats-new.md` is only a pointer at the blog and GitHub Releases, so there is nothing to add there. ## Path protection rules @@ -334,7 +336,7 @@ When adding a new endpoint: - [ ] Swagger block on the handler: `@Summary`, `@Tags`, `@Param`, `@Success`, `@Router` - [ ] If new capability area (new swagger tag): entry in `instructionDefs` in `core/http/endpoints/localai/api_instructions.go` + test count bumped in `api_instructions_test.go` - [ ] If new `FLAG_*` usecase flag: matching `CAP_*` symbol exported from `core/http/react-ui/src/utils/capabilities.js` -- [ ] `docs/content/features/.md` created; cross-links from related feature pages; entry in `docs/content/whats-new.md` +- [ ] `docs/content/features/.md` created; cross-links from related feature pages; capability covered in the release blog post (see [preparing-a-release.md](preparing-a-release.md)) **Quality** - [ ] Error responses use `schema.ErrorResponse` format (or `echo.NewHTTPError` with a mapped gRPC status — see the `mapBackendError` helper in `core/http/endpoints/localai/images.go`) diff --git a/.agents/backend-signing.md b/.agents/backend-signing.md index 3abb31d7a..98c32d3e9 100644 --- a/.agents/backend-signing.md +++ b/.agents/backend-signing.md @@ -16,8 +16,7 @@ side (`pkg/oci/cosignverify` plus the gallery YAML). per-arch manifest before checking signatures. - **Storage:** Signatures are written as OCI 1.1 referrers (`--registry-referrers-mode=oci-1-1`) in the new Sigstore bundle format - (current cosign releases do this by default; no `--new-bundle-format` - flag). No `:sha256-.sig` tag clutter. + (`--new-bundle-format`). No `:sha256-.sig` tag clutter. - **Consumer:** `pkg/oci/cosignverify` discovers the bundle via the referrers API, hands it to `sigstore-go`, and verifies it against the policy declared in the gallery YAML (`Gallery.Verification`). @@ -34,14 +33,15 @@ to sign. The job needs: - `permissions: { id-token: write, contents: read }` at the job level so the runner can exchange its GitHub OIDC token for a Fulcio cert. -- `sigstore/cosign-installer@v3` step (current cosign releases already - default to the new bundle format). +- `sigstore/cosign-installer@v3` step (the pinned cosign v2 release needs + `--new-bundle-format` explicitly). - After each `docker buildx imagetools create`, resolve the resulting list digest with `docker buildx imagetools inspect --format '{{.Manifest.Digest}}'` and sign: ```sh cosign sign --yes --recursive \ + --new-bundle-format \ --registry-referrers-mode=oci-1-1 \ "${REGISTRY_REPO}@${DIGEST}" ``` @@ -70,7 +70,7 @@ entry (`backend/index.yaml`): url: github:mudler/LocalAI/backend/index.yaml@master verification: issuer: "https://token.actions.githubusercontent.com" - identity_regex: "^https://github\\.com/mudler/LocalAI/\\.github/workflows/backend_merge\\.yml@refs/heads/master$" + identity_regex: "^https://github\\.com/mudler/LocalAI/\\.github/workflows/backend_merge\\.yml@refs/(heads/master|tags/.+)$" # Optional revocation cutoff; advance during incident response. # not_before: "2026-06-01T00:00:00Z" ``` diff --git a/.agents/ci-caching.md b/.agents/ci-caching.md index 17cc6001d..1aa9c54ac 100644 --- a/.agents/ci-caching.md +++ b/.agents/ci-caching.md @@ -122,18 +122,89 @@ The per-backend prefix match only sees files under a backend's own directory, so | Changed path | Rebuilds | |---|---| -| `backend/backend.proto` | everything (all languages compile or copy it) | +| `backend/backend.proto` | nothing if the edit is additive-only, otherwise everything (see below) | | `backend/Dockerfile.` | the Linux entries whose `dockerfile:` names it | | `backend/python/common/` | Python, Linux + Darwin | -| `scripts/build/package-gpu-libs.sh` | Python, Linux only | +| `scripts/build/package-gpu-libs.sh` | every Linux entry (Python, Go and C++ all run it) | | `scripts/build/-darwin.sh` | the Darwin entries that build target routes to | | `.github/workflows/backend_build[_darwin].yml` | everything on that OS | | anything else under `scripts/build/` (except `*_test.sh`) | everything — conservative default for unclassified packaging inputs | 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. +## Content-blind PRs skip the workflows that cannot see them + +`backend_pr.yml` and `test-extra.yml` filter themselves (matrix generation and a `detect-changes` job), so a gallery-only or docs-only PR costs them about one job each. The Go and image workflows had no filter of any kind, so a one-line `gallery/index.yaml` edit queued 20 jobs, and a docs-only PR queued the same. + +This is worth more than it looks. Measured over the week to 2026-07-30, **97% of CI wall-clock is queueing, 3% is execution** (median queue ~5h against a 4-20min median job). Cutting job count is therefore the only lever that shortens feedback time; making individual jobs faster moves 3%. + +The volume is real: 13 gallery-only PRs merged that week with 10 open at once, and 78 of the 137 PRs opened were bot-generated. + +`paths-ignore` on the PR trigger of `image-pr.yml` (7 jobs), `build-test.yaml` (3), `lint.yml` (2) and `tests-e2e.yml` (1) drops 13 of those 20. The excluded set: + +| Path | Why no image or Go build can see it | +|---|---| +| `gallery/**` | Model-gallery metadata, parsed at runtime, never copied into an image | +| `docs/**`, `examples/**`, `**/*.md` | Never enter an image or a binary. `lint.yml` already excluded these before gallery was added | + +### `backend/{cpp,go,python}/**` on `image-pr.yml` and `build-test.yaml` only + +Version-pin bumps dominate PR volume: 48 `update/*` PRs in the week to 2026-07-30, from 16 pins, each a two-line diff. Most edit nothing but one `backend/*//Makefile`. + +Neither of those two workflows can observe such a change. `make build` is `go build ./cmd/local-ai`, GoReleaser builds the same plus `./cmd/launcher`, and the core image's final stage ships only `entrypoint.sh`, `healthcheck.sh` and that binary. The per-backend trees are copied into the builder but nothing in them reaches the output. + +What still triggers a full run, because none of it lives under those prefixes: + +- `backend/backend.proto` — feeds `protogen-go`, so it does change the binary. +- `go.mod` / `go.sum` — the `go mod tidy` before-hook. +- `backend/Dockerfile.*` and anything else directly under `backend/`. + +Deliberately **not** applied to: + +| Workflow | Why it must keep seeing `backend/**` | +|---|---| +| `test.yml` | `TEST_PATHS` explicitly includes `./backend/go/cloud-proxy/...`, `./backend/go/local-store/...` and `./backend/go/valkey-store/...` | +| `lint.yml` | `.golangci.yml` carries `backend/`-scoped rules, so golangci-lint covers that tree | +| `tests-e2e.yml` | The e2e suite drives real backends over gRPC | +| `backend_pr.yml` | This is the workflow whose entire job is to rebuild the changed backend | + +What still runs, and why it has to: + +| Workflow | Why it keeps running | +|---|---| +| `test.yml` (`tests`) | `core/gallery/variants_lint_test.go` reads the real `gallery/index.yaml` and asserts the index invariants (no duplicate entry names, no build claimed by two parents). This is the only schema-level check the gallery has. | +| `yaml-check.yml` (`Yamllint`) | Lints `gallery/` for syntax. | +| `backend_pr.yml`, `test-extra.yml` | Already self-filtering; they stop after the detect step. | + +Two properties this relies on: + +- `paths-ignore` skips a run only when **every** changed file matches, so a PR touching the gallery *and* Go code still runs everything. That is what makes the exclusion safe rather than a hole. +- `master` carries no branch protection and no rulesets, so a skipped workflow reports no status and nothing waits on it. If required status checks are ever introduced, these four entries must be excluded from the required set or PRs will hang on "Expected — Waiting for status to be reported". + +### `image.yml` on master push is gated too, by a job rather than a path filter + +The same reasoning applies to master pushes, and the volume is larger there: on 2026-07-30, **12 of the 23 queued `image.yml` runs** were commits like "add 1 new model to gallery" or a docs fix, each rebuilding all 18 container images. + +`image.yml` now has a `changes` job that decides once whether the push can affect any image; the other 11 jobs carry `needs: changes` plus an `if:` on its output. Verified against the shipped `Dockerfile`: the final stage copies only `entrypoint.sh`, `healthcheck.sh` and the `local-ai` binary, there is no `go:embed` of `gallery/` or `docs/`, and the gallery is fetched at runtime from `github:mudler/LocalAI/gallery/index.yaml@master`. A gallery-only commit therefore produces byte-identical images, and the gallery change reaches users through GitHub immediately whether or not an image is rebuilt. + +Two properties to preserve if you touch it: + +- **It is a job gate, not `paths-ignore`.** `paths-ignore` on `push` also applies to tag pushes, and a tag created on an existing commit carries an empty commits list, which would silently skip the release image build. The gate short-circuits to "build" for `refs/tags/*`, and for any push whose base commit is missing, zero, or unresolvable. +- **The merge jobs must name the gate explicitly.** They use `if: ${{ !cancelled() && ... }}`, and `!cancelled()` is true when a dependency is *skipped*, so without the extra condition they would run and try to merge manifest lists for images that were never built. + ## The `DEPS_REFRESH` cache-buster (Python backends) Every Python backend goes through the shared `backend/Dockerfile.python`, which ends with: @@ -169,15 +240,38 @@ RUN --mount=type=cache,target=/root/.ccache,id=-ccache-${TARGETARCH}-${ bash /usr/local/sbin/compile.sh ``` -The compile script exports `CMAKE_C/CXX/CUDA_COMPILER_LAUNCHER=ccache` so CMake threads ccache through gcc/g++/nvcc. `cache-to: type=registry,mode=max` exports the cache mount data into the registry cache, so subsequent builds restore it. +The compile script exports `CMAKE_C/CXX/CUDA_COMPILER_LAUNCHER=ccache` so CMake threads ccache through gcc/g++/nvcc. Cache scope is per `(TARGETARCH, BUILD_TYPE)` so e.g. cublas-12 doesn't share with cublas-13 (their CUDA headers differ; cross-pollination would just be cache misses anyway). -On a `LLAMA_VERSION` bump, most translation units are byte-identical to the previous version's preprocessed source — ccache returns the previous `.o` and skips the real compile. Same for LocalAI source changes that don't actually touch llama.cpp's CMake inputs. Cache scope is per `(TARGETARCH, BUILD_TYPE)` so e.g. cublas-12 doesn't share with cublas-13 (their CUDA headers differ; cross-pollination would just be cache misses anyway). +### ⚠️ This ccache does nothing in CI today + +This section previously claimed that `cache-to: type=registry,mode=max` "exports the cache mount data into the registry cache, so subsequent builds restore it". **That is not true.** BuildKit does not export the contents of a `--mount=type=cache` to a registry cache export. A cache mount lives in the builder's local state, and every CI job gets a fresh runner with a fresh builder, so `/root/.ccache` starts empty on every single build. + +Measured on 2026-07-30 from the `ccache -s` output the compile script already prints (it runs `ccache -z` first, so the numbers are per-build): + +| Job | Commit touched | Build time | ccache | +|---|---|---|---| +| 89766266951 (llama-cpp, cublas 13) | `backend/go/magpie-tts-cpp/Makefile` only | 6369s | **0 / 889 hits**, and 0 / 1778 | +| 89766267281 (llama-cpp, hipblas) | same commit | 8160s | **0 / 537 hits** | +| 90210828110 (llama-cpp, cublas 12.8) | `LLAMA_VERSION` bump | 5673s | **0 / 813 hits** | + +The first two are the decisive control: commit `90355cd44` changed exactly one file, `backend/go/magpie-tts-cpp/Makefile`, nowhere near llama.cpp. The engine source was byte-identical to the previous build, which is precisely the case this section says ccache should serve, and the hit rate was still **0.00%**. A cache that was being restored but merely matching poorly would show partial hits; 0-of-N is the signature of an empty cache. + +So the paragraph above about `LLAMA_VERSION` bumps reusing previous `.o` files describes an intended design that is not in effect. `Dockerfile.{llama-cpp,ik-llama-cpp,turboquant,bonsai,ds4,privacy-filter}` pay the ccache wrapper overhead and get nothing back. Multi-hour C++ rebuilds are recompiling identical translation units from scratch. + +**Do not "fix" this by adding cache mounts to more Dockerfiles.** Wiring the same mount into `Dockerfile.golang` (215 of the 434 matrix entries) was measured locally at 18% faster on a rebuild after a source edit, with a 71.5% ccache hit rate — but only because the local test reused one builder across both builds. In CI it would be a no-op for exactly the reason above. + +Making this actually work needs the cache to live outside the builder. The options, none of them free: + +- **ccache `remote_storage`** (ccache ≥ 4.4, HTTP or Redis backend) or **sccache** with an S3/GCS/Redis backend. Genuinely works across runners; needs a cache service to point at. quay.io is a registry, not a blob store, so the existing infra does not cover it. +- **Round-trip the cache dir through `actions/cache` on the runner**: restore it, pass it in, and export it back out via a build stage output. No external infra, but clunky, and the repo already sits at GitHub's 10 GB cache ceiling while the llama-cpp ccache alone is capped at 5 GB. + +Until one of those lands, treat C++ backend builds as always-cold and spend the effort on not running them instead (path filtering, see above). ## Composite actions Two composite actions handle runner-side prep: -- **`.github/actions/free-disk-space/action.yml`** — wraps `jlumbroso/free-disk-space@main` plus an explicit apt purge of dotnet/android/ghc/mono/etc. Reclaims ~6–10 GB on `ubuntu-latest`. No-op on self-hosted runners. Used by `backend_build.yml`, `image_build.yml`, `test.yml`, `tests-aio.yml`, etc. +- **`.github/actions/free-disk-space/action.yml`** — wraps `jlumbroso/free-disk-space@main` plus an explicit apt purge of dotnet/android/ghc/mono/etc. Reclaims ~6–10 GB on `ubuntu-latest`. No-op on self-hosted runners. Used by `backend_build.yml`, `image_build.yml` and `base-images.yml` — the jobs that actually build images. Deliberately **not** used by `test.yml`, which runs no buildx step. - **`.github/actions/setup-build-disk/action.yml`** — relocates Docker's data-root to `/mnt` on hosted X64 runners. GHA hosted `ubuntu-latest` ships ~75 GB of unused space at `/mnt`; combined with the free-disk-space cleanup this gives ~100 GB working space — enough for ROCm dev image + vLLM torch install + flash-attn intermediate layers. No-op on self-hosted and on non-X64 hosted runners. Used by `backend_build.yml`, `image_build.yml`, `base-images.yml`. Both actions run before any docker buildx step. @@ -218,10 +312,20 @@ Eviction is rarely needed in normal operation — `DEPS_REFRESH` handles weekly ## What the cache does **not** cover -- The `free-disk-space` and `setup-build-disk` composite actions run on every job — these reclaim runner-state, not Docker layers, so BuildKit caches don't apply. +- The `free-disk-space` and `setup-build-disk` composite actions run on every job — these reclaim runner-state, not Docker layers, so BuildKit caches don't apply. `test.yml` deliberately does **not** use `free-disk-space`: it runs no buildx step, and the multi-GB fixture downloads that once justified it left `make test` in the test-suite reorg. - Intermediate artifacts of `Build (PR)` are not pushed anywhere — PRs only build for verification. - Darwin builds (see below) — macOS runners have no Docker daemon, so the registry-backed BuildKit cache cannot apply. +### The Linux Go workflows set `cache: false` on purpose + +`test.yml`, `lint.yml`, `tests-e2e.yml` and friends pass `cache: false` to `actions/setup-go@v5`, unlike the darwin jobs. This looks like an oversight and is not. + +Measured over the week to 2026-07-30, the `Set up Go` step has a **median of 11 seconds** on these runners. There is essentially nothing to win: the module download is not where the time goes. The expensive steps are compilation and test execution (`Test (with coverage gate)` at ~18.6min, `Test Backend E2E` at ~14.5min), and Go's build cache would have to survive across runners to touch those. + +Enabling it also has a real cost. GitHub caps Actions cache at **10 GB per repo and the repo already sits at that ceiling** (31 entries), so every `setup-go` entry written by a branch with a distinct `go.sum` (222-375 MB on Linux, up to 1.4 GB on macOS) evicts something else. See the darwin cache budget below. + +Before re-enabling this, measure `Set up Go` again and confirm it has actually become slow. If room is needed in the 10 GB budget, the cheapest evictions are the `docker.io--tonistiigi--binfmt` entries (~30 MB each, trivially re-fetched). + ## Darwin native caches `backend_build_darwin.yml` runs natively on `macOS-14` GitHub-hosted runners — there is no Docker, no BuildKit, no cross-job registry cache. Instead, the reusable workflow uses `actions/cache@v4` for four native caches that mirror the spirit of the Linux cache (warm by default, weekly refresh for unpinned Python deps, PRs read-only). @@ -255,6 +359,26 @@ GitHub Actions caches are limited to 10 GB per repo. Steady-state worst case: ~8 One residual self-hosted reference remains in `test-extra.yml` (`tests-vibevoice-cpp-grpc-transcription` uses `bigger-runner` for the 30s JFK-decode timeout headroom). That's a separate concern. +### Small always-on jobs routed to `arc-runner-set` + +The hosted pool is shared across the whole *account*, not per repo, so a burst in one repo starves the others. On 2026-07-31 it went to **zero scheduled jobs for 35 consecutive minutes** with 39 jobs queued, while `arc-runner-set` completed 12 jobs without interruption over the same window. Actions was healthy globally at the time (other public repos were scheduling normally), so this is an account-level throttle, not an outage. + +`gh-pages.yml` (`build` + `deploy`) is therefore routed to `arc-runner-set` when `github.repository == 'mudler/LocalAI'`. It needs no fork-safety clause because it only triggers on push-to-master and `workflow_dispatch`, so it never executes pull-request code. The repository guard keeps forks (which have no such runner label) from queueing forever. It fetches its own toolchains via `setup-go` / `actions-hugo` and uses no `sudo`/`apt`. + +#### What the `arc-runner-set` image actually contains + +Measured 2026-07-31 on run `30637392862` by a preflight step, not assumed: + +| present | **absent** | +|---|---| +| `git`, `curl`, `unzip`, `tar`, `ldd`, `python3` | **`make`**, **`gcc`** | + +That is why `lint.yml` is **not** on the self-hosted pool. Both of its jobs were routed there and both failed in one second: `golangci-lint` needs `make` (for `make protogen-go`, itself needing `curl`+`unzip` to fetch protoc, and for `make lint`), and `build-scripts` additionally needs a C toolchain because the packaging-script tests compile a throwaway binary and inspect it with `ldd`. Both jobs are back on `ubuntu-latest`. + +The preflight steps were deliberately left in place. They cost about a second on the hosted pool and mean that whenever the runner image gains `make` + `gcc`, re-routing is one `runs-on:` line per job and any remaining gap reports itself by name rather than as an opaque mid-build failure. + +Note for any future re-route: `lint.yml` also triggers on `pull_request`, and a fork PR runs untrusted contributor code. That must never reach a persistent self-hosted runner, so any re-route has to stay push-only, e.g. `${{ (github.event_name == 'push' && github.repository == 'mudler/LocalAI') && 'arc-runner-set' || 'ubuntu-latest' }}`. + ## Touching the cache pipeline When changing `image_build.yml`, `backend_build.yml`, any of the `backend/Dockerfile.*` files, `Dockerfile.base-grpc-builder`, `.docker/install-base-deps.sh`, `.docker/-compile.sh`, or `scripts/changed-backends.js`: diff --git a/.agents/coding-style.md b/.agents/coding-style.md index ddee45ae2..75ab026c8 100644 --- a/.agents/coding-style.md +++ b/.agents/coding-style.md @@ -70,3 +70,37 @@ The project documentation is located in `docs/content`. When adding new features - **Configuration**: If you modify configuration options, update the relevant sections in `docs/content/`. - **Examples**: providing concrete examples (like YAML configuration blocks) is highly encouraged to help users get started quickly. - **Shortcodes**: Use `{{% notice note %}}`, `{{% notice tip %}}`, or `{{% notice warning %}}` for callout boxes. Do **not** use `{{% alert %}}` — that shortcode does not exist in this project's Hugo theme and will break the docs build. + +## React UI styling + +The React UI ships a design system in `core/http/react-ui/src/App.css`: design +tokens, form grids, data tables, stat cards, callouts, plus a small semantic +primitive layer (`.stack`, `.hstack`, `.text-note`, `.text-meta`, `.tone-*`, +`.icon-chip`). **Use it instead of `style={{ ... }}`.** Inline styles are a +spacing or colour decision made in one file, so no two pages end up sharing a +rhythm, which is the main reason the app reads as unfinished. + +Inline styles are still correct for values that are genuinely computed at +runtime: `width: ${pct}%`, a data-driven `background`, a tooltip's coordinates. +Everything else belongs in a class. + +A ratchet enforces this: + +```sh +cd core/http/react-ui +npm run lint:inline-styles # fails if the count went UP +npm run lint:inline-styles:report # per-file counts, worst first +npm run lint:inline-styles:write # refresh the baseline after converting +``` + +The gate also fails on **duplicate `className` attributes on one element**. JSX +keeps the last and silently drops the first, so `` loses its icon while passing lint, the build and the e2e +suite. Converting a style to a class on an element that already has a +`className` is the usual way to introduce one; merge them into a single +attribute instead. + +When converting a page, prefer naming the shapes it actually has +(`.p2p-diagram`, `.usage-tile`) over adding more utilities, and check whether an +existing block already covers it: the Nodes page reuses the P2P setup shapes, +and Model Editor reuses the Settings section rail. diff --git a/.agents/preparing-a-release.md b/.agents/preparing-a-release.md new file mode 100644 index 000000000..9f6f9bcfd --- /dev/null +++ b/.agents/preparing-a-release.md @@ -0,0 +1,26 @@ +# Preparing a Release + +A release is not finished when the tag is pushed. The GitHub release, the blog post and the demo clips ship together, because the changelog says what moved and the post and the clips are what make anyone care. + +## What a release must include + +1. **Labels on the merged PRs.** GitHub generates the raw notes from PR labels, so label first, generate second. Wrong labels mean a miscategorised changelog that has to be edited by hand. +2. **`RELEASE_NOTES_vX.Y.Z.md`** at the repository root, in the house style: what changed, why it matters, PR numbers so people can read the diffs. +3. **A blog post under `website/content/blog/`.** One post per release, front matter with `title`, `date`, `author`, `category: "Release"`, `tags`, `summary` and `extracss: ["blog.css"]`. Cover the two or three changes that alter what a user does day to day, not the whole changelog, and link the PR numbers. See `website/content/blog/what-landed-in-localai-4-8.md` for the shape. +4. **Demo clips for the notable features.** Anything visible (a new backend, a UI change, a new endpoint, a measured speedup) gets a short screen recording. Put the file in `website/static/media/`, reference it from the blog post, and reuse it on the marketing pages where it fits. + +A release without a post and without clips is incomplete, in the same way a user-facing code change without a docs update is incomplete. + +## Clip conventions + +- MP4, H.264, no audio track unless the feature is about audio. Keep them short (10 to 30 seconds) and loopable. +- Record the real thing. A clip from the engine's own benchmark suite or a real session, never a mockup. +- Where the change is a speedup, record both sides on the same machine on the same input, so the comparison is honest. +- Name the file after the feature, not the release (`vllm-race.mp4`, not `v4-8-demo.mp4`), so it stays reusable once the release is old. +- The marketing site plays clips with `muted loop playsinline preload="none"` and a `data-lazy` attribute, which the site's IntersectionObserver uses to play and pause them on scroll. Follow that pattern for anything you add. + +## Order of work + +Label the PRs, generate and edit the release notes, cut the draft release, record the clips while the branch is still fresh in your head, then write the post against the notes and the clips. Publishing the release and merging the post should happen on the same day. + +The `creating-localai-releases` skill drives steps 1 to 3 and captures the React UI screenshots that go into the notes. diff --git a/.agents/vllm-backend.md b/.agents/vllm-backend.md index a2b9e614e..a68bd506e 100644 --- a/.agents/vllm-backend.md +++ b/.agents/vllm-backend.md @@ -21,6 +21,17 @@ options: - reasoning_parser:qwen3 ``` +## `Options[]` doubles as CLI-style engine flags + +Beyond the parser names above, `Options[]` carries `--` prefixed engine flags (`--enable-prefix-caching`, `--kv-cache-dtype:fp8_e5m2`). `apply_options_to_engine_args` in `backend/python/common/vllm_utils.py` maps them onto `AsyncEngineArgs` fields, and it must run **before** `AsyncLLMEngine.from_engine_args()` - applying them afterwards is a silent no-op, which is exactly what issue #11130 was. + +Things to keep straight when touching this: + +- Precedence is typed proto fields → `options:` → `engine_args:`. `applyEngineArgDefaults` in `core/config/hooks_vllm.go` therefore skips seeding a production default whose key the user already set as an option, otherwise the later `engine_args:` pass would silently override them. +- Only `--` prefixed entries are engine flags; `tool_parser:`/`reasoning_parser:` and friends keep their meaning. Parser lookups accept both spellings via `normalize_option_key`. +- Unknown or uncoercible flags warn and are skipped, unlike `engine_args:` which is strict - `Options[]` is a shared bag and knows entries this mapping doesn't. +- Field types come from the annotation's *base* (`Literal["auto","float16"]` is not a float). The helper's tests are stdlib-only: `make test-python-helpers`. + Auto-defaults for known model families live in `core/config/parser_defaults.json` and are applied: - at gallery import time by `core/gallery/importers/vllm.go` - at model load time by the `vllm` / `vllm-omni` backend hook in `core/config/hooks_vllm.go` diff --git a/.docker/install-base-deps.sh b/.docker/install-base-deps.sh index 2b0e7e0c6..4331921ca 100755 --- a/.docker/install-base-deps.sh +++ b/.docker/install-base-deps.sh @@ -113,6 +113,54 @@ if [ "${BUILD_TYPE:-}" = "vulkan" ] && [ "${SKIP_DRIVERS:-false}" = "false" ]; t rm -rf /var/lib/apt/lists/* fi +# --- 2b. Intel graphics driver (BUILD_TYPE=sycl*) --- +# The Intel oneAPI base image brings the compilers and the oneAPI libraries, but +# not the driver that talks to the graphics card. The packaging step copies that +# driver into the backend, so that the backend works on a machine which has no +# Intel graphics packages of its own, for the same reason the Vulkan section +# above installs the Mesa drivers. Install it here so there is something to copy. +# +# Only the sycl builds are covered, because those are the ones whose packaging +# copies the driver. See package_intel_libs in scripts/build/package-gpu-libs.sh. +# +# The driver comes from Intel's own package repository, not from the Ubuntu +# archive. The archive has 23.43 from late 2023, which does not know any card +# released since, so a machine with a recent Intel GPU would end up carrying a +# driver that cannot drive it. Intel's repository has 25.18 for the same Ubuntu +# release. +# +# Anything that goes wrong here fails the build, on purpose. An unreachable +# repository is a passing problem that a retry fixes, whereas carrying a +# different driver than intended, or none, is a difference nobody would notice +# until a user reports an idle GPU. +if case "${BUILD_TYPE:-}" in sycl*) true;; *) false;; esac \ + && [ "${SKIP_DRIVERS:-false}" = "false" ]; then + # Ubuntu release name, which is what the repository is indexed by. + ubuntu_codename=$(. /etc/os-release && echo "${VERSION_CODENAME:-}") + if [ -z "$ubuntu_codename" ]; then + echo "ERROR: cannot tell which Ubuntu release this image is, so cannot pick the Intel driver repository" >&2 + exit 1 + fi + + # The key is armored text, which apt reads directly from a .asc file, so + # there is no need for gnupg here. "unified" is the component Intel ships + # its current driver in. + mkdir -p /usr/share/keyrings + curl -fsSL https://repositories.intel.com/gpu/intel-graphics.key \ + -o /usr/share/keyrings/intel-graphics.asc + echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.asc] https://repositories.intel.com/gpu/ubuntu ${ubuntu_codename} unified" \ + > /etc/apt/sources.list.d/intel-graphics.list + apt-get update + # The first package holds the driver OpenCL talks to, the second the driver + # Level Zero talks to. Between them they pull in the compiler and the memory + # manager that both need. + apt-get install -y --no-install-recommends \ + intel-opencl-icd \ + libze-intel-gpu1 + apt-get clean + rm -rf /var/lib/apt/lists/* +fi + # --- 3. CUDA toolkit (BUILD_TYPE=cublas|l4t) --- if { [ "${BUILD_TYPE:-}" = "cublas" ] || [ "${BUILD_TYPE:-}" = "l4t" ]; } && [ "${SKIP_DRIVERS:-false}" = "false" ]; then apt-get update diff --git a/.docker/llama-cpp-build-target.sh b/.docker/llama-cpp-build-target.sh new file mode 100755 index 000000000..efe5f0ad1 --- /dev/null +++ b/.docker/llama-cpp-build-target.sh @@ -0,0 +1,32 @@ +#!/usr/bin/env bash +set -euo pipefail + +arch=${1:?target architecture is required} +build_type=${2-} + +# SYCL compiles the whole tree with icpx -fsycl, and icpx never finishes +# ggml-cpu/arch/x86/repack.cpp at -march=sapphirerapids: the job sits on that one +# translation unit until GitHub kills it at 6h. gcc builds the same file in +# seconds, so only the SYCL images have to give up the CPU variant matrix. +# +# ROCm runs out of the same 6h budget for a different reason: volume, not a +# stall. hipcc compiles ggml's HIP kernels once per entry in AMDGPU_TARGETS, +# which is eleven architectures (gfx908 through gfx1201), and the CPU variant +# matrix lands on top of that. The job built in 2h27m before it was added and +# has been killed at exactly 6h00m on every run since, so no ROCm llama-cpp +# image has been published since 2026-08-01. +case "$build_type" in + sycl*|hipblas*) + echo llama-cpp-fallback + exit 0 + ;; +esac + +# GPU arm64 base images do not consistently provide the gcc-14 toolchain needed +# to compile ggml's armv9.2 CPU variants. Keep their portable fallback until the +# builder images can supply that compiler. +if [ "$arch" = "arm64" ] && [ -n "$build_type" ]; then + echo llama-cpp-fallback +else + echo llama-cpp-cpu-all +fi diff --git a/.docker/llama-cpp-compile.sh b/.docker/llama-cpp-compile.sh index 647a1c448..32ff2a239 100755 --- a/.docker/llama-cpp-compile.sh +++ b/.docker/llama-cpp-compile.sh @@ -18,27 +18,27 @@ if [[ -n "${CUDA_DOCKER_ARCH:-}" ]]; then fi cd /LocalAI/backend/cpp/llama-cpp -if [ -z "${BUILD_TYPE:-}" ]; then - # Pure CPU image (BUILD_TYPE empty): one build with ggml CPU_ALL_VARIANTS replaces the - # per-microarch binaries (x86: avx/avx2/avx512/fallback; arm64: armv8.x/armv9.x). ggml - # dlopens the best libggml-cpu-*.so at runtime by probing host CPU features. +BUILD_TARGET=$(/LocalAI/.docker/llama-cpp-build-target.sh "${TARGETARCH}" "${BUILD_TYPE:-}") +if [ "$BUILD_TARGET" = "llama-cpp-cpu-all" ]; then + # One build with ggml CPU_ALL_VARIANTS replaces the per-microarch binaries (x86: + # avx/avx2/avx512/fallback; arm64: armv8.x/armv9.x). BUILD_TYPE remains in the + # environment, so GPU builds retain their accelerator backend while ggml dlopens the + # best CPU library when work is offloaded to the host. # # arm64: the CPU_ALL_VARIANTS table includes armv9.2 SME variants whose -march=...+sme is # rejected by the Ubuntu 24.04 default gcc-13. gcc-14 accepts it, so build the arm64 # 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 - make llama-cpp-cpu-all -else - # GPU build (cublas/hipblas/sycl/vulkan/...): the accelerator does the compute, so a - # single fallback CPU build is enough - no per-microarch CPU variants needed. (This also - # keeps the heavy GPU backend compile from also building the whole CPU variant matrix, - # and avoids the gcc-14 apt step on GPU base images such as nvidia l4t.) - make llama-cpp-fallback fi +make "$BUILD_TARGET" make llama-cpp-grpc make llama-cpp-rpc-server diff --git a/.docker/turboquant-build-target.sh b/.docker/turboquant-build-target.sh new file mode 100755 index 000000000..8447e9347 --- /dev/null +++ b/.docker/turboquant-build-target.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +set -euo pipefail + +arch=${1:?target architecture is required} +build_type=${2-} + +# SYCL compiles the whole tree with icpx -fsycl, and icpx never finishes +# ggml-cpu/arch/x86/repack.cpp at -march=sapphirerapids: the job sits on that one +# translation unit until GitHub kills it at 6h. gcc builds the same file in +# seconds, so only the SYCL images have to give up the CPU variant matrix. +case "$build_type" in + sycl*) + echo turboquant-fallback + exit 0 + ;; +esac + +# GPU arm64 base images do not consistently provide the gcc-14 toolchain needed +# to compile ggml's armv9.2 CPU variants. Keep their portable fallback until the +# builder images can supply that compiler. +if [ "$arch" = "arm64" ] && [ -n "$build_type" ]; then + echo turboquant-fallback +else + echo turboquant-cpu-all +fi diff --git a/.docker/turboquant-compile.sh b/.docker/turboquant-compile.sh index ca6cf2690..b2ba5ffce 100755 --- a/.docker/turboquant-compile.sh +++ b/.docker/turboquant-compile.sh @@ -19,20 +19,18 @@ fi cd /LocalAI/backend/cpp/turboquant -if [ -z "${BUILD_TYPE:-}" ]; then - # Pure CPU image: one ggml CPU_ALL_VARIANTS build replaces the per-microarch binaries. +BUILD_TARGET=$(/LocalAI/.docker/turboquant-build-target.sh "${TARGETARCH}" "${BUILD_TYPE:-}") +if [ "$BUILD_TARGET" = "turboquant-cpu-all" ]; then + # BUILD_TYPE remains in the environment, so GPU builds retain their accelerator while + # ggml selects the best CPU library when model work is offloaded to the host. # arm64: the armv9.2 SME variants need gcc-14 (gcc-13 rejects +sme). if [ "${TARGETARCH}" = "arm64" ]; then + 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 - make turboquant-cpu-all -else - # GPU build (cublas/hipblas/sycl/vulkan/...): single fallback CPU build, the accelerator - # does the compute. Keeps the GPU compile from also building the CPU variant matrix and - # avoids the gcc-14 apt step on GPU base images such as nvidia l4t. - make turboquant-fallback fi +make "$BUILD_TARGET" make turboquant-grpc make turboquant-rpc-server diff --git a/.dockerignore b/.dockerignore index 159f97514..befd3706a 100644 --- a/.dockerignore +++ b/.dockerignore @@ -40,6 +40,16 @@ backend/cpp/privacy-filter/build backend/cpp/privacy-filter/grpc-server backend/cpp/privacy-filter/package +# audio-cpp: same in-place pattern. The Makefile clones audio.cpp at the pinned +# AUDIO_CPP_VERSION and the `audio.cpp:` target is the directory itself, so a +# stale host checkout COPY'd in makes the build compile against whatever commit +# the host had. build/ is worse than stale: its CMakeCache.txt records the host +# source, prefix and compiler paths, and cmake refuses to reconfigure from it. +backend/cpp/audio-cpp/audio.cpp +backend/cpp/audio-cpp/build +backend/cpp/audio-cpp/grpc-server +backend/cpp/audio-cpp/package + # Rust backend build output (sources are tracked; target/ is generated) backend/rust/*/target diff --git a/.github/backend-matrix.yml b/.github/backend-matrix.yml index 387fbdc20..f95311be8 100644 --- a/.github/backend-matrix.yml +++ b/.github/backend-matrix.yml @@ -66,6 +66,34 @@ 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: "" @@ -728,6 +756,19 @@ 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" @@ -1688,6 +1729,19 @@ 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" @@ -1701,6 +1755,19 @@ 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" @@ -3076,6 +3143,97 @@ include: dockerfile: "./backend/Dockerfile.privacy-filter" context: "./" ubuntu-version: '2404' + # audio-cpp: 0xShug0/audio.cpp, a multi-family ggml audio engine (TTS, ASR, + # VAD, diarization, source separation, music generation). + # + # These entries deliberately carry NO builder-base-image, unlike the + # privacy-filter and llama-cpp blocks above. The prebuilt + # quay.io/go-skynet/ci-cache:base-grpc-* images ship a from-source gRPC whose + # protobuf is v26, and protobuf has depended on abseil since v22. audio.cpp + # links sentencepiece with SPM_PROTOBUF_PROVIDER=package (needed to stop + # sentencepiece's vendored protobuf 3.14 from colliding with the 3.21 the + # generated backend.pb.cc is built against, which broke every nested-message + # parse), so sentencepiece then sees real abseil's + # `absl::lts_20240116::internal` alongside its own vendored plain + # `absl::internal` and every `absl::internal::` reference becomes ambiguous. + # Verified, not theorised: building against base-grpc-amd64 fails at + # sentencepiece-static.dir/error.cc.o with "reference to 'internal' is + # ambiguous". Dockerfile.audio-cpp therefore installs Ubuntu Noble's apt + # gRPC/protobuf 3.21.12 itself and has a single `builder` stage, so the + # BUILDER_BASE_IMAGE / BUILDER_TARGET / SKIP_DRIVERS build-args are never + # consumed. Same reason CUDA needs its toolkit in base-image rather than in a + # builder image: this is the ds4 shape, not the llama-cpp one. + # + # No ROCm entry: upstream has no HIP configuration. No CUDA arm64 or L4T + # entry: upstream documents and validates CUDA on x86 only. Darwin/Metal is in + # the includeDarwin matrix below, built by scripts/build/audio-cpp-darwin.sh. + # + # No vulkan entry either, though Dockerfile.audio-cpp and the backend Makefile + # both handle BUILD_TYPE=vulkan for local builds. Every other vulkan backend + # gets its Mesa ICD drivers from .docker/install-base-deps.sh, which installs + # mesa-vulkan-drivers so package-gpu-libs.sh can bundle them; this Dockerfile + # calls neither, so the image would ship a Vulkan loader that finds no GPU. No + # CI job runs a vulkan image against real hardware, so it would pass green and + # fail in users' hands. The entry comes back once the ICD question is settled. + - build-type: '' + cuda-major-version: "" + cuda-minor-version: "" + platforms: 'linux/amd64' + platform-tag: 'amd64' + tag-latest: 'auto' + tag-suffix: '-cpu-audio-cpp' + runs-on: 'ubuntu-latest' + base-image: "ubuntu:24.04" + skip-drivers: 'true' + backend: "audio-cpp" + dockerfile: "./backend/Dockerfile.audio-cpp" + context: "./" + ubuntu-version: '2404' + - build-type: '' + cuda-major-version: "" + cuda-minor-version: "" + platforms: 'linux/arm64' + platform-tag: 'arm64' + tag-latest: 'auto' + tag-suffix: '-cpu-audio-cpp' + runs-on: 'ubuntu-24.04-arm' + base-image: "ubuntu:24.04" + skip-drivers: 'true' + backend: "audio-cpp" + dockerfile: "./backend/Dockerfile.audio-cpp" + context: "./" + ubuntu-version: '2404' + # cuda-major-version is forwarded into the build (Dockerfile.audio-cpp -> the + # backend Makefile) and picks the CMAKE_CUDA_ARCHITECTURES list, which upstream + # otherwise sets to `native` and no CI runner can enumerate. cuda-minor-version + # and the base-image tag encode the same toolkit and must move together; + # nothing checks that for you. + - build-type: 'cublas' + cuda-major-version: "12" + cuda-minor-version: "8" + platforms: 'linux/amd64' + tag-latest: 'auto' + tag-suffix: '-gpu-nvidia-cuda-12-audio-cpp' + runs-on: 'ubuntu-latest' + base-image: "nvidia/cuda:12.8.1-devel-ubuntu24.04" + skip-drivers: 'true' + backend: "audio-cpp" + dockerfile: "./backend/Dockerfile.audio-cpp" + 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-audio-cpp' + runs-on: 'ubuntu-latest' + base-image: "nvidia/cuda:13.0.0-devel-ubuntu24.04" + skip-drivers: 'true' + backend: "audio-cpp" + dockerfile: "./backend/Dockerfile.audio-cpp" + context: "./" + ubuntu-version: '2404' - build-type: '' cuda-major-version: "" cuda-minor-version: "" @@ -3239,6 +3397,35 @@ 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: "" @@ -3564,6 +3751,34 @@ 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" @@ -3577,6 +3792,19 @@ 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" @@ -5408,6 +5636,35 @@ 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: "" @@ -5949,6 +6206,10 @@ 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" @@ -6030,6 +6291,18 @@ includeDarwin: - backend: "privacy-filter" tag-suffix: "-metal-darwin-arm64-privacy-filter" lang: "go" + # audio-cpp is the same shape: a C++/ggml backend built by a bespoke darwin + # script (make backends/audio-cpp-darwin), which reuses the backend's own + # package.sh so the Darwin package keeps the root-level layout the Linux image + # has (grpc-server, run.sh and assets/ in one directory, dylibs in lib/). + # No build-type: the backend Makefile turns ENGINE_ENABLE_METAL on from + # uname -s. lang=go drives runner/toolchain selection only - there is no + # backend/go/audio-cpp, which is why backend_build_darwin.yml and + # DARWIN_BESPOKE_BUILDERS in scripts/lib/backend-filter.mjs both route this + # backend away from the generic Go path. + - backend: "audio-cpp" + tag-suffix: "-metal-darwin-arm64-audio-cpp" + lang: "go" # LocalVQE has no Metal path; on Apple Silicon it builds CPU-only (GGML_METAL # OFF) but is still a native arm64 image. Uses the darwin/metal build profile. - backend: "localvqe" @@ -6126,6 +6399,10 @@ 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" diff --git a/.github/ci/gen-redirects.sh b/.github/ci/gen-redirects.sh new file mode 100755 index 000000000..14dda6f15 --- /dev/null +++ b/.github/ci/gen-redirects.sh @@ -0,0 +1,76 @@ +#!/usr/bin/env bash +# +# Generate client-side redirects for the documentation URLs that used to live at +# the site root. +# +# Until this site existed, the Hugo docs site WAS localai.io, so pages +# were published at /features/..., /getting-started/..., /faq/ and so on. The +# docs now build under /docs/, and GitHub Pages serves static files only: there +# is no server-side rewrite, no .htaccess, no _redirects. The only way to keep +# every published, bookmarked and search-indexed URL alive is to leave a real +# HTML file at the old address that sends the browser to the new one. +# +# Anything the main site already publishes wins: it owns /, /engines/, +# /blog/ and friends, so an existing file is never replaced. +# +# Usage: gen-redirects.sh [base-url] +# public-dir merged output directory (main site with docs/ inside it) +# base-url absolute or root-relative prefix the deployment is served from, +# trailing slash optional (default "/") + +set -euo pipefail + +PUBLIC_DIR=${1:?usage: gen-redirects.sh [base-url]} +BASE_URL=${2:-/} + +# Normalise to exactly one trailing slash so concatenation below is predictable. +BASE_URL="${BASE_URL%/}/" + +DOCS_DIR="${PUBLIC_DIR}/docs" + +if [ ! -d "$DOCS_DIR" ]; then + echo "gen-redirects: no docs output at ${DOCS_DIR}" >&2 + exit 1 +fi + +created=0 +skipped=0 + +# Every .html file is a reachable old URL, not just directory indexes: the +# generated model gallery ships as a bare gallery.html and used to sit at the +# root too. +while IFS= read -r src; do + rel=${src#"$DOCS_DIR"/} + dst="${PUBLIC_DIR}/${rel}" + + if [ -e "$dst" ]; then + skipped=$((skipped + 1)) + continue + fi + + # Link to the directory, not to its index.html, so the redirect target is the + # canonical URL the docs site itself advertises. + target="${BASE_URL}docs/${rel%index.html}" + + mkdir -p "$(dirname "$dst")" + printf '%s' ' + + + +Moved + + + + + +

This page moved to '"$target"'.

+ + +' > "$dst" + + created=$((created + 1)) +done </dev/null | tr -d '\r' | grep -i '^link:' || true) + if [ -z "$link" ]; then + # No Link header means a single page, so count that page directly. + gh api "${path}?per_page=100" --jq 'length' + return + fi + last=$(sed -n 's/.*[?&]page=\([0-9]*\)>; rel="last".*/\1/p' <<<"$link") + [ -n "$last" ] || { gh api "${path}?per_page=100" --jq 'length'; return; } + printf '%s\n' "$last" +} + +read -r stars forks < <(gh api "repos/${REPO}" --jq '"\(.stargazers_count) \(.forks_count)"') +contributors=$(count_via_link_header "repos/${REPO}/contributors") +releases=$(count_via_link_header "repos/${REPO}/releases") + +# Not derivable from the GitHub API, so keep whatever is already on disk. +discord=$(sed -n 's/^discord: *\([0-9]*\).*/\1/p' "$OUT" 2>/dev/null | head -1) +discord="${discord:-0}" + +for n in stars forks contributors releases; do + v="${!n}" + [[ "$v" =~ ^[0-9]+$ ]] && [ "$v" -gt 0 ] || { + echo "refusing to write: ${n} came back as '${v}'" >&2 + exit 1 + } +done + +cat > "$OUT" < docs/static/gallery.html - - name: Build site + # Two Hugo sites, one Pages artifact: the main site owns the root, + # the docs site is nested under /docs/. + - name: Build the main site + working-directory: website + run: hugo --minify --baseURL "${{ steps.pages.outputs.base_url }}/" + + - name: Build documentation site working-directory: docs run: | mkdir -p layouts/_default - hugo --minify --baseURL "${{ steps.pages.outputs.base_url }}/" + hugo --minify --baseURL "${{ steps.pages.outputs.base_url }}/docs/" + + - name: Merge documentation into the main site + run: | + mkdir -p website/public/docs + cp -R docs/public/. website/public/docs/ + + # Keeps the pre-split URLs alive; see the script header. + - name: Generate legacy URL redirects + run: .github/ci/gen-redirects.sh website/public "${{ steps.pages.outputs.base_url }}/" - name: Upload artifact uses: actions/upload-pages-artifact@v5 with: - path: docs/public + path: website/public deploy: environment: name: github-pages url: ${{ steps.deployment.outputs.page_url }} - runs-on: ubuntu-latest + # Same routing as build: a hosted slot for a ~10s deploy is exactly the kind + # of job that should not block on a starved pool. deploy-pages authenticates + # with the job's OIDC token (id-token: write above), which self-hosted + # runners issue the same way hosted ones do. + runs-on: ${{ github.repository == 'mudler/LocalAI' && 'arc-runner-set' || 'ubuntu-latest' }} needs: build steps: - name: Deploy to GitHub Pages diff --git a/.github/workflows/image-pr.yml b/.github/workflows/image-pr.yml index 146ad22cf..2605c2541 100644 --- a/.github/workflows/image-pr.yml +++ b/.github/workflows/image-pr.yml @@ -3,7 +3,28 @@ on: pull_request: - + # None of these seven image builds can observe a diff confined to these + # paths. Gallery metadata is parsed at runtime and never copied into an + # image; docs and markdown never enter one at all. Gallery content is + # still checked by yaml-check.yml and by + # core/gallery/variants_lint_test.go under 'tests'. + # + # backend/{cpp,go,python}/**: this workflow builds the core image, whose + # only compiled output is `make build` -> `go build ./cmd/local-ai`. The + # per-backend trees are copied into the builder but nothing in them reaches + # the binary or the final stage. backend/backend.proto is deliberately not + # listed: it feeds protogen-go and so does change the binary, and it does + # not live under any of these prefixes, so it still triggers a full run. + # See .agents/ci-caching.md. + paths-ignore: + - 'gallery/**' + - 'docs/**' + - 'examples/**' + - '**/*.md' + - 'backend/cpp/**' + - 'backend/go/**' + - 'backend/python/**' + concurrency: group: ci-${{ github.event.pull_request.number || github.sha }}-${{ github.repository }} cancel-in-progress: ${{ github.event_name == 'pull_request' }} diff --git a/.github/workflows/image.yml b/.github/workflows/image.yml index 655ef1b43..cd90e849b 100644 --- a/.github/workflows/image.yml +++ b/.github/workflows/image.yml @@ -13,8 +13,53 @@ cancel-in-progress: ${{ github.event_name == 'pull_request' }} jobs: - hipblas-jobs: + # Decide once whether this push can change any image. Gallery metadata is + # fetched at runtime and never baked into an image, and docs/markdown never + # enter one, so a push confined to those paths produces byte-identical + # images. On 2026-07-30, 12 of the 23 queued runs of this workflow were + # commits like "add 1 new model to gallery" or a docs fix, each rebuilding + # all 18 images. + # + # A job-level gate rather than `paths-ignore` on the trigger: paths-ignore + # would also apply to tag pushes, and a tag created on an existing commit + # carries an empty commits list, which would silently skip the release image + # build. Tags short-circuit to "build" below, as does a push whose base + # commit cannot be resolved -- the same run-everything posture the backend + # matrix filter takes for a truncated diff. + changes: if: github.repository == 'mudler/LocalAI' + runs-on: ubuntu-latest + outputs: + build: ${{ steps.decide.outputs.build }} + steps: + - uses: actions/checkout@v7 + with: + fetch-depth: 0 + - id: decide + env: + BEFORE: ${{ github.event.before }} + AFTER: ${{ github.sha }} + run: | + set -euo pipefail + emit() { echo "$2"; echo "build=$1" >> "$GITHUB_OUTPUT"; exit 0; } + case "${GITHUB_REF}" in + refs/tags/*) emit true "tag push: building every image" ;; + esac + if [ -z "${BEFORE:-}" ] || [ "${BEFORE}" = "0000000000000000000000000000000000000000" ] \ + || ! git cat-file -e "${BEFORE}^{commit}" 2>/dev/null; then + emit true "no resolvable base commit: building every image" + fi + files="$(git diff --name-only "${BEFORE}" "${AFTER}")" + echo "changed files:"; echo "${files:-}" + [ -z "${files}" ] && emit true "empty diff: building every image" + if echo "${files}" | grep -qvE '^(gallery/|docs/|examples/)|\.md$'; then + emit true "push touches image-visible content: building" + fi + emit false "only gallery/docs/markdown changed: images identical, skipping" + + hipblas-jobs: + needs: changes + if: github.repository == 'mudler/LocalAI' && needs.changes.outputs.build == 'true' uses: ./.github/workflows/image_build.yml with: tag-latest: ${{ matrix.tag-latest }} @@ -47,7 +92,8 @@ ubuntu-codename: 'noble' core-image-build: - if: github.repository == 'mudler/LocalAI' + needs: changes + if: github.repository == 'mudler/LocalAI' && needs.changes.outputs.build == 'true' uses: ./.github/workflows/image_build.yml with: tag-latest: ${{ matrix.tag-latest }} @@ -155,8 +201,8 @@ # merge whenever any matrix cell of the parent build fails or is # cancelled. Same fix as backend.yml's merge jobs — we still want to # publish the manifest list for tag-suffixes whose legs all succeeded. - if: ${{ !cancelled() && github.repository == 'mudler/LocalAI' }} - needs: core-image-build + if: ${{ !cancelled() && github.repository == 'mudler/LocalAI' && needs.changes.outputs.build == 'true' }} + needs: [changes, core-image-build] uses: ./.github/workflows/image_merge.yml with: tag-latest: 'auto' @@ -168,8 +214,8 @@ quayPassword: ${{ secrets.LOCALAI_REGISTRY_PASSWORD }} gpu-vulkan-image-merge: - if: ${{ !cancelled() && github.repository == 'mudler/LocalAI' }} - needs: core-image-build + if: ${{ !cancelled() && github.repository == 'mudler/LocalAI' && needs.changes.outputs.build == 'true' }} + needs: [changes, core-image-build] uses: ./.github/workflows/image_merge.yml with: tag-latest: 'auto' @@ -187,8 +233,8 @@ # Each merge job needs only its parent build matrix and is filtered by # tag-suffix in image_merge.yml's artifact-download pattern. gpu-nvidia-cuda-12-image-merge: - if: ${{ !cancelled() && github.repository == 'mudler/LocalAI' }} - needs: core-image-build + if: ${{ !cancelled() && github.repository == 'mudler/LocalAI' && needs.changes.outputs.build == 'true' }} + needs: [changes, core-image-build] uses: ./.github/workflows/image_merge.yml with: tag-latest: 'auto' @@ -200,8 +246,8 @@ quayPassword: ${{ secrets.LOCALAI_REGISTRY_PASSWORD }} gpu-nvidia-cuda-13-image-merge: - if: ${{ !cancelled() && github.repository == 'mudler/LocalAI' }} - needs: core-image-build + if: ${{ !cancelled() && github.repository == 'mudler/LocalAI' && needs.changes.outputs.build == 'true' }} + needs: [changes, core-image-build] uses: ./.github/workflows/image_merge.yml with: tag-latest: 'auto' @@ -213,8 +259,8 @@ quayPassword: ${{ secrets.LOCALAI_REGISTRY_PASSWORD }} gpu-intel-image-merge: - if: ${{ !cancelled() && github.repository == 'mudler/LocalAI' }} - needs: core-image-build + if: ${{ !cancelled() && github.repository == 'mudler/LocalAI' && needs.changes.outputs.build == 'true' }} + needs: [changes, core-image-build] uses: ./.github/workflows/image_merge.yml with: tag-latest: 'auto' @@ -226,8 +272,8 @@ quayPassword: ${{ secrets.LOCALAI_REGISTRY_PASSWORD }} gpu-hipblas-image-merge: - if: ${{ !cancelled() && github.repository == 'mudler/LocalAI' }} - needs: hipblas-jobs + if: ${{ !cancelled() && github.repository == 'mudler/LocalAI' && needs.changes.outputs.build == 'true' }} + needs: [changes, hipblas-jobs] uses: ./.github/workflows/image_merge.yml with: tag-latest: 'auto' @@ -239,8 +285,8 @@ quayPassword: ${{ secrets.LOCALAI_REGISTRY_PASSWORD }} nvidia-l4t-arm64-image-merge: - if: ${{ !cancelled() && github.repository == 'mudler/LocalAI' }} - needs: gh-runner + if: ${{ !cancelled() && github.repository == 'mudler/LocalAI' && needs.changes.outputs.build == 'true' }} + needs: [changes, gh-runner] uses: ./.github/workflows/image_merge.yml with: tag-latest: 'auto' @@ -252,8 +298,8 @@ quayPassword: ${{ secrets.LOCALAI_REGISTRY_PASSWORD }} nvidia-l4t-arm64-cuda-13-image-merge: - if: ${{ !cancelled() && github.repository == 'mudler/LocalAI' }} - needs: gh-runner + if: ${{ !cancelled() && github.repository == 'mudler/LocalAI' && needs.changes.outputs.build == 'true' }} + needs: [changes, gh-runner] uses: ./.github/workflows/image_merge.yml with: tag-latest: 'auto' @@ -265,7 +311,8 @@ quayPassword: ${{ secrets.LOCALAI_REGISTRY_PASSWORD }} gh-runner: - if: github.repository == 'mudler/LocalAI' + needs: changes + if: github.repository == 'mudler/LocalAI' && needs.changes.outputs.build == 'true' uses: ./.github/workflows/image_build.yml with: tag-latest: ${{ matrix.tag-latest }} diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 572764aea..35493fe4b 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -8,6 +8,9 @@ on: - 'examples/**' - 'README.md' - '**/*.md' + # golangci-lint runs new-from-merge-base, so a diff with no touched Go + # lines can only ever be a no-op. See .agents/ci-caching.md. + - 'gallery/**' push: branches: - master @@ -18,8 +21,41 @@ concurrency: jobs: golangci-lint: + # Self-hosted for PUSH only, and only in the canonical repo. + # + # This workflow also runs on pull_request, which for a fork PR means + # executing untrusted contributor code. That must never land on a + # self-hosted runner, so anything that is not a push to mudler/LocalAI stays + # on the ephemeral hosted pool. Pushes to master are trusted code that has + # already been reviewed and merged. + # + # Why at all: the hosted pool is shared account-wide and starved for 35 + # straight minutes on 2026-07-31 while arc-runner-set kept completing jobs. + # Lint is small and runs on every commit, so it is a good candidate to move + # off the contended pool. + # REVERTED to hosted: the arc-runner-set image has git, curl, unzip, tar, + # ldd and python3, but NOT make (nor gcc). Measured on run 30637392862, + # where the preflight below named both. Re-route here once the runner image + # ships a C toolchain and make; the preflight stays so the next attempt + # fails by name in one second instead of opaquely mid-build. runs-on: ubuntu-latest steps: + - name: Preflight - required host tools + # The hosted images ship these; a self-hosted container image may not. + # Check up front so a missing tool reports itself by name instead of + # surfacing as an opaque failure inside `make protogen-go` (which needs + # curl + unzip for protoc) or `make lint`. + run: | + missing="" + for t in git curl unzip make tar; do + command -v "$t" >/dev/null 2>&1 || missing="$missing $t" + done + echo "runner: ${RUNNER_NAME:-unknown} os: $(uname -sm)" + if [ -n "$missing" ]; then + echo "::error::missing required tools on this runner:$missing" + exit 1 + fi + echo "all required tools present" - uses: actions/checkout@v7 with: # Full history so golangci-lint's new-from-merge-base can reach @@ -52,8 +88,30 @@ jobs: # container build (a missing transitive dep, a partial cuDNN family). Their # shell tests need nothing but bash + gcc + ldd, so run them on every PR # rather than waiting on a multi-GB cross-arch backend image build. + # + # Push-only self-hosted routing, same fork-safety reasoning as + # golangci-lint above. + # REVERTED to hosted: the arc-runner-set image has git, curl, unzip, tar, + # ldd and python3, but NOT make (nor gcc). Measured on run 30637392862, + # where the preflight below named both. Re-route here once the runner image + # ships a C toolchain and make; the preflight stays so the next attempt + # fails by name in one second instead of opaquely mid-build. runs-on: ubuntu-latest steps: + - name: Preflight - required host tools + # This job additionally needs a C toolchain: the packaging-script tests + # compile a throwaway binary and inspect it with ldd. + run: | + missing="" + for t in git make gcc ldd python3; do + command -v "$t" >/dev/null 2>&1 || missing="$missing $t" + done + echo "runner: ${RUNNER_NAME:-unknown} os: $(uname -sm)" + if [ -n "$missing" ]; then + echo "::error::missing required tools on this runner:$missing" + exit 1 + fi + echo "all required tools present" - uses: actions/checkout@v7 - name: run packaging script tests run: make test-build-scripts @@ -66,3 +124,9 @@ jobs: node-version: '20' - name: run CI script tests run: make test-ci-scripts + + # The shared python backend helpers (Options[] parsing, engine-arg + # mapping, model reference resolution) are stdlib-only, so their tests + # ride along here instead of waiting on a multi-GB backend image build. + - name: run shared python backend helper tests + run: make test-python-helpers diff --git a/.github/workflows/refresh-site-counters.yml b/.github/workflows/refresh-site-counters.yml new file mode 100644 index 000000000..f8c3626b9 --- /dev/null +++ b/.github/workflows/refresh-site-counters.yml @@ -0,0 +1,44 @@ +name: Refresh site counters + +# The landing page shows a star count, a contributor count and a release +# count. They were typed in by hand, so they drifted the moment somebody +# forgot. This pulls the real numbers once a week and commits them only when +# they have actually moved, which in turn triggers the usual Pages deploy. + +on: + schedule: + # Mondays, 06:17 UTC. Off the hour on purpose, since the scheduler queues + # everything that asks for :00 and drops what it cannot run. + - cron: '17 6 * * 1' + workflow_dispatch: + +permissions: + contents: write + +concurrency: + group: refresh-site-counters + cancel-in-progress: false + +jobs: + refresh: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Read the counts off the GitHub API + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: ./.github/ci/refresh-site-counters.sh + + - name: Commit only if something moved + run: | + if git diff --quiet -- website/data/stats.yaml; then + echo "counters unchanged, nothing to commit" + exit 0 + fi + git diff --unified=0 -- website/data/stats.yaml + git config user.name "github-actions[bot]" + git config user.email "41898282+github-actions[bot]@users.noreply.github.com" + git add website/data/stats.yaml + git commit -m "chore(website): refresh the counters" + git push diff --git a/.github/workflows/secscan.yaml b/.github/workflows/secscan.yaml index 96daa5afa..2e4ae6d97 100644 --- a/.github/workflows/secscan.yaml +++ b/.github/workflows/secscan.yaml @@ -28,9 +28,9 @@ jobs: steps: - name: Checkout Source uses: actions/checkout@v7 - if: ${{ github.actor != 'dependabot[bot]' }} + if: ${{ !github.repository.fork && github.actor != 'dependabot[bot]' }} - name: Run Gosec Security Scanner - if: ${{ github.actor != 'dependabot[bot]' }} + if: ${{ !github.repository.fork && github.actor != 'dependabot[bot]' }} uses: securego/gosec@v2.27.1 with: # we let the report trigger content trigger a failure using the GitHub Security features. @@ -39,7 +39,7 @@ jobs: # noise, G104 unhandled errors) are inherent to that upstream code, not ours to rewrite. args: '-no-fail -exclude-dir=backend/go/supertonic -fmt sarif -out results.sarif ./...' - name: Upload SARIF file - if: ${{ github.actor != 'dependabot[bot]' }} + if: ${{ !github.repository.fork && github.actor != 'dependabot[bot]' }} uses: github/codeql-action/upload-sarif@v4 with: # Path to SARIF file relative to the root of the repository diff --git a/.github/workflows/test-extra.yml b/.github/workflows/test-extra.yml index 899b98e60..96eb57155 100644 --- a/.github/workflows/test-extra.yml +++ b/.github/workflows/test-extra.yml @@ -38,6 +38,7 @@ 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 }} @@ -935,6 +936,41 @@ 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 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 7e702a7dd..7081d630c 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -24,8 +24,13 @@ jobs: uses: actions/checkout@v7 with: submodules: true - - name: Free disk space - uses: ./.github/actions/free-disk-space + # No free-disk-space step here on purpose. That action exists to make room + # for docker buildx layers, and this job runs no buildx step. It was also + # sized for a `make test` that downloaded multi-GB GGUF/whisper fixtures + # and built llama-cpp/whisper/stablediffusion-ggml; after the test-suite + # reorg it does neither (see the Makefile test target). It cost ~3min of + # every run, and its tool-cache:true wipe also forced setup-go and + # setup-node to re-download toolchains that ship preinstalled. - name: Setup Go ${{ matrix.go-version }} uses: actions/setup-go@v5 with: diff --git a/.github/workflows/tests-e2e.yml b/.github/workflows/tests-e2e.yml index eef0a1a1c..3c1cb711c 100644 --- a/.github/workflows/tests-e2e.yml +++ b/.github/workflows/tests-e2e.yml @@ -3,6 +3,14 @@ name: 'E2E Backend Tests' on: pull_request: + # The e2e suite drives backends over gRPC directly and reads none of these + # paths, so a diff confined to them cannot move it. + # See .agents/ci-caching.md. + paths-ignore: + - 'gallery/**' + - 'docs/**' + - 'examples/**' + - '**/*.md' push: branches: - master diff --git a/.gitignore b/.gitignore index 5286461b3..a4c84a0e8 100644 --- a/.gitignore +++ b/.gitignore @@ -30,6 +30,7 @@ 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 @@ -61,6 +62,11 @@ prepare /ggml-metal.metal docs/static/gallery.html +# Hugo build output and lock files (docs/ and website/) +docs/public/ +website/public/ +.hugo_build.lock + # Protobuf generated files *.pb.go *pb2.py @@ -118,3 +124,8 @@ formal-verification/out/ # package directory itself and untrack the source. /apexentries /.github/ci/apexentries/apexentries + +# Runtime state written by `local-ai run` when it is started from the repo +# root, which is what a contributor testing a build does. Nothing under here is +# source: it is the instance's own models, outputs, traces and identity. +/data/ diff --git a/ADOPTERS.md b/ADOPTERS.md new file mode 100644 index 000000000..08169038c --- /dev/null +++ b/ADOPTERS.md @@ -0,0 +1,48 @@ +# Adopters + +Organisations running LocalAI, listed by the people who run it. + +If your organisation uses LocalAI and you are happy to say so publicly, open a +pull request adding a row to the table below. That pull request is how we know +we have permission to list you, which is why we do not add anybody ourselves. + +You do not need to be a large company, and you do not need to disclose anything +sensitive. A sentence on what you use it for is more useful to other readers +than a logo. + +## How to add yourself + +1. Add a row to the table, in alphabetical order. +2. Use your organisation's usual name and a link to your site. +3. Say briefly what you use LocalAI for, and whether it is in production. +4. Open the pull request from an account that makes it plausible you speak for + the organisation, or say in the description who you are. We may ask. + +To be removed, open a pull request deleting your row, or email +[info@localai.io](mailto:info@localai.io). We will not ask why. + +## Who is using LocalAI + + + +| Organisation | What they use it for | Status | +|---|---|---| +| _Your organisation here_ | | | + +## What this list is not + +This is not a list of everyone who has ever starred the repository, and it is +not a list of the employers of people who have contributed a patch. Both of +those are easy to scrape and neither means what a logo wall implies. + +The website shows two separate things, both of which are checkable without +anybody's permission: + +- **Engineers from these companies have contributed code.** Evidence is the + commit history plus the employer on that person's public GitHub profile. It + is a claim about a person, not about their employer. +- **These projects integrate LocalAI.** Evidence is a reference to LocalAI in + that project's own repository or documentation. + +Those two lists live in [`website/data/ecosystem.yaml`](website/data/ecosystem.yaml). +This file is the third, stronger thing: organisations that chose to say so. diff --git a/AGENTS.md b/AGENTS.md index d997813cd..dd2c79125 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -32,6 +32,7 @@ LocalAI follows the Linux kernel project's [guidelines for AI coding assistants] | [.agents/adding-gallery-models.md](.agents/adding-gallery-models.md) | Adding GGUF models from HuggingFace to the model gallery | | [.agents/localai-assistant-mcp.md](.agents/localai-assistant-mcp.md) | LocalAI Assistant chat modality — adding admin tools to the in-process MCP server, editing skill prompts, keeping REST + MCP + skills in sync | | [.agents/backend-signing.md](.agents/backend-signing.md) | Backend OCI image signing (keyless cosign + sigstore-go) — producer-side CI setup, consumer-side gallery `verification:` block, strict mode (`LOCALAI_REQUIRE_BACKEND_INTEGRITY`), revocation via `not_before` | +| [.agents/preparing-a-release.md](.agents/preparing-a-release.md) | Cutting a release: PR labels, `RELEASE_NOTES_vX.Y.Z.md`, the blog post under `website/content/blog/`, and the demo clips under `website/static/media/` | ## Quick Reference @@ -42,6 +43,7 @@ LocalAI follows the Linux kernel project's [guidelines for AI coding assistants] - **Docs (docs-with-code rule)**: When you change user-facing behavior (API endpoints, CLI flags, config keys, or features), update the corresponding page under `docs/content/` in the SAME change, not as a follow-up. A user-facing change without a matching docs update is incomplete. See also the documentation conventions in [.agents/coding-style.md](.agents/coding-style.md). - **New API endpoints**: LocalAI advertises its capability surface in several independent places — swagger `@Tags`, `/api/instructions` registry, auth `RouteFeatureRegistry`, React UI `capabilities.js`, docs. Read [.agents/api-endpoints-and-auth.md](.agents/api-endpoints-and-auth.md) and follow its checklist — missing any surface means clients, admins, and the UI won't know the endpoint exists. - **Admin endpoints → MCP tool**: every admin endpoint that an admin would manage conversationally (install/list/edit/toggle/upgrade) MUST also be exposed as an MCP tool in `pkg/mcp/localaitools/`. The LocalAI Assistant chat modality and the standalone `local-ai mcp-server` consume that package; drift between REST and MCP is a real risk. Read [.agents/localai-assistant-mcp.md](.agents/localai-assistant-mcp.md) — the `TestToolHTTPRouteMappingComplete` test fails until you wire the new tool and update the route map. +- **Releases ship with a post and clips**: a release is not done at the tag. It needs labelled PRs, `RELEASE_NOTES_vX.Y.Z.md`, a blog post under `website/content/blog/`, and a short demo clip in `website/static/media/` for each notable feature. See [.agents/preparing-a-release.md](.agents/preparing-a-release.md). - **Build**: Inspect `Makefile` and `.github/workflows/` — ask the user before running long builds - **Backend OS coverage**: a new backend must target every OS it can build for, not just Linux. `.github/backend-matrix.yml` has two matrices — `include:` (Linux) and `includeDarwin:` (macOS / Apple Silicon). Most C/C++/GGML and many Python backends build on Darwin too — wire the `includeDarwin` entry + `backend/index.yaml` `metal:` entries, or say in the PR why an OS is unsupported. See the darwin checklist in [.agents/adding-backends.md](.agents/adding-backends.md). - **Gallery variant ranking**: a gallery entry can declare `variants` (alternative builds of the same weights), and LocalAI ranks the ones a host can run by engine preference first, size second. A new backend that should be preferred on some hardware must be listed in `engineNamePreferenceRules` in `pkg/system/capabilities.go`; the sibling `backendBuildTagPreferenceRules` speaks build tags rather than engine names, and using the wrong table matches nothing without erroring. See [.agents/adding-backends.md](.agents/adding-backends.md). diff --git a/Makefile b/Makefile index 81b00d7c8..6d64540b8 100644 --- a/Makefile +++ b/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/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 +.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/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 backends/audio-cpp backends/audio-cpp-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/... +TEST_PATHS?=./api/... ./pkg/... ./core/... ./backend/go/cloud-proxy/... ./backend/go/local-store/... ./backend/go/valkey-store/... ## Coverage output and the committed baseline that CI compares against. ## The gate is strict: total coverage must never decrease (no tolerance). @@ -172,6 +172,15 @@ build-dev: ## Run LocalAI in dev mode with live reload dev-dist: $(GORELEASER) build --snapshot --clean +## PR-time variant of dev-dist: builds only the host platform instead of all +## three release targets (linux/amd64, linux/arm64, darwin/arm64). The point of +## running goreleaser on a PR is to catch a broken config or a broken +## before-hook (protogen-go, react-ui, go mod tidy), and --single-target still +## exercises every one of those. Nothing consumes a PR's cross-compiled +## binaries. master pushes and tags still run the full dev-dist/dist. +dev-dist-single: + $(GORELEASER) build --snapshot --clean --single-target + dist: $(GORELEASER) build --clean @@ -222,6 +231,14 @@ test-build-scripts: test-ci-scripts: @set -e; for t in scripts/lib/*_test.mjs; do echo "== $$t"; node --test "$$t"; done +## Runs the unit tests for the shared python backend helpers. These modules are +## pure stdlib on purpose so they run without any backend venv; the list is +## explicit because their siblings (model_identity_test) import grpc and the +## generated protobufs, which only exist inside a built backend. +PYTHON_HELPER_TESTS?=python_utils_test vllm_utils_test model_utils_test mlx_utils_test parent_watch_test +test-python-helpers: + cd backend/python/common && python3 -m unittest $(PYTHON_HELPER_TESTS) + ## Runs the core suite ($(TEST_PATHS)) with statement-coverage instrumentation ## and writes a merged profile to $(COVERAGE_PROFILE). Deliberately omits ## --fail-fast so a single failure doesn't truncate the coverage number, and @@ -386,6 +403,15 @@ 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 @@ -594,6 +620,8 @@ 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 @@ -626,6 +654,8 @@ 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. @@ -1188,6 +1218,10 @@ backends/privacy-filter-darwin: build bash ./scripts/build/privacy-filter-darwin.sh ./local-ai backends install "ocifile://$(abspath ./backend-images/privacy-filter.tar)" +backends/audio-cpp-darwin: build + bash ./scripts/build/audio-cpp-darwin.sh + ./local-ai backends install "ocifile://$(abspath ./backend-images/audio-cpp.tar)" + build-darwin-python-backend: build bash ./scripts/build/python-darwin.sh @@ -1218,6 +1252,10 @@ 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 @@ -1241,14 +1279,21 @@ BACKEND_DS4 = ds4|ds4|.|false|false # openai-privacy-filter PII/NER token classifier) — the TokenClassify RPC for # the PII redactor tier, on stock ggml with no llama.cpp carry-patches. BACKEND_PRIVACY_FILTER = privacy-filter|privacy-filter|.|false|false +# audio-cpp wraps 0xShug0/audio.cpp, a multi-family ggml audio inference engine +# (TTS, ASR, VAD, diarization, source separation, music generation). Builds +# against apt gRPC/protobuf rather than a prebuilt base-grpc image; the reason +# is on the audio-cpp block in .github/backend-matrix.yml. +BACKEND_AUDIO_CPP = audio-cpp|audio-cpp|.|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 @@ -1342,12 +1387,15 @@ $(eval $(call generate-docker-build-target,$(BACKEND_TURBOQUANT))) $(eval $(call generate-docker-build-target,$(BACKEND_BONSAI))) $(eval $(call generate-docker-build-target,$(BACKEND_DS4))) $(eval $(call generate-docker-build-target,$(BACKEND_PRIVACY_FILTER))) +$(eval $(call generate-docker-build-target,$(BACKEND_AUDIO_CPP))) $(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))) @@ -1408,7 +1456,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-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-audio-cpp ######################################################## ### Mock Backend for E2E Tests @@ -1500,7 +1548,12 @@ swagger: gen-assets: $(GOCMD) run core/dependencies_manager/manager.go webui_static.yaml core/http/static/assets -## Documentation +## Documentation and website +# The published site is two Hugo sites: website/ owns the root, docs/ is nested +# under /docs/. Serve them separately while editing; use `make site` to get the +# merged tree (including the legacy URL redirects) that GitHub Pages deploys. +SITE_BASE_URL?=http://localhost:8000 + docs/layouts/_default: mkdir -p docs/layouts/_default @@ -1512,12 +1565,30 @@ docs/public: docs/layouts/_default docs/static/gallery.html docs-clean: rm -rf docs/public + rm -rf website/public rm -rf docs/static/gallery.html .PHONY: docs docs: docs/static/gallery.html cd docs && hugo serve +.PHONY: website +website: + cd website && hugo serve + +.PHONY: site +site: docs/static/gallery.html + rm -rf website/public docs/public + cd website && hugo --minify --baseURL "$(SITE_BASE_URL)/" + cd docs && hugo --minify --baseURL "$(SITE_BASE_URL)/docs/" + mkdir -p website/public/docs + cp -R docs/public/. website/public/docs/ + ./.github/ci/gen-redirects.sh website/public "$(SITE_BASE_URL)/" + +.PHONY: site-serve +site-serve: site + cd website/public && python3 -m http.server 8000 + ######################################################## ## Platform-specific builds ######################################################## diff --git a/README.md b/README.md index d2d0c1479..290378751 100644 --- a/README.md +++ b/README.md @@ -161,7 +161,7 @@ local-ai run https://gist.githubusercontent.com/.../phi-2.yaml local-ai run oci://localai/phi-2:latest ``` -To test a running LocalAI server from the terminal, open an interactive chat session from another shell. Inside the prompt, `/models` lists installed models and `/model ` switches between them. +To work with a running LocalAI server from the terminal, start the built-in agent from another shell. It answers questions, reads your files and runs commands on your machine, asking you to approve anything that changes state. Inside a session, `/models` lists installed models and `/model ` switches between them. See the [Terminal agent](https://localai.io/docs/features/terminal-agent/) docs. ```bash # Terminal 1 @@ -195,7 +195,7 @@ For more details, see the [Getting Started guide](https://localai.io/basics/gett - **August 2025**: MLX, MLX-VLM, Diffusers, llama.cpp now supported on Apple Silicon - **July 2025**: All backends migrated outside the main binary — [lightweight, modular architecture](https://github.com/mudler/LocalAI/releases/tag/v3.2.0) -For older news and full release notes, see [GitHub Releases](https://github.com/mudler/LocalAI/releases) and the [News page](https://localai.io/basics/news/). +For older news and full release notes, see [GitHub Releases](https://github.com/mudler/LocalAI/releases) and the [blog](https://localai.io/blog/). ## Features @@ -238,13 +238,14 @@ Most backends wrap a best-in-class upstream engine. A handful of them are native | [magpie-tts.cpp](https://github.com/mudler/magpie-tts.cpp) | C++/GGML port of NVIDIA's Magpie TTS Multilingual 357M: 22.05 kHz mono text-to-speech in 5 voices and 9+ languages, with the NanoCodec neural codec and tokenizer/G2P embedded in a single GGUF | | [ced.cpp](https://github.com/localai-org/ced.cpp) | C++/GGML port of the CED audio-tagging models: sound-event classification (527-class AudioSet) over REST and the realtime API for live recognition | | [voice-detect.cpp](https://github.com/localai-org/voice-detect.cpp) | Speaker recognition and voice analysis (ECAPA-TDNN, WeSpeaker, ERes2Net, CAM++, wav2vec2 age/gender/emotion), replacing the Python speaker-recognition backend | -| [voxtral-tts.c](https://github.com/mudler/voxtral-tts.c) | Voxtral Realtime 4B speech-to-text in pure C | +| [voxtral-tts.c](https://github.com/mudler/voxtral-tts.c) | Mistral Voxtral-4B-TTS text-to-speech in pure C: 20 preset voices across 9 languages, 24 kHz WAV output, no dependencies beyond libc | | [vibevoice.cpp](https://github.com/mudler/vibevoice.cpp) | Native port of Microsoft VibeVoice for TTS (voice cloning) and long-form ASR with speaker diarization | | [rf-detr.cpp](https://github.com/localai-org/rf-detr.cpp) | Native RF-DETR object detection and instance segmentation | | [locate-anything.cpp](https://github.com/mudler/locate-anything.cpp) | Open-vocabulary object detection and visual grounding (LocateAnything-3B) | | [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) | @@ -259,7 +260,7 @@ We also maintain [apex-quant](https://github.com/localai-org/apex-quant), a per- - [Kubernetes installation](https://localai.io/basics/getting_started/#run-localai-in-kubernetes) - [Integrations & community projects](https://localai.io/docs/integrations/) - [Installation video walkthrough](https://www.youtube.com/watch?v=cMVNnlqwfw4) -- [Media & blog posts](https://localai.io/basics/news/#media-blogs-social) +- [Blog: release write-ups, benchmarks and engineering notes](https://localai.io/blog/) - [Examples](https://github.com/mudler/LocalAI-examples) — including the [realtime voice assistant demo](https://github.com/localai-org/localai-realtime-demo) (Go client for the Realtime API with tool calling) ## Team diff --git a/backend/Dockerfile.audio-cpp b/backend/Dockerfile.audio-cpp new file mode 100644 index 000000000..993e2e3b6 --- /dev/null +++ b/backend/Dockerfile.audio-cpp @@ -0,0 +1,120 @@ +ARG BASE_IMAGE=ubuntu:24.04 +ARG APT_MIRROR="" +ARG APT_PORTS_MIRROR="" + +# audio-cpp: 0xShug0/audio.cpp, a ggml audio inference framework covering TTS, +# ASR, VAD, diarization, source separation and music generation, wrapped as a +# LocalAI gRPC backend. +# +# BASE_IMAGE is ubuntu:24.04 for cpu and vulkan builds, or +# nvidia/cuda:-devel-ubuntu24.04 for cublas builds; both ship apt and +# Ubuntu Noble packages, and the CUDA base additionally provides +# /usr/local/cuda. BUILD_TYPE selects the engine backend in the Makefile: +# "" = portable CPU with all ggml CPU variants, "cublas" -> +# -DENGINE_ENABLE_CUDA=ON, "vulkan" -> -DENGINE_ENABLE_VULKAN=ON. Darwin +# (Metal) builds bypass this Dockerfile entirely. +# +# Upstream needs GCC 13 or newer, which ubuntu:24.04 and the CUDA 12/13 +# devel-ubuntu24.04 images all provide. +# +# THIS BACKEND CANNOT USE .docker/install-base-deps.sh OR THE PREBUILT +# quay.io/go-skynet/ci-cache:base-grpc-* IMAGES, AND THAT IS NOT A STYLE CHOICE. +# +# Both supply gRPC v1.65 built from source at /opt/grpc, which downstream +# Dockerfiles copy to /usr/local. That gRPC vendors protobuf v26, and protobuf +# has depended on abseil since v22: google/protobuf/message_lite.h includes +# absl/strings/cord.h. audio.cpp links sentencepiece, and our CMakeLists sets +# SPM_PROTOBUF_PROVIDER=package so sentencepiece uses the same protobuf the +# generated backend.pb.cc was built against (the alternative broke every +# nested-message parse; the full account is in backend/cpp/audio-cpp/CMakeLists.txt). +# That makes sentencepiece's init.h include the external message_lite.h while it +# still includes its own vendored mini-abseil from third_party/absl. The vendored +# copy declares `namespace absl { namespace internal { ... } }` and real abseil +# declares `namespace absl { inline namespace lts_20240116 { namespace internal +# { ... } } }`, so every `absl::internal::` reference becomes ambiguous and the +# compile dies in absl/base/casts.h. Verified, not theorised: building this image +# against the base-grpc-amd64 prebuilt fails at +# sentencepiece-static/error.cc.o with "reference to 'internal' is ambiguous". +# +# Ubuntu Noble's apt protobuf is 3.21.12, which predates the abseil dependency, +# so message_lite.h pulls in no abseil and the vendored copy is the only one in +# scope. That is also the exact protobuf/gRPC pair every unit and end-to-end run +# of this backend has been verified against. Keep it: a from-source gRPC here +# does not buy a faster build, it buys a broken one. +# +# The install-base-deps path is additionally unsafe because it drops protoc 27.1 +# into /usr/local/bin, which shadows apt's protoc on PATH and would generate +# protobuf-27 sources to be compiled against 3.21 headers. +FROM ${BASE_IMAGE} AS builder +ARG BUILD_TYPE +ARG TARGETARCH +ARG TARGETVARIANT +ARG APT_MIRROR +ARG APT_PORTS_MIRROR +# Selects the CUDA architecture list in backend/cpp/audio-cpp/Makefile. It has +# to be forwarded: upstream compiles engine_runtime for `native` when +# CMAKE_CUDA_ARCHITECTURES is unset, and no CI runner has a GPU to enumerate. +# The value is the same cuda-major-version the matrix entry declares. +ARG CUDA_MAJOR_VERSION + +ENV BUILD_TYPE=${BUILD_TYPE} \ + CUDA_MAJOR_VERSION=${CUDA_MAJOR_VERSION} \ + APT_MIRROR=${APT_MIRROR} \ + APT_PORTS_MIRROR=${APT_PORTS_MIRROR} \ + DEBIAN_FRONTEND=noninteractive \ + PATH=/usr/local/cuda/bin:${PATH} + +WORKDIR /build + +# gRPC/protobuf from apt, deliberately; see the block above. libgrpc++-dev ships +# a CMake config so find_package(gRPC CONFIG) resolves, and libprotobuf-dev +# lands in the layout CMake's FindProtobuf module expects, which matters because +# sentencepiece runs a bare find_package(Protobuf REQUIRED) with no CONFIG +# fallback of its own. +# +# BUILD_TYPE=vulkan additionally needs the loader headers and glslc; both are in +# Noble. The CUDA toolkit for BUILD_TYPE=cublas comes from BASE_IMAGE. +RUN --mount=type=bind,source=.docker/apt-mirror.sh,target=/usr/local/sbin/apt-mirror \ + sh /usr/local/sbin/apt-mirror && \ + apt-get update && \ + apt-get install -y --no-install-recommends \ + git cmake build-essential pkg-config ca-certificates \ + libgrpc++-dev libprotobuf-dev protobuf-compiler protobuf-compiler-grpc && \ + if [ "${BUILD_TYPE}" = "vulkan" ]; then \ + apt-get install -y --no-install-recommends libvulkan-dev glslc; \ + fi && \ + if [ "${TARGETARCH}" = "arm64" ]; then \ + apt-get install -y --no-install-recommends gcc-14 g++-14; \ + fi && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + +COPY . /LocalAI + +# gcc-14 on arm64, for the same reason llama-cpp does it in +# .docker/llama-cpp-compile.sh: ggml's CPU_ALL_VARIANTS table includes armv9.2 +# variants built with -march=...+sme, and Noble's default gcc-13 rejects that +# feature modifier outright ("invalid feature modifier 'sme'"). Every variant in +# the table has to COMPILE even though a host only ever dlopens the one its own +# CPU supports, so one unbuildable variant fails the whole image. +# +# ON EVERY arm64 BUILD_TYPE, which is where this differs from llama-cpp's script +# and why that difference is spelled out rather than assumed. llama-cpp only +# needs gcc-14 for its pure-CPU image because its GPU builds run +# llama-cpp-fallback, which has no variant table at all. This backend's Makefile +# sets ENGINE_ENABLE_CPU_ALL_VARIANTS for every non-Darwin build, GPU included, +# so an arm64 GPU image would hit the identical compile error. Gating this on an +# empty BUILD_TYPE would leave that trap armed for the first arm64 GPU entry +# added to the matrix, which today has none. +RUN --mount=type=cache,target=/root/.ccache,id=audio-cpp-ccache-${TARGETARCH}-${BUILD_TYPE},sharing=locked \ + if [ "${TARGETARCH}" = "arm64" ]; then \ + export CC=gcc-14 CXX=g++-14; \ + fi && \ + make -C /LocalAI/backend/cpp/audio-cpp BUILD_TYPE=${BUILD_TYPE} \ + CUDA_MAJOR_VERSION=${CUDA_MAJOR_VERSION} NATIVE=false grpc-server package + +# The package directory is the whole image: run.sh, grpc-server, the dlopened +# ggml CPU variants, the bundled loader and its library closure, and the +# bundled silero_vad / marblenet_vad assets. Nothing else exists at run time. +FROM scratch +COPY --from=builder /LocalAI/backend/cpp/audio-cpp/package/. ./ diff --git a/backend/Dockerfile.golang b/backend/Dockerfile.golang index c7dcac400..fea0804f3 100644 --- a/backend/Dockerfile.golang +++ b/backend/Dockerfile.golang @@ -248,10 +248,48 @@ RUN </dev/null 2>&1; then \ + echo "==> prebuilding engine for ${BACKEND} (cacheable layer)" && \ + make engine; \ + else \ + echo "==> ${BACKEND} has no engine target; it builds with the backend"; \ + fi + +COPY . /LocalAI + +# The engine variants built above survive this COPY (they are build outputs, not +# tracked files) and are newer than the pinned clone, so make treats them as up +# to date and goes straight to the Go binary. RUN cd /LocalAI && make protogen-go && make -C /LocalAI/backend/go/${BACKEND} build FROM scratch diff --git a/backend/Dockerfile.llama-cpp b/backend/Dockerfile.llama-cpp index 8e725ef62..2f21aaa8e 100644 --- a/backend/Dockerfile.llama-cpp +++ b/backend/Dockerfile.llama-cpp @@ -111,6 +111,10 @@ 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 diff --git a/backend/README.md b/backend/README.md index 0e92a0f03..a458da35f 100644 --- a/backend/README.md +++ b/backend/README.md @@ -56,6 +56,7 @@ 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 diff --git a/backend/backend.proto b/backend/backend.proto index 25f45eac8..2de6c8711 100644 --- a/backend/backend.proto +++ b/backend/backend.proto @@ -15,7 +15,9 @@ service Backend { rpc PredictStream(PredictOptions) returns (stream Reply) {} rpc Embedding(PredictOptions) returns (EmbeddingResult) {} rpc GenerateImage(GenerateImageRequest) returns (Result) {} + rpc UpscaleImage(UpscaleImageRequest) 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 @@ -34,6 +36,7 @@ service Backend { rpc TTSStream(TTSRequest) returns (stream Reply) {} rpc SoundGeneration(SoundGenerationRequest) returns (Result) {} rpc TokenizeString(PredictOptions) returns (TokenizationResponse) {} + rpc Detokenize(DetokenizeRequest) returns (DetokenizeResponse) {} rpc Status(HealthMessage) returns (StatusResponse) {} rpc Detect(DetectOptions) returns (DetectResponse) {} // SoundDetection runs an audio-tagging / sound-event-classification model @@ -181,6 +184,13 @@ 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 @@ -493,6 +503,11 @@ 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 @@ -508,6 +523,12 @@ 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 { @@ -617,6 +638,12 @@ message GenerateImageRequest { string ModelIdentity = 13; } +message UpscaleImageRequest { + string src = 1; // input image path + string dst = 2; // output image path + int32 scale = 3; // upscale factor (e.g. 2 or 4) +} + message GenerateVideoRequest { string prompt = 1; string negative_prompt = 2; // Negative prompt for video generation @@ -640,6 +667,20 @@ 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 params = 9; +} + message TTSRequest { string text = 1; string model = 2; @@ -763,6 +804,14 @@ message TokenizationResponse { repeated int32 tokens = 2; } +message DetokenizeRequest { + repeated int32 tokens = 1; +} + +message DetokenizeResponse { + string content = 1; +} + message MemoryUsageData { uint64 total = 1; map breakdown = 2; @@ -1089,11 +1138,28 @@ message AudioTransformRequest { string ModelIdentity = 5; } +// One named output of a transform that produces several from a single run. +// Source separation is the case that needs it: htdemucs yields drums, bass, +// other and vocals from one pass over the input. +message AudioTransformStem { + string name = 1; // the model's own stem id, e.g. "vocals" + string dst = 2; // path of the file written for that stem +} + message AudioTransformResult { string dst = 1; int32 sample_rate = 2; int32 samples = 3; bool reference_provided = 4; + // Every named output the run produced, in the model's own order, including + // the one copied into dst. Empty for a transform with a single output. + // + // It exists because dst carries one file while separation produces several, + // and running the model once per stem would cost four full separations of + // the same audio. The backend runs once, writes each stem beside dst, and + // names them here; without this field the other stems are on disk but no + // caller can find them, which is the same as not having produced them. + repeated AudioTransformStem stems = 5; } // Bidirectional streaming audio transform. The first message MUST carry a diff --git a/backend/cpp/audio-cpp/.gitignore b/backend/cpp/audio-cpp/.gitignore new file mode 100644 index 000000000..2b848d6d8 --- /dev/null +++ b/backend/cpp/audio-cpp/.gitignore @@ -0,0 +1,8 @@ +audio.cpp/ +build/ +package/ +grpc-server +backend.pb.cc +backend.pb.h +backend.grpc.pb.cc +backend.grpc.pb.h diff --git a/backend/cpp/audio-cpp/CMakeLists.txt b/backend/cpp/audio-cpp/CMakeLists.txt new file mode 100644 index 000000000..18dd5e981 --- /dev/null +++ b/backend/cpp/audio-cpp/CMakeLists.txt @@ -0,0 +1,331 @@ +cmake_minimum_required(VERSION 3.20) +project(audio-cpp-grpc-server LANGUAGES C CXX) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(TARGET grpc-server) + +set(AUDIO_CPP_DIR "${CMAKE_CURRENT_SOURCE_DIR}/audio.cpp" + CACHE PATH "Path to the pinned audio.cpp checkout") +option(AUDIO_CPP_GRPC_BUILD_TESTS "Build engine-linked ctest binaries" OFF) + +if(NOT EXISTS "${AUDIO_CPP_DIR}/CMakeLists.txt") + message(FATAL_ERROR + "AUDIO_CPP_DIR does not contain an audio.cpp checkout: ${AUDIO_CPP_DIR}. " + "Run 'make audio.cpp' first.") +endif() + +if(APPLE) + # Homebrew installs protobuf/grpc under a non-default prefix. + if(CMAKE_HOST_SYSTEM_PROCESSOR MATCHES "arm64") + set(HOMEBREW_DEFAULT_PREFIX "/opt/homebrew") + else() + set(HOMEBREW_DEFAULT_PREFIX "/usr/local") + endif() + link_directories("${HOMEBREW_DEFAULT_PREFIX}/lib") + include_directories("${HOMEBREW_DEFAULT_PREFIX}/include") +endif() + +find_package(Threads REQUIRED) +find_package(Protobuf CONFIG QUIET) +if(NOT Protobuf_FOUND) + find_package(Protobuf REQUIRED) +endif() +find_package(gRPC CONFIG QUIET) +if(NOT gRPC_FOUND) + # Reached only on distros whose grpc++ packaging ships no CMake config. + # Ubuntu's libgrpc-dev does ship one, so this is dead code on LocalAI's own + # build distro. Kept for the distros that do not. + find_library(GRPCPP_LIB grpc++ REQUIRED) + find_library(GRPCPP_REFLECTION_LIB grpc++_reflection REQUIRED) + add_library(gRPC::grpc++ INTERFACE IMPORTED) + set_target_properties(gRPC::grpc++ PROPERTIES + INTERFACE_LINK_LIBRARIES "${GRPCPP_LIB}") + add_library(gRPC::grpc++_reflection INTERFACE IMPORTED) + set_target_properties(gRPC::grpc++_reflection PROPERTIES + INTERFACE_LINK_LIBRARIES "${GRPCPP_REFLECTION_LIB}") +endif() + +find_program(_PROTOC NAMES protoc REQUIRED) +find_program(_GRPC_CPP_PLUGIN NAMES grpc_cpp_plugin REQUIRED) + +get_filename_component(HW_PROTO "${CMAKE_CURRENT_SOURCE_DIR}/../../backend.proto" ABSOLUTE) +get_filename_component(HW_PROTO_PATH "${HW_PROTO}" PATH) + +set(HW_PROTO_SRCS "${CMAKE_CURRENT_BINARY_DIR}/backend.pb.cc") +set(HW_PROTO_HDRS "${CMAKE_CURRENT_BINARY_DIR}/backend.pb.h") +set(HW_GRPC_SRCS "${CMAKE_CURRENT_BINARY_DIR}/backend.grpc.pb.cc") +set(HW_GRPC_HDRS "${CMAKE_CURRENT_BINARY_DIR}/backend.grpc.pb.h") + +add_custom_command( + OUTPUT "${HW_PROTO_SRCS}" "${HW_PROTO_HDRS}" "${HW_GRPC_SRCS}" "${HW_GRPC_HDRS}" + COMMAND ${_PROTOC} + ARGS --grpc_out "${CMAKE_CURRENT_BINARY_DIR}" + --cpp_out "${CMAKE_CURRENT_BINARY_DIR}" + -I "${HW_PROTO_PATH}" + --plugin=protoc-gen-grpc="${_GRPC_CPP_PLUGIN}" + "${HW_PROTO}" + DEPENDS "${HW_PROTO}") + +add_library(hw_grpc_proto STATIC + ${HW_GRPC_SRCS} ${HW_GRPC_HDRS} + ${HW_PROTO_SRCS} ${HW_PROTO_HDRS}) +target_include_directories(hw_grpc_proto PUBLIC ${CMAKE_CURRENT_BINARY_DIR}) +# Required on macOS: without these the Homebrew protobuf/grpc include dirs never +# reach this target and google/protobuf/runtime_version.h is not found. +target_link_libraries(hw_grpc_proto PUBLIC protobuf::libprotobuf gRPC::grpc++) + +# TWO PROTOBUF RUNTIMES IN ONE BINARY, AND THE ONE THAT WON WAS THE WRONG ONE. +# +# engine_runtime links sentencepiece, whose default SPM_PROTOBUF_PROVIDER +# ("internal") builds the protobuf-lite 3.14.0 sources vendored under +# external/sentencepiece/third_party/protobuf-lite. Our generated backend.pb.cc +# is compiled against the toolchain's protobuf 3.21.12 headers and links +# libprotobuf.so 3.21.12. Both used to end up in the executable: 476 +# google::protobuf:: symbols from that archive, 278 of them also defined by +# libprotobuf.so. +# +# The binding is decided at STATIC LINK time. Once ld pulls a sentencepiece +# member in for sentencepiece's own code, that member's protobuf definitions are +# in the executable and references from libhw_grpc_proto.a bind to them. Do NOT +# reach for -Wl,--exclude-libs: it flips those symbols to LOCAL in .dynsym and +# the breakage is unchanged, because no visibility flag revisits a static +# binding already made. +# +# What broke, measured rather than assumed: +# google::protobuf::internal::ParseContext::ParseMessage(MessageLite*, const char*) +# is what every generated _InternalParse calls for a SUBMESSAGE field and for +# nothing else. Bound to the 3.14 definition it fails, so a flat message parsed +# and every nested one did not: a TranscriptResult carrying segments serialized +# to correct bytes that the same process could not read back, and +# TranscriptLiveRequest, a oneof of submessages, could not have been parsed at +# all. 3.21 generated code was also running 3.14 arena, ArenaStringPtr and +# ExtensionSet code, which is an ABI mismatch rather than a missing feature, so +# "not observed to bite yet" was never a reason to leave it. +# +# "package" makes sentencepiece use the protobuf found above, which is the one +# the generated code was built against. It must be set before add_subdirectory, +# since that is when sentencepiece's own cache entry is created. +# +# WHAT THAT BUYS IS ONE PROTOBUF RUNTIME, not an executable free of protobuf +# symbols, and the difference matters to whoever checks this next. Measured with +# nm -C --defined-only on the linked grpc-server, 2515 google::protobuf:: +# symbols are still DEFINED in it, and that is what should be there: they are +# generated code, sentencepiece::ModelProto's own _InternalParse and +# CheckTypeAndMergeFrom among them, which name protobuf types in their +# signatures and are compiled into every user of a .proto. Expecting zero would +# send a reader looking for a regression that is not one. +# +# The claim that decides whether the ABI mismatch above is gone is the RUNTIME +# one, and it holds: google::protobuf::internal::ParseContext::ParseMessage is +# UNDEFINED in the executable, so every generated _InternalParse resolves it to +# libprotobuf.so at load instead of to a vendored 3.14 copy. No vendored +# protobuf-lite archive is pulled in at all, and citrinet_asr, which parses a +# SentencePiece ModelProto at load, tokenizes correctly as a result. +set(SPM_PROTOBUF_PROVIDER "package" CACHE STRING + "Make sentencepiece use the found protobuf, not its vendored 3.14 copy" FORCE) + +# Upstream's global add_compile_options(-Wall -Wextra -Wpedantic -pedantic-errors) +# is a directory property of the subdirectory and does not reach our targets. +# +# EXCLUDE_FROM_ALL is load-bearing, do not drop it: upstream's default target set +# includes its CLI, server, converter and test binaries, none of which we ship. +# Without it every build would compile all of them. The targets we do name in +# target_link_libraries below are still built on demand, so nothing is lost. +add_subdirectory("${AUDIO_CPP_DIR}" "${CMAKE_CURRENT_BINARY_DIR}/audio-cpp" EXCLUDE_FROM_ALL) + +add_executable(${TARGET} + grpc-server.cpp + model_options.cpp + capability_routing.cpp + family_gate.cpp + loaded_model.cpp + audio_io.cpp + audio_units.cpp + transcript_assembly.cpp + result_map.cpp + stem_selection.cpp + generation_request.cpp + stream_delta.cpp + wav_header.cpp + inference_lane.cpp + live_watchdog.cpp +) + +# Two files carry a switch over an enum with no `default:` label, deliberately, +# so that -Wswitch reports an enumerator nobody handled. -Wswitch is only a +# warning by default, and a warning in a 600-file build log is a warning nobody +# reads, so it is promoted to an error on exactly these two translation units. +# Not project-wide: upstream's own sources are not held to this, and they are +# where the churn is. +# +# loaded_model.cpp mirrors engine::runtime::VoiceTaskKind onto its own Task enum. +# Its static_asserts catch an insertion or a reorder, but an enumerator APPENDED +# after the last one shifts no value, so no assertion can see it. What does see +# it is from_engine_task's switch over the engine enum. This is the difference +# between a build failure and a backend that silently runs the wrong task. +# +# capability_routing.cpp's unsupported_surface() switches UnsupportedRpc onto the +# row of unsupported_surfaces() that explains it. Left as a warning, a sixth +# enumerator added without a row BUILDS AND SHIPS, and its trailing +# `return surfaces[0];` then answers the new RPC with AudioEncode's codec reason: +# a confident, specific and false statement about audio.cpp, on the wire, on the +# one code path whose entire job is to be truthful about what this backend +# cannot do. Verified rather than assumed: adding a sixth enumerator and building +# the shipping target produced exit 0, a binary, and one warning. A compile-time +# check is the better trade than the runtime fallback it replaced only if it is +# fatal, so here it is fatal. +if(NOT MSVC) + set_source_files_properties(loaded_model.cpp capability_routing.cpp + PROPERTIES COMPILE_OPTIONS "-Werror=switch") +endif() + +target_include_directories(${TARGET} PRIVATE + "${AUDIO_CPP_DIR}/include" + "${CMAKE_CURRENT_SOURCE_DIR}") + +# The shipping binary is held to the same bar as the tests below. Upstream's own +# add_compile_options is a property of its directory and never reached this +# target, so until now "the build was clean" meant only that nothing was being +# checked. +if(NOT MSVC) + target_compile_options(${TARGET} PRIVATE -Wall -Wextra -Wpedantic) +endif() + +target_link_libraries(${TARGET} PRIVATE + hw_grpc_proto + engine_runtime + ggml + gRPC::grpc++ + gRPC::grpc++_reflection + protobuf::libprotobuf + Threads::Threads) + +# ENGINE_ENABLE_CPU_ALL_VARIANTS builds ggml backends as shared objects that sit +# next to the binary in the package, so the binary must search its own directory. +# BUILD_WITH_INSTALL_RPATH keeps the build-tree binary at exactly "$ORIGIN". +# Upstream sets CMAKE_BUILD_WITH_INSTALL_RPATH in its own directory scope, which +# does not reach ours, so without this CMake also appends its build-tree library +# directory. That absolute build-host path would survive into the copied binary +# and let a package.sh that forgot to bundle libggml*.so still pass on the build +# machine while failing everywhere else. +set_target_properties(${TARGET} PROPERTIES + BUILD_RPATH "$ORIGIN" + INSTALL_RPATH "$ORIGIN" + BUILD_WITH_INSTALL_RPATH TRUE) + +if(AUDIO_CPP_GRPC_BUILD_TESTS) + enable_testing() + + # These are the units whose tests CANNOT run under + # backend/cpp/run-unit-tests.sh, because that script compiles each + # *_test.cpp standalone with no protobuf and no audio.cpp include path. + # They are named *_ctest.cpp so the script's glob does not pick them up and + # fail every backend's suite; everything that can be stdlib-only still is, + # and still lives in a *_test.cpp beside its unit. + add_executable(result_map_ctest + result_map_ctest.cpp + result_map.cpp + transcript_assembly.cpp + audio_units.cpp) + target_include_directories(result_map_ctest PRIVATE + "${AUDIO_CPP_DIR}/include" + "${CMAKE_CURRENT_SOURCE_DIR}" + # session.h reaches ggml.h through core/backend.h. Every other target + # here inherits that directory from the ggml target it links; this one + # links no ggml, so it has to name it. + "${AUDIO_CPP_DIR}/external/ggml/include") + # No engine_runtime: result_map touches only the plain structs in + # engine/framework/runtime/session.h, so the header is all it needs. + target_link_libraries(result_map_ctest PRIVATE + hw_grpc_proto + protobuf::libprotobuf + Threads::Threads) + target_compile_options(result_map_ctest PRIVATE -Wall -Wextra -Wpedantic) + add_test(NAME result_map COMMAND result_map_ctest) + + # Same shape as result_map_ctest: generation_request touches only the plain + # structs in engine/framework/runtime/session.h plus the generated protobuf + # messages, so the headers are all it needs and no engine_runtime is linked. + add_executable(generation_request_ctest + generation_request_ctest.cpp + generation_request.cpp) + target_include_directories(generation_request_ctest PRIVATE + "${AUDIO_CPP_DIR}/include" + "${CMAKE_CURRENT_SOURCE_DIR}" + # session.h reaches ggml.h through core/backend.h, and this target links + # no ggml, so it has to name the include directory itself. + "${AUDIO_CPP_DIR}/external/ggml/include") + target_link_libraries(generation_request_ctest PRIVATE + hw_grpc_proto + protobuf::libprotobuf + Threads::Threads) + target_compile_options(generation_request_ctest PRIVATE -Wall -Wextra -Wpedantic) + add_test(NAME generation_request COMMAND generation_request_ctest) + + add_executable(audio_io_ctest + audio_io_ctest.cpp + audio_io.cpp) + target_include_directories(audio_io_ctest PRIVATE + "${AUDIO_CPP_DIR}/include" + "${CMAKE_CURRENT_SOURCE_DIR}") + target_link_libraries(audio_io_ctest PRIVATE + engine_runtime + ggml + Threads::Threads) + target_compile_options(audio_io_ctest PRIVATE -Wall -Wextra -Wpedantic) + set_target_properties(audio_io_ctest PROPERTIES + BUILD_RPATH "$ORIGIN" + INSTALL_RPATH "$ORIGIN" + BUILD_WITH_INSTALL_RPATH TRUE) + add_test(NAME audio_io COMMAND audio_io_ctest) + + # The streaming drivers live in loaded_model.cpp, which links the engine, so + # this cannot be a standalone *_test.cpp. It builds no model and reads no + # file: LoadedModel::Session is a plain struct holding a pointer to an + # engine interface, so the drivers are exercised against fake sessions. + add_executable(streaming_driver_ctest + streaming_driver_ctest.cpp + loaded_model.cpp + capability_routing.cpp + family_gate.cpp + model_options.cpp + inference_lane.cpp) + target_include_directories(streaming_driver_ctest PRIVATE + "${AUDIO_CPP_DIR}/include" + "${CMAKE_CURRENT_SOURCE_DIR}") + target_link_libraries(streaming_driver_ctest PRIVATE + engine_runtime + ggml + Threads::Threads) + target_compile_options(streaming_driver_ctest PRIVATE -Wall -Wextra -Wpedantic) + # No "$ORIGIN" rpath override here, unlike the shipping target and unlike + # audio_io_ctest. loaded_model.cpp reaches make_default_registry, so this + # binary genuinely links libggml, and CMake's own build-tree rpath is what + # finds it: the ggml shared objects land in ${CMAKE_CURRENT_BINARY_DIR}/bin + # while the test binary sits one directory up. A build-host absolute path in + # a test binary is harmless, since package.sh ships only grpc-server, and + # forcing "$ORIGIN" here means ctest cannot start the binary at all. + add_test(NAME streaming_driver COMMAND streaming_driver_ctest) + + # Asserts that the upstream ABSENCES capability_routing.cpp's refusal + # messages rest on are still absences, by querying make_default_registry() + # rather than by re-reading upstream. This is what makes an AUDIO_CPP_VERSION + # bump that adds a codec task kind, an spk family or a streaming converter + # fail the build instead of leaving a false statement on the wire. + # + # It links engine_runtime purely to run that query, which is why it lives + # here rather than with the standalone *_test.cpp files, and it needs no + # "$ORIGIN" rpath override for the same reason streaming_driver_ctest does + # not: see the note above. + add_executable(upstream_absence_ctest upstream_absence_ctest.cpp) + target_include_directories(upstream_absence_ctest PRIVATE + "${AUDIO_CPP_DIR}/include" + "${CMAKE_CURRENT_SOURCE_DIR}") + target_link_libraries(upstream_absence_ctest PRIVATE + engine_runtime + ggml + Threads::Threads) + target_compile_options(upstream_absence_ctest PRIVATE -Wall -Wextra -Wpedantic) + add_test(NAME upstream_absence COMMAND upstream_absence_ctest) +endif() diff --git a/backend/cpp/audio-cpp/Makefile b/backend/cpp/audio-cpp/Makefile new file mode 100644 index 000000000..e33ec8d4f --- /dev/null +++ b/backend/cpp/audio-cpp/Makefile @@ -0,0 +1,172 @@ +# audio.cpp backend Makefile. +# +# Upstream pin lives below in the AUDIO_CPP_VERSION variable, so +# .github/bump_deps.sh can find and update it, matching the llama-cpp / ds4 +# convention. That script seds every line matching the variable name followed by +# an assignment, so this comment deliberately spells the name on its own: a +# comment repeating the full assignment token gets rewritten and mangled by the +# first auto-bump (backend/cpp/ds4/Makefile shows the damage). The clone +# recipe is a make target (not a prepare.sh) so 'make purge && make' is a clean +# rebuild and so the bump bot can see the pin. + +AUDIO_CPP_VERSION?=7efbb58def443722ea540d931dd3debee3e4d5e8 +AUDIO_CPP_REPO?=https://github.com/0xShug0/audio.cpp + +CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST)))) +BUILD_DIR := build + +BUILD_TYPE ?= +NATIVE ?= false +JOBS ?= $(shell nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4) + +UNAME_S := $(shell uname -s) + +# AUDIOCPP_DEPLOYMENT_BUILD compiles the model_specs/*.json catalog into +# engine_runtime, so the shipped package needs no model_specs directory and a +# safetensors model tree still resolves its family spec. +CMAKE_ARGS ?= -DCMAKE_BUILD_TYPE=Release -DAUDIOCPP_DEPLOYMENT_BUILD=ON + +# CMAKE_CUDA_ARCHITECTURES must be set explicitly for a cublas build, and this +# is not a tuning knob: upstream's CMakeLists sets CUDA_ARCHITECTURES to +# `native` on the engine_runtime target whenever the root-scope variable is +# unset (audio.cpp/CMakeLists.txt, the `if (CMAKE_CUDA_ARCHITECTURES)` branch +# next to the istft/torch_random .cu sources), and docs/build/linux.md says so +# outright: "Leave CMAKE_CUDA_ARCHITECTURES unset to build for the GPUs present +# at build time (native)". No CI runner has a GPU, so `native` has nothing to +# enumerate. ggml's own default (external/ggml/src/ggml-cuda/CMakeLists.txt) +# does not rescue this: it list(APPEND)s in the ggml subdirectory scope, which +# never reaches the root scope where the engine_runtime property is decided. +# +# The values below are ggml's list for the matching toolkit, copied rather than +# invented, so the two targets compile for exactly the same set: +# - CUDA 13 drops the Maxwell/Pascal/Volta virtual archs (50/61/70). +# - 121a-real needs CUDA >= 12.9, so the CUDA 12 list (built against 12.8) +# stops at 120a-real. +# - `a`-suffixed archs are used rather than ggml's rejected 120f-virtual: the +# `f` suffix needs CMake >= 3.31.8, and Ubuntu Noble ships 3.28.3. The +# 3.28 validator (Modules/Internal/CMakeCUDAArchitecturesValidate.cmake) +# accepts `[0-9]+a?(-real|-virtual)?`. +# +# Setting it here also pins ggml's copy, since its default is guarded by +# `if (NOT DEFINED CMAKE_CUDA_ARCHITECTURES)`. CUDA_MAJOR_VERSION is the CI +# build-arg, forwarded by Dockerfile.audio-cpp. +# +# An EMPTY major maps to `native`, NOT to the CUDA 12 list. Only CI declares a +# major; a developer running `BUILD_TYPE=cublas make` locally declares none, and +# the CUDA 12 list contains 120a-real, which needs nvcc >= 12.8. Falling through +# to it turned every local build on a CUDA 12.0-12.7 host into a compile error, +# where upstream's documented behaviour ("Leave CMAKE_CUDA_ARCHITECTURES unset +# to build for the GPUs present at build time") worked. `native` restores that. +# It does require a GPU to enumerate, so the escape hatch for a GPU-less local +# cross-build is to set CUDA_ARCHITECTURES on the command line, which the ?= +# assignments below leave untouched. +CUDA_MAJOR_VERSION ?= +ifeq ($(CUDA_MAJOR_VERSION),13) + CUDA_ARCHITECTURES ?= 75-virtual;80-virtual;86-real;89-real;120a-real;121a-real +else ifeq ($(CUDA_MAJOR_VERSION),12) + CUDA_ARCHITECTURES ?= 50-virtual;61-virtual;70-virtual;75-virtual;80-virtual;86-real;89-real;120a-real +else ifeq ($(CUDA_MAJOR_VERSION),) + CUDA_ARCHITECTURES ?= native +else ifeq ($(BUILD_TYPE),cublas) + # Gated on cublas because the variable means nothing to any other build, so a + # stray CUDA_MAJOR_VERSION in the environment must not break `make clean` or + # a CPU build. It does still error for `BUILD_TYPE=cublas make clean`, which + # is the right trade: that invocation is asking about a CUDA build tree. + $(error CUDA_MAJOR_VERSION=$(CUDA_MAJOR_VERSION) has no architecture list here (12 and 13 do). Leave it empty for a native build, or pass CUDA_ARCHITECTURES explicitly.) +endif + +ifeq ($(BUILD_TYPE),cublas) + CMAKE_ARGS += -DENGINE_ENABLE_CUDA=ON "-DCMAKE_CUDA_ARCHITECTURES=$(CUDA_ARCHITECTURES)" +else ifeq ($(BUILD_TYPE),vulkan) + CMAKE_ARGS += -DENGINE_ENABLE_VULKAN=ON +else ifeq ($(UNAME_S),Darwin) + # Metal. ggml embeds the shader library by default (GGML_METAL_EMBED_LIBRARY + # defaults to GGML_METAL), so the package needs no .metallib beside the + # binary. Darwin builds go through scripts/build/audio-cpp-darwin.sh. + CMAKE_ARGS += -DENGINE_ENABLE_METAL=ON + # AppleClang ships no OpenMP runtime and Homebrew's libomp is keg-only, so + # neither libomp.dylib nor omp.h is symlinked into /opt/homebrew and CMake's + # FindOpenMP cannot find them on its own (the workflow's `brew link libomp` + # is a no-op for a keg-only formula, and its failure is swallowed). + # audio.cpp calls find_package(OpenMP REQUIRED COMPONENTS CXX) whenever + # ENGINE_ENABLE_OPENMP is ON, so with no hint the macOS build dies at + # configure time before compiling anything. OpenMP_ROOT is honoured by the + # find_library/find_path calls inside FindOpenMP under CMP0074, which is NEW + # here because audio.cpp requires CMake 3.20. + # + # If the keg is absent, turn OpenMP off rather than fail: the tree's only + # include is guarded by #ifdef _OPENMP and a #pragma omp without + # -fopenmp is simply ignored, so an OpenMP-less build is CORRECT. It is not + # cheap, though: 108 `#pragma omp` directives across ~30 files (roformer, + # demucs, chatterbox, moss, supertonic, seed_vc, framework/audio/dsp) are + # compiled out, and clang says nothing about an ignored omp pragma unless + # -Wsource-uses-openmp is on. A green package that is quietly single-threaded + # in every host DSP loop gets blamed on Metal, not on packaging, so the + # fallback announces itself. + ifeq ($(origin LIBOMP_PREFIX),undefined) + LIBOMP_PREFIX := $(shell brew --prefix libomp 2>/dev/null) + endif + # Nested ifneq rather than $(and ...): $(and) needs GNU make 3.81, and while + # that is what Apple ships, an older make expands it to empty and would take + # the OpenMP-OFF branch with no way to tell that from a genuinely missing + # keg. Two plain conditionals cannot fail that way. + LIBOMP_USABLE := + ifneq ($(wildcard $(LIBOMP_PREFIX)/lib/libomp.dylib),) + ifneq ($(wildcard $(LIBOMP_PREFIX)/include/omp.h),) + LIBOMP_USABLE := yes + endif + endif + ifeq ($(LIBOMP_USABLE),yes) + CMAKE_ARGS += "-DOpenMP_ROOT=$(LIBOMP_PREFIX)" + else + $(warning audio-cpp: libomp not found at '$(LIBOMP_PREFIX)'; building without OpenMP (single-threaded host DSP). Install it with `brew install libomp`, or set LIBOMP_PREFIX.) + CMAKE_ARGS += -DENGINE_ENABLE_OPENMP=OFF + endif +else + # Portable Linux CPU. Upstream wires this to GGML_BACKEND_DL + + # GGML_CPU_ALL_VARIANTS + $ORIGIN rpath, so one build serves every CPU + # tier instead of an AVX-tier image fan-out. + CMAKE_ARGS += -DENGINE_ENABLE_CPU_ALL_VARIANTS=ON +endif + +ifneq ($(NATIVE),true) + CMAKE_ARGS += -DENGINE_ENABLE_NATIVE_CPU=OFF +endif + +.PHONY: all grpc-server package test test-engine clean purge +all: grpc-server + +# Clone the upstream source at the pinned commit. The directory is the target +# so make only re-clones when it is missing. After bumping AUDIO_CPP_VERSION, +# run 'make purge && make' to refetch. +audio.cpp: + mkdir -p audio.cpp + cd audio.cpp && \ + git init -q && \ + git remote add origin $(AUDIO_CPP_REPO) && \ + git fetch --depth 1 origin $(AUDIO_CPP_VERSION) && \ + git checkout FETCH_HEAD + +grpc-server: audio.cpp + mkdir -p $(BUILD_DIR) + cd $(BUILD_DIR) && cmake $(CMAKE_ARGS) $(CURRENT_MAKEFILE_DIR) && \ + cmake --build . --config Release -j $(JOBS) + cp $(BUILD_DIR)/grpc-server grpc-server + +package: grpc-server + bash package.sh + +test: + @echo "audio-cpp: standalone unit tests run from the repo root via 'make test-backend-cpp'" + +# Engine-linked tests. Needs the upstream checkout and a full engine build. +test-engine: audio.cpp + mkdir -p $(BUILD_DIR) + cd $(BUILD_DIR) && cmake $(CMAKE_ARGS) -DAUDIO_CPP_GRPC_BUILD_TESTS=ON $(CURRENT_MAKEFILE_DIR) && \ + cmake --build . --config Release -j $(JOBS) && ctest --output-on-failure --no-tests=error + +clean: + rm -rf $(BUILD_DIR) grpc-server package + +purge: clean + rm -rf audio.cpp diff --git a/backend/cpp/audio-cpp/audio_io.cpp b/backend/cpp/audio-cpp/audio_io.cpp new file mode 100644 index 000000000..709d1864e --- /dev/null +++ b/backend/cpp/audio-cpp/audio_io.cpp @@ -0,0 +1,124 @@ +#include "audio_io.h" + +#include "loaded_model.h" + +#include "engine/framework/audio/conversion.h" +#include "engine/framework/audio/wav_reader.h" +#include "engine/framework/audio/wav_writer.h" + +#include +#include +#include + +namespace audiocpp_backend { + +engine::runtime::AudioBuffer read_audio_file(const std::string &path, + int target_sample_rate) { + if (path.empty()) { + throw ConfigError("audio-cpp: no input audio path was supplied"); + } + std::error_code ec; + const bool present = std::filesystem::exists(std::filesystem::path(path), ec); + if (ec) { + // exists() returning false with ec set does NOT mean the file is + // absent, it means the question could not be answered: most often a + // parent directory is not searchable. Reporting that as "does not + // exist" sends the operator after the file when the fault is the + // permissions on the directory above it. + throw ConfigError("audio-cpp: cannot stat input audio " + path + ": " + + ec.message()); + } + if (!present) { + throw ConfigError("audio-cpp: input audio does not exist: " + path); + } + engine::audio::WavData wav; + try { + wav = engine::audio::read_wav_f32(std::filesystem::path(path)); + } catch (const std::exception &err) { + throw ConfigError("audio-cpp: cannot read " + path + + " as WAV: " + err.what()); + } + if (wav.sample_rate <= 0) { + throw ConfigError("audio-cpp: " + path + + " declares a non-positive sample rate; every " + "timestamp derived from it would be zero"); + } + // AudioBuffer's own default is 1, and a reader that reports 0 channels + // still gave us an interleaving of one. Normalised before the conversion + // below rather than after, because mixdown_interleaved_to_mono_average + // throws on a non-positive channel count. + if (wav.channels <= 0) { + wav.channels = 1; + } + + engine::runtime::AudioBuffer buffer; + if (target_sample_rate <= 0) { + buffer.sample_rate = wav.sample_rate; + buffer.channels = wav.channels; + buffer.samples = std::move(wav.samples); + return buffer; + } + + buffer.sample_rate = target_sample_rate; + buffer.channels = 1; + try { + // A no-op copy when the rates already match, so the common 16 kHz + // upload pays only the mono mixdown it would have paid inside the + // family anyway. + buffer.samples = + engine::audio::convert_wav_to_mono_linear_resampled(wav, target_sample_rate); + } catch (const std::exception &err) { + // ConfigError, so this is INVALID_ARGUMENT rather than INTERNAL. What + // reaches here is a malformed input: a sample count that is not a whole + // number of frames is the realistic one, and it is the uploader's file + // that is truncated, not this backend that is broken. + throw ConfigError("audio-cpp: cannot resample " + path + " from " + + std::to_string(wav.sample_rate) + " Hz to " + + std::to_string(target_sample_rate) + + " Hz: " + err.what()); + } + return buffer; +} + +void write_audio_file(const std::string &path, + const engine::runtime::AudioBuffer &audio) { + if (path.empty()) { + throw ConfigError("audio-cpp: no output path was supplied"); + } + const std::filesystem::path destination(path); + if (destination.has_parent_path()) { + // Best effort: a failure here shows up as a write failure below, with a + // message naming the file the caller actually asked for. + std::error_code ec; + std::filesystem::create_directories(destination.parent_path(), ec); + } + try { + engine::audio::write_pcm16_wav(destination, audio.sample_rate, + audio.channels > 0 ? audio.channels : 1, + audio.samples); + } catch (const std::exception &err) { + // NOT a ConfigError, and the distinction is not cosmetic. The + // destination is chosen by LocalAI rather than by the caller: it is a + // unique name inside GeneratedContentDir. A failure to write it is a + // full disk, a permission fault on the server's own directory, or a bad + // mount, none of which the caller can fix or is to blame for. As a + // ConfigError this surfaced as INVALID_ARGUMENT, which tells a client + // its request was wrong and not to retry; a plain runtime_error maps to + // INTERNAL, which is both true and retryable. The empty path above + // stays INVALID_ARGUMENT, because that one really is a malformed + // request. + throw std::runtime_error("audio-cpp: cannot write " + path + ": " + + err.what()); + } +} + +engine::runtime::AudioBuffer buffer_from_mono(std::vector samples, + int sample_rate) { + engine::runtime::AudioBuffer buffer; + buffer.sample_rate = sample_rate; + buffer.channels = 1; + buffer.samples = std::move(samples); + return buffer; +} + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/audio_io.h b/backend/cpp/audio-cpp/audio_io.h new file mode 100644 index 000000000..be01734b6 --- /dev/null +++ b/backend/cpp/audio-cpp/audio_io.h @@ -0,0 +1,61 @@ +#pragma once + +// Thin wrappers over the framework's public audio IO. Engine-linked, so this +// unit is built and tested through the CMake target rather than by +// backend/cpp/run-unit-tests.sh. The pure part of the arithmetic these +// wrappers feed lives in audio_units, which is stdlib-only and does have a +// standalone test. + +#include "engine/framework/runtime/session.h" + +#include +#include + +namespace audiocpp_backend { + +// Reads a WAV file. Throws ConfigError when the file is missing, is not +// readable as WAV, or declares a non-positive sample rate: all three are +// user-fixable input problems rather than backend faults. +// +// A declared sample rate of zero is refused rather than passed on, because +// every downstream conversion in audio_units answers 0 for a non-positive rate. +// Accepting it would turn a corrupt header into a response full of zero +// timestamps, which reads as a real answer. +// +// `target_sample_rate` is the rate the CALLER needs, in Hz: +// +// 0 (or negative) keep the file's own rate and channel count. +// positive downmix to mono and resample to that rate. Resampling is +// skipped when the file already declares it, so passing the +// rate a route needs costs nothing on the common input. +// +// It is a parameter, and not a constant inside this function, because the +// routes that read audio do not agree on an answer. Speech routes want 16 kHz +// mono; source separation does not, and folding a 44.1 kHz stereo input to +// 16 kHz mono for demucs or roformer would destroy the very thing they separate +// (both refuse a rate other than their own outright). Making the caller name +// the rate keeps that decision where the route is known. +// +// Downmixing along with the resample is not an extra liberty: every family a +// positive rate is used for (silero_vad, sortformer_diar and every ASR family) +// begins by calling the same mixdown_interleaved_to_mono_average on whatever it +// is given. Doing it once here produces the identical samples and halves the +// buffer that is then moved through the request. +engine::runtime::AudioBuffer read_audio_file(const std::string &path, + int target_sample_rate); + +// Writes 16-bit PCM WAV, creating parent directories. +// +// Throws ConfigError, i.e. INVALID_ARGUMENT, ONLY for an empty path, which is a +// malformed request. Every other failure throws a plain runtime_error, i.e. +// INTERNAL: the destination is LocalAI's own generated-content directory and +// not anything the caller named, so a full disk or a permission fault there is +// a server fault and is worth retrying, which is the opposite of what +// INVALID_ARGUMENT tells a client. +void write_audio_file(const std::string &path, + const engine::runtime::AudioBuffer &audio); + +engine::runtime::AudioBuffer buffer_from_mono(std::vector samples, + int sample_rate); + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/audio_io_ctest.cpp b/backend/cpp/audio-cpp/audio_io_ctest.cpp new file mode 100644 index 000000000..8e6b1e1a6 --- /dev/null +++ b/backend/cpp/audio-cpp/audio_io_ctest.cpp @@ -0,0 +1,217 @@ +// Tests for audio_io's reading contract, and in particular for the resampling +// that keeps a 44.1 or 48 kHz upload from reaching a family that only accepts +// 16 kHz. +// +// NAMED _ctest AND NOT _test ON PURPOSE: see the note at the top of +// result_map_ctest.cpp. This file links the audio.cpp engine, so it is built +// and run by ctest, not by backend/cpp/run-unit-tests.sh. +// +// make -C backend/cpp/audio-cpp test-engine + +#include "audio_io.h" + +#include "loaded_model.h" + +#include +#include +#include +#include +#include +#include + +static int failures = 0; + +static void check(bool ok, const std::string &name) { + if (!ok) { + failures++; + fprintf(stderr, "FAIL: %s\n", name.c_str()); + } else { + fprintf(stderr, "ok: %s\n", name.c_str()); + } +} + +using namespace audiocpp_backend; + +// A one-second tone, interleaved across `channels`. Real audio rather than +// silence so a resample that dropped its input would be visible as a flat +// buffer, not just as a different length. +static engine::runtime::AudioBuffer tone(int sample_rate, int channels, + float seconds) { + engine::runtime::AudioBuffer buffer; + buffer.sample_rate = sample_rate; + buffer.channels = channels; + const auto frames = + static_cast(static_cast(sample_rate) * seconds); + buffer.samples.reserve(frames * static_cast(channels)); + for (size_t frame = 0; frame < frames; ++frame) { + const float value = 0.5f * std::sin(2.0f * 3.14159265f * 220.0f * + static_cast(frame) / + static_cast(sample_rate)); + for (int channel = 0; channel < channels; ++channel) { + buffer.samples.push_back(value); + } + } + return buffer; +} + +static float peak(const std::vector &samples) { + float highest = 0.0f; + for (const float sample : samples) { + highest = std::max(highest, std::abs(sample)); + } + return highest; +} + +static std::filesystem::path scratch_dir() { + const auto dir = std::filesystem::temp_directory_path() / "audiocpp-io-ctest"; + std::filesystem::create_directories(dir); + return dir; +} + +// The I2 fixture. Before the resample this returned a 44.1 kHz buffer, which +// silero_vad and sortformer_diar both reject with a plain runtime_error, which +// the server maps to INTERNAL. A 44.1 kHz WAV is an ordinary upload. +static void test_441k_stereo_is_read_as_16k_mono() { + const auto path = scratch_dir() / "input-44100-stereo.wav"; + write_audio_file(path.string(), tone(44100, 2, 1.0f)); + + const auto audio = read_audio_file(path.string(), 16000); + check(audio.sample_rate == 16000, "44.1 kHz input is resampled to 16 kHz"); + check(audio.channels == 1, "stereo input is downmixed to mono"); + // Linear resampling lands within a sample or two of the exact ratio. + const auto frames = static_cast(audio.samples.size()); + check(frames > 15990 && frames < 16010, + "one second in stays one second out"); + check(peak(audio.samples) > 0.2f, + "the resampled buffer still carries the signal"); +} + +static void test_48k_is_read_as_16k() { + const auto path = scratch_dir() / "input-48000-mono.wav"; + write_audio_file(path.string(), tone(48000, 1, 0.5f)); + + const auto audio = read_audio_file(path.string(), 16000); + check(audio.sample_rate == 16000, "48 kHz input is resampled to 16 kHz"); + const auto frames = static_cast(audio.samples.size()); + check(frames > 7990 && frames < 8010, "half a second in, half a second out"); +} + +// The common case: the upload is already 16 kHz mono, and nothing is resampled. +static void test_16k_mono_passes_through_unchanged() { + const auto path = scratch_dir() / "input-16000-mono.wav"; + const auto source = tone(16000, 1, 1.0f); + write_audio_file(path.string(), source); + + const auto audio = read_audio_file(path.string(), 16000); + check(audio.sample_rate == 16000, "16 kHz stays 16 kHz"); + check(audio.channels == 1, "mono stays mono"); + check(audio.samples.size() == source.samples.size(), + "a matching rate resamples nothing"); +} + +// Rate 0 means "give me the file as it is", which is what a source separation +// route needs: demucs and roformer refuse anything but their own 44.1 kHz and +// work on stereo, so the reader must not force them to 16 kHz mono. +static void test_zero_target_keeps_the_native_format() { + const auto path = scratch_dir() / "input-native.wav"; + write_audio_file(path.string(), tone(44100, 2, 0.25f)); + + const auto audio = read_audio_file(path.string(), 0); + check(audio.sample_rate == 44100, "a zero target keeps the file's rate"); + check(audio.channels == 2, "a zero target keeps the file's channels"); +} + +static void test_missing_file_is_a_config_error() { + bool threw_config_error = false; + try { + read_audio_file((scratch_dir() / "does-not-exist.wav").string(), 16000); + } catch (const ConfigError &) { + threw_config_error = true; + } catch (const std::exception &) { + // Any other type maps to INTERNAL, which is what this asserts against. + } + check(threw_config_error, "a missing input file is INVALID_ARGUMENT, not INTERNAL"); +} + +static void test_unreadable_file_is_a_config_error() { + const auto path = scratch_dir() / "not-a-wav.wav"; + { + FILE *file = fopen(path.string().c_str(), "wb"); + if (file != nullptr) { + fputs("this is not a RIFF header", file); + fclose(file); + } + } + bool threw_config_error = false; + try { + read_audio_file(path.string(), 16000); + } catch (const ConfigError &) { + threw_config_error = true; + } catch (const std::exception &) { + } + check(threw_config_error, "a non-WAV input is INVALID_ARGUMENT, not INTERNAL"); +} + +// The write side of the same distinction. The destination is LocalAI's own +// generated-content directory, not a caller-supplied path, so a failure to +// write it is a server fault: INTERNAL, which a client may retry, and not +// INVALID_ARGUMENT, which tells it the request itself was wrong. +static void test_write_failure_is_not_a_config_error() { + // A regular file where a directory has to be. ENOTDIR defeats root as well + // as an ordinary user, unlike a chmod, which CI running as root would walk + // straight through. + const auto blocker = scratch_dir() / "blocking-file"; + { + FILE *file = fopen(blocker.string().c_str(), "wb"); + if (file != nullptr) { + fputs("not a directory", file); + fclose(file); + } + } + const auto path = blocker / "nested" / "out.wav"; + + bool threw_config_error = false; + bool threw_something = false; + try { + write_audio_file(path.string(), tone(16000, 1, 0.05f)); + } catch (const ConfigError &) { + threw_config_error = true; + threw_something = true; + } catch (const std::exception &) { + threw_something = true; + } + check(threw_something, "an unwritable destination is reported at all"); + check(!threw_config_error, + "a failed write is INTERNAL, not INVALID_ARGUMENT: the caller did not " + "choose the destination and cannot fix it"); + check(!std::filesystem::exists(path), "and nothing was written"); +} + +static void test_empty_output_path_is_a_config_error() { + // The one write failure that IS the caller's: no path at all. + bool threw_config_error = false; + try { + write_audio_file("", tone(16000, 1, 0.05f)); + } catch (const ConfigError &) { + threw_config_error = true; + } catch (const std::exception &) { + } + check(threw_config_error, "an empty output path stays INVALID_ARGUMENT"); +} + +int main() { + test_441k_stereo_is_read_as_16k_mono(); + test_48k_is_read_as_16k(); + test_16k_mono_passes_through_unchanged(); + test_zero_target_keeps_the_native_format(); + test_missing_file_is_a_config_error(); + test_unreadable_file_is_a_config_error(); + test_write_failure_is_not_a_config_error(); + test_empty_output_path_is_a_config_error(); + if (failures) { + fprintf(stderr, "%d check(s) failed\n", failures); + return 1; + } + fprintf(stderr, "all audio_io checks passed\n"); + return 0; +} diff --git a/backend/cpp/audio-cpp/audio_units.cpp b/backend/cpp/audio-cpp/audio_units.cpp new file mode 100644 index 000000000..99a096cbd --- /dev/null +++ b/backend/cpp/audio-cpp/audio_units.cpp @@ -0,0 +1,139 @@ +#include "audio_units.h" + +#include +#include +#include + +namespace audiocpp_backend { + +std::int64_t interleaved_frame_count(std::size_t sample_count, int channels) { + const std::size_t lanes = channels > 0 ? static_cast(channels) + : static_cast(1); + // Truncating division is deliberate: a trailing partial frame is not a + // position every channel reached, so counting it would overstate the length. + return static_cast(sample_count / lanes); +} + +std::int64_t samples_to_nanoseconds(std::int64_t samples, int sample_rate) { + if (sample_rate <= 0) { + return 0; + } + // Split into whole seconds plus a remainder so the intermediate product + // cannot overflow on long recordings, and so rates like 44100 stay exact. + // The remainder division truncates deliberately: that matches Go's + // time.Duration conventions and keeps successive sample indices monotonic. + const std::int64_t rate = static_cast(sample_rate); + const std::int64_t whole_seconds = samples / rate; + const std::int64_t remainder = samples % rate; + return whole_seconds * 1000000000LL + (remainder * 1000000000LL) / rate; +} + +float samples_to_seconds(std::int64_t samples, int sample_rate) { + if (sample_rate <= 0) { + return 0.0f; + } + return static_cast(static_cast(samples) / + static_cast(sample_rate)); +} + +std::int64_t seconds_to_samples(double seconds, int sample_rate) { + // !(seconds > 0.0) rather than seconds <= 0.0: every comparison against NaN + // is false, so the <= form lets NaN reach the cast below, which is undefined + // behaviour and lands on INT64_MIN in practice. This is the one entry point + // fed by untrusted-shaped input (a float-seconds timestamp off the wire, or + // a boundary from a model that diverged), and a hugely negative sample index + // used later as an offset or a length is a wild pointer rather than merely a + // wrong timestamp. + if (sample_rate <= 0 || !(seconds > 0.0)) { + return 0; + } + const double scaled = seconds * static_cast(sample_rate); + // Bound before the cast for the same reason: converting a double at or above + // 2^63 (infinity included) is undefined behaviour, so saturate instead. + const double limit = + static_cast(std::numeric_limits::max()); + if (scaled >= limit) { + return std::numeric_limits::max(); + } + // Round rather than truncate: these functions exist to cross the float + // seconds boundary the VAD and diarize messages use, so a value that came + // from samples_to_seconds converts back to the sample it started as. + // Truncation lost one sample about half the time, starting at n=1. + // + // That round trip is exact only below roughly 2^23 samples. Past that the + // float samples_to_seconds returns can no longer resolve adjacent indices + // and the trip fails whatever the rounding. Both the first failing INDEX + // and the duration it stands for depend on the rate, so they are listed per + // rate rather than folded into one range; measured: + // + // 16 kHz 16384001 samples 17.1 min + // 44.1 kHz 11289602 samples 4.3 min + // 48 kHz 12288002 samples 4.3 min + // 96 kHz 12288002 samples 2.1 min + // + // The shortest recording this bites is therefore a couple of minutes of + // 96 kHz audio. It is a property of the float seconds API itself, not of + // the rounding here, and it is why nothing should use these to carry a + // sample-accurate position in a long recording. + return static_cast(std::llround(scaled)); +} + +std::vector s16le_to_f32(const std::string &bytes) { + std::vector samples; + const size_t count = bytes.size() / 2; + samples.reserve(count); + for (size_t i = 0; i < count; ++i) { + const auto low = static_cast(bytes[i * 2]); + const auto high = static_cast(bytes[i * 2 + 1]); + const auto raw = static_cast( + static_cast(low) | + (static_cast(high) << 8)); + // 32768 on decode against 32767 on encode is deliberate, not a typo. + // 32768 is what keeps INT16_MIN at exactly -1.0 and every other code + // inside the [-1, 1] range this header promises; dividing by 32767 + // would decode INT16_MIN to -1.00003. See f32_to_s16le for the other + // half of the pair. The cost is that a round trip shrinks a sample by + // 32767/32768, well under one LSB. + samples.push_back(static_cast(raw) / 32768.0f); + } + return samples; +} + +std::string f32_to_s16le(const std::vector &samples) { + std::string bytes; + bytes.reserve(samples.size() * 2); + for (const float sample : samples) { + // NaN maps to silence. A NaN sample rendered as a full-scale click is + // worse audio than a dropped one, and this unit converts audio that may + // have originated off the wire. + // + // This guard also removes what used to be a spelling hazard in the + // clamp below. std::min and std::max return their first argument when + // the comparison is false, and every comparison against NaN is false, + // so before this branch existed the choice of spelling silently decided + // whether a NaN reached std::lround, whose result is unspecified for + // NaN. These three leaked it, the last being the idiomatic C++17 way to + // write a clamp and so the likeliest future edit: + // std::min(std::max(sample, -1.0f), 1.0f) + // std::max(std::min(sample, 1.0f), -1.0f) + // std::clamp(sample, -1.0f, 1.0f) + // The order is no longer load-bearing now that the guard runs first, + // but the history is why the guard is here, so do not drop it. + if (std::isnan(sample)) { + bytes.push_back(0); + bytes.push_back(0); + continue; + } + const float clamped = std::max(-1.0f, std::min(1.0f, sample)); + // 32767 rather than 32768 so +1.0 saturates at INT16_MAX instead of + // overflowing to INT16_MIN. See s16le_to_f32 for why decode differs. + const auto value = + static_cast(std::lround(clamped * 32767.0f)); + const auto raw = static_cast(value); + bytes.push_back(static_cast(raw & 0xFF)); + bytes.push_back(static_cast((raw >> 8) & 0xFF)); + } + return bytes; +} + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/audio_units.h b/backend/cpp/audio-cpp/audio_units.h new file mode 100644 index 000000000..376d12d4c --- /dev/null +++ b/backend/cpp/audio-cpp/audio_units.h @@ -0,0 +1,50 @@ +#pragma once + +// Time and sample-format conversion between audio.cpp's runtime types (sample +// indices, float PCM) and LocalAI's proto types. Standard library only. +// +// LocalAI uses three different time units: +// TranscriptSegment / TranscriptWord start,end : int64 nanoseconds +// VADSegment start,end : float seconds +// DiarizeSegment start,end : float seconds + +#include +#include +#include +#include + +namespace audiocpp_backend { + +// Frames in an interleaved buffer of `sample_count` floats laid out across +// `channels` channels. A frame is one per-channel position, which is the unit +// every duration and every span boundary in this backend is expressed in, so a +// stereo buffer must not report twice its real length: feeding sample_count +// straight to samples_to_seconds makes a 3 second stereo clip come back as 6. +// +// A non-positive channel count is treated as mono, matching +// engine::runtime::AudioBuffer's own default of 1 and keeping a reader that +// reports 0 channels from dividing by zero. +std::int64_t interleaved_frame_count(std::size_t sample_count, int channels); + +// Returns 0 when sample_rate is not positive rather than dividing by zero. +// Uses integer arithmetic so 44.1 kHz does not lose precision. +std::int64_t samples_to_nanoseconds(std::int64_t samples, int sample_rate); + +float samples_to_seconds(std::int64_t samples, int sample_rate); + +// Rounds to nearest. Negative seconds and NaN both yield 0, and a value too +// large to convert saturates at INT64_MAX rather than overflowing. Round trips +// with samples_to_seconds only below roughly 2^23 samples, past which the float +// seconds can no longer resolve adjacent sample indices. +std::int64_t seconds_to_samples(double seconds, int sample_rate); + +// Decodes little-endian signed 16-bit PCM. A trailing odd byte is dropped. +std::vector s16le_to_f32(const std::string &bytes); + +// Encodes to little-endian signed 16-bit PCM, clamping to [-1, 1] first so an +// overshooting sample saturates instead of wrapping to the opposite sign. +// A NaN sample encodes to 0, on the grounds that silence beats a full-scale +// click. +std::string f32_to_s16le(const std::vector &samples); + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/audio_units_test.cpp b/backend/cpp/audio-cpp/audio_units_test.cpp new file mode 100644 index 000000000..ec74eeee0 --- /dev/null +++ b/backend/cpp/audio-cpp/audio_units_test.cpp @@ -0,0 +1,226 @@ +// Unit tests for audio_units. Standard library only. The harness compiles this +// as a single translation unit, so the implementation is included directly. + +#include "audio_units.cpp" + +#include +#include +#include +#include +#include +#include + +static int failures = 0; + +static void check(bool ok, const std::string &name) { + if (!ok) { + failures++; + fprintf(stderr, "FAIL: %s\n", name.c_str()); + } else { + fprintf(stderr, "ok: %s\n", name.c_str()); + } +} + +static bool close_to(float a, float b, float tol) { return std::fabs(a - b) <= tol; } + +using namespace audiocpp_backend; + +static void test_nanoseconds() { + // LocalAI TranscriptSegment/TranscriptWord times are nanoseconds + // (Go reads them as time.Duration). + check(samples_to_nanoseconds(16000, 16000) == 1000000000LL, "1s at 16k is 1e9 ns"); + check(samples_to_nanoseconds(8000, 16000) == 500000000LL, "0.5s at 16k"); + check(samples_to_nanoseconds(0, 16000) == 0, "zero samples is zero ns"); + check(samples_to_nanoseconds(1000, 0) == 0, "zero sample rate yields zero, not UB"); + // 44.1 kHz must not lose precision to float arithmetic. + check(samples_to_nanoseconds(44100, 44100) == 1000000000LL, "1s at 44.1k"); + check(samples_to_nanoseconds(22050, 44100) == 500000000LL, "0.5s at 44.1k"); + // The cases above all land on values a float happens to hold exactly, so + // they do not actually rule float arithmetic out. These do: + // a fraction that does not divide evenly, and a duration whose magnitude + // exceeds a float's 24-bit mantissa at nanosecond resolution. + check(samples_to_nanoseconds(44099, 44100) == 999977324LL, + "44.1k fraction is exact, not rounded through a float"); + check(samples_to_nanoseconds(44100LL * 3600, 44100) == 3600000000000LL, + "one hour at 44.1k is exact to the nanosecond"); + // A naive samples * 1e9 would overflow int64 here; the split into whole + // seconds plus a remainder is what keeps this correct. + check(samples_to_nanoseconds(44100LL * 360000, 44100) == 360000000000000LL, + "100 hours at 44.1k does not overflow"); + // Double arithmetic is close enough to pass everything above, but still + // truncates this one a nanosecond short. Integer division does not. + check(samples_to_nanoseconds(4004, 8000) == 500500000LL, + "0.5005s at 8k is exact to the nanosecond"); + // Truncation, not rounding: this matches Go's time.Duration conventions and + // keeps successive sample indices monotonic. The exact value here is + // 22675.7...; rounding to nearest would give 22676. + check(samples_to_nanoseconds(1, 44100) == 22675LL, + "a sub-nanosecond fraction truncates rather than rounding up"); +} + +static void test_seconds() { + check(close_to(samples_to_seconds(24000, 24000), 1.0f, 1e-6f), "1s at 24k"); + check(close_to(samples_to_seconds(12000, 24000), 0.5f, 1e-6f), "0.5s at 24k"); + check(close_to(samples_to_seconds(100, 0), 0.0f, 1e-6f), "zero sample rate is 0s"); + check(seconds_to_samples(1.0, 16000) == 16000, "1s to samples at 16k"); + check(seconds_to_samples(0.5, 16000) == 8000, "0.5s to samples at 16k"); + check(seconds_to_samples(1.0, 0) == 0, "zero sample rate yields zero samples"); + check(seconds_to_samples(-1.0, 16000) == 0, "negative seconds clamps to zero"); + + // seconds_to_samples is the one entry point fed by untrusted-shaped input: + // a float-seconds timestamp off the wire, or a VAD boundary from a model + // that diverged. A hugely negative sample index used later as an offset or + // a length is a wild pointer, not merely a wrong timestamp. + const double nan_seconds = std::numeric_limits::quiet_NaN(); + const double inf_seconds = std::numeric_limits::infinity(); + const std::int64_t max_samples = std::numeric_limits::max(); + check(seconds_to_samples(nan_seconds, 16000) == 0, "NaN seconds yields zero"); + check(seconds_to_samples(inf_seconds, 16000) == max_samples, + "infinite seconds saturates instead of overflowing"); + check(seconds_to_samples(1e30, 16000) == max_samples, + "out of range seconds saturates instead of overflowing"); + check(seconds_to_samples(-inf_seconds, 16000) == 0, + "negative infinity clamps to zero"); + + // Crossing the float-seconds boundary and back is the expected round trip + // for the VAD and diarize messages, so it must not lose a sample. + // Truncation loses one about half the time, starting at n=1. + check(seconds_to_samples(samples_to_seconds(1, 44100), 44100) == 1, + "one sample survives the seconds round trip at 44.1k"); + check(seconds_to_samples(samples_to_seconds(1, 16000), 16000) == 1, + "one sample survives the seconds round trip at 16k"); + check(seconds_to_samples(samples_to_seconds(4001, 8000), 8000) == 4001, + "4001 samples survive the seconds round trip at 8k"); +} + +static void test_s16le_round_trip() { + const std::vector original = {0.0f, 0.5f, -0.5f, 1.0f, -1.0f}; + const std::string encoded = f32_to_s16le(original); + check(encoded.size() == original.size() * 2, "two bytes per sample"); + + const std::vector decoded = s16le_to_f32(encoded); + check(decoded.size() == original.size(), "round trip keeps the sample count"); + for (size_t i = 0; i < original.size(); ++i) { + // 16-bit quantisation: one LSB is ~3.05e-5. Guard the index so a short + // result reports a named failure instead of aborting the whole suite. + check(i < decoded.size() && close_to(decoded[i], original[i], 1e-4f), + "round trip preserves sample " + std::to_string(i)); + } +} + +static void test_s16le_endianness() { + // 0.5 encodes to 16384 = 0x4000, little endian is 0x00 0x40. + const std::string encoded = f32_to_s16le({0.5f}); + check(encoded.size() == 2, "one sample is two bytes"); + check(static_cast(encoded[0]) == 0x00, "low byte first"); + check(static_cast(encoded[1]) == 0x40, "high byte second"); +} + +static void test_s16le_clamping() { + // Values outside [-1, 1] must clamp, not wrap around to the opposite sign. + const std::string encoded = f32_to_s16le({2.0f, -2.0f}); + const std::vector decoded = s16le_to_f32(encoded); + check(decoded.size() == 2, "two samples survive clamping"); + check(decoded.size() > 0 && decoded[0] > 0.99f, + "positive overshoot clamps to full scale"); + check(decoded.size() > 1 && decoded[1] < -0.99f, + "negative overshoot clamps to full scale"); +} + +static void test_s16le_decode_range() { + // INT16_MIN is the one value that pins the decode scale. Dividing by 32767 + // instead of 32768 would decode it to -1.00003, outside the [-1, 1] range + // the header promises, and every other test would still pass. + const std::vector decoded = s16le_to_f32(std::string("\x00\x80", 2)); + check(decoded.size() == 1, "INT16_MIN decodes to one sample"); + check(decoded.size() == 1 && decoded[0] == -1.0f, + "INT16_MIN decodes to exactly -1.0, not past full scale"); +} + +static void test_s16le_nan_input() { + // A NaN sample must not reach std::lround, whose result is unspecified for + // NaN. Asserting a range is not enough to pin this: the three outcomes the + // plausible clamp spellings produce (full scale, negative full scale, zero) + // are all finite and all inside [-1, 1], so a range check passes for every + // one of them. Only an exact value distinguishes them. + // NaN maps to silence, not to full scale: a NaN sample rendered as a + // full-scale click is worse audio than a dropped one, and this unit + // converts audio that may have originated off the wire. + // + // volatile so the NaN cannot be constant-folded, which would let the + // compiler evaluate the conversion at compile time and raise no + // floating-point exception at run time for the check below to observe. + volatile float nan_source = std::numeric_limits::quiet_NaN(); + const std::vector input = {nan_source}; + + std::feclearexcept(FE_ALL_EXCEPT); + const std::string encoded = f32_to_s16le(input); + const bool raised_invalid = std::fetestexcept(FE_INVALID) != 0; + const std::vector decoded = s16le_to_f32(encoded); + + check(decoded.size() == 1, "a NaN sample still encodes to one sample"); + check(decoded.size() == 1 && decoded[0] == 0.0f, + "a NaN sample encodes to exactly zero, not to a full-scale click"); + // Independent of the value: a quiet NaN raises invalid-operation as soon as + // it reaches any ordered comparison, which is what std::min and std::max + // use, so this fails unless the NaN is diverted before the clamp runs at + // all. That is what stops the explicit guard from being dropped in favour + // of a clamp spelling that happens to yield zero. + check(!raised_invalid, + "encoding a NaN sample raises no invalid-operation exception"); +} + +static void test_s16le_odd_length() { + // A truncated frame must drop the dangling byte rather than read past it. + const std::string odd(5, '\0'); + check(s16le_to_f32(odd).size() == 2, "odd byte count drops the trailing byte"); + check(s16le_to_f32(std::string()).empty(), "empty input yields no samples"); +} + +static void test_interleaved_frame_count() { + // Mono is a pass-through, which is the only case the VAD path exercises. + check(interleaved_frame_count(16000, 1) == 16000, "mono frames equal samples"); + // The case that matters: a stereo buffer holds two floats per position, so a + // one second 16 kHz stereo clip is 32000 floats and still one second. Handing + // the raw float count to samples_to_seconds reports two seconds instead. + check(interleaved_frame_count(32000, 2) == 16000, + "stereo frames are half the samples"); + check(samples_to_seconds(interleaved_frame_count(32000, 2), 16000) == 1.0f, + "a one second stereo clip measures one second, not two"); + check(interleaved_frame_count(48000, 3) == 16000, + "three channels divide by three"); + // engine::runtime::AudioBuffer defaults channels to 1, but a reader is free + // to report 0, and dividing by that is undefined rather than merely wrong. + check(interleaved_frame_count(1000, 0) == 1000, + "zero channels is treated as mono"); + check(interleaved_frame_count(1000, -2) == 1000, + "a negative channel count is treated as mono"); + // A dangling partial frame is not a position every channel reached. + check(interleaved_frame_count(3, 2) == 1, + "a trailing partial frame is not counted"); + check(interleaved_frame_count(0, 2) == 0, "an empty buffer has no frames"); + // Past 2^32 floats, so a size_t narrowed to 32 bits on the way in, or a + // signed 32-bit intermediate, shows up here rather than in a multi-hour + // recording nobody tests with. + check(interleaved_frame_count(static_cast(9000000000ULL), 2) == + 4500000000LL, + "a buffer beyond 2^32 floats counts frames without truncating"); +} + +int main() { + test_interleaved_frame_count(); + test_nanoseconds(); + test_seconds(); + test_s16le_round_trip(); + test_s16le_endianness(); + test_s16le_clamping(); + test_s16le_decode_range(); + test_s16le_nan_input(); + test_s16le_odd_length(); + if (failures) { + fprintf(stderr, "%d check(s) failed\n", failures); + return 1; + } + fprintf(stderr, "all audio_units checks passed\n"); + return 0; +} diff --git a/backend/cpp/audio-cpp/capability_routing.cpp b/backend/cpp/audio-cpp/capability_routing.cpp new file mode 100644 index 000000000..5fc2b63a7 --- /dev/null +++ b/backend/cpp/audio-cpp/capability_routing.cpp @@ -0,0 +1,411 @@ +#include "capability_routing.h" + +#include + +namespace audiocpp_backend { +namespace { + +struct NamedTask { + Task task; + const char *name; +}; + +// Short names are exactly the strings audio.cpp prints and parses in +// framework/runtime/session.cpp, so a name pinned here survives conversion at +// the engine boundary and a name copied out of audio.cpp is accepted here. All +// thirteen have an upstream name; only "spk" is absent from the --task table in +// docs/usage.md. +const NamedTask kTaskNames[] = { + {Task::Vad, "vad"}, + {Task::Asr, "asr"}, + {Task::Diarization, "diar"}, + {Task::SourceSeparation, "sep"}, + {Task::AudioGeneration, "gen"}, + {Task::Tts, "tts"}, + {Task::VoiceCloning, "clon"}, + {Task::VoiceConversion, "vc"}, + {Task::SpeechToSpeech, "s2s"}, + {Task::Alignment, "align"}, + {Task::VoiceDesign, "vdes"}, + {Task::SpeakerRecognition, "spk"}, + {Task::Svc, "svc"}, +}; + +// Accepted on input but never emitted. "spkrec" was this backend's own earlier +// name for the kind; upstream only ever knew "spk". +const NamedTask kTaskAliases[] = { + {Task::SpeakerRecognition, "spkrec"}, +}; + +// First match wins, which is safe because a Capabilities value holds at most one +// entry per task: it mirrors upstream runtime::TaskCapability (model.h), which +// pairs one kind with a modes vector, and no loader's supported_tasks list +// repeats a kind. +bool family_supports(const Capabilities &caps, Task task, Mode mode) { + for (const auto &capability : caps.tasks) { + if (capability.task != task) { + continue; + } + return std::find(capability.modes.begin(), capability.modes.end(), mode) != + capability.modes.end(); + } + return false; +} + +// Mode preference per RPC. Only AudioTranscriptionStream has a fallback: a +// server-streaming transcription can be satisfied by an offline run that emits +// one delta then the final result. Live transcription cannot, because it is +// bidirectional and must consume audio incrementally. +std::vector mode_candidates(Rpc rpc) { + switch (rpc) { + case Rpc::TtsStream: + case Rpc::AudioTranscriptionLive: + return {Mode::Streaming}; + case Rpc::AudioTranscriptionStream: + return {Mode::Streaming, Mode::Offline}; + default: + return {Mode::Offline}; + } +} + +std::vector task_candidates(Rpc rpc, const RequestShape &shape) { + switch (rpc) { + case Rpc::Tts: + case Rpc::TtsStream: + // A supplied speaker clip is the strongest signal: the caller named the + // voice they want. Free-form instructions come next. Both fall back to + // plain Tts so a family without the specialised task still answers. + if (shape.has_voice_reference) { + return {Task::VoiceCloning, Task::Tts, Task::VoiceDesign}; + } + if (shape.has_instructions) { + return {Task::VoiceDesign, Task::Tts, Task::VoiceCloning}; + } + return {Task::Tts, Task::VoiceCloning, Task::VoiceDesign}; + case Rpc::AudioTranscription: + case Rpc::AudioTranscriptionStream: + case Rpc::AudioTranscriptionLive: + // Asr first: `prompt` is also whisper-style decoding context, so its + // presence must not hijack a real ASR family into forced alignment. + if (shape.has_prompt_text) { + return {Task::Asr, Task::Alignment}; + } + return {Task::Asr}; + case Rpc::Vad: + return {Task::Vad}; + case Rpc::Diarize: + return {Task::Diarization}; + case Rpc::SoundGeneration: + return {Task::AudioGeneration}; + case Rpc::AudioTransform: + // Svc is listed for completeness but is unreachable by auto-routing, by + // design: the only families advertising it (seed_vc, vevo2) also + // advertise VoiceConversion, which always wins, and no request signal + // means "this input is singing". Singing voice conversion therefore + // requires an explicit task:svc pin. + return {Task::SourceSeparation, Task::VoiceConversion, Task::Svc, + Task::SpeechToSpeech}; + } + return {}; +} + +// The reasons behind unsupported_surfaces(), spelled once because AudioEncode +// and AudioDecode share theirs. Each is phrased in terms of what upstream does +// and does not have, so a reader can check it against the pinned checkout +// rather than take it on trust. Every one of them was checked against +// audio.cpp e800d435d130dc776baf6f3e6129bb62b1495c89, and one of the four +// claims this backend was planned against did not survive that check: see +// kTransformStreamReason. +// +// A latent upstream inconsistency worth knowing about but deliberately NOT put +// on the wire, because it would mislead: model_spec/schema.cpp's task-string +// whitelist does accept "codec" (and "dialogue"), while +// model_spec/metadata.cpp's parse_task_kind has no branch for either and +// throws "unknown model spec task". So a spec declaring "codec" validates and +// then fails to load. That is a hole in upstream's own validation, not a codec +// task this backend could reach. +const char *const kCodecReason = + "audio.cpp's VoiceTaskKind has no codec entry, so no family can be asked to " + "turn PCM into codec frames or back; miocodec carries a Codec tag in " + "upstream's README but its loader advertises only vc and s2s"; +// NOT "streaming exists for tts and asr only", which is what this backend was +// planned to say and is false: silero_vad advertises vad with RunMode::Streaming +// (src/models/silero_vad/session.cpp). The claim that actually holds is the +// narrower one below, about the four tasks AudioTransform routes to. +// The trailing clause is not padding. The premise is an absence, and an absence +// does not on its own make the RPC impossible: an offline sep family could be +// buffered and emitted as a stream, which is what several LocalAI backends do. +// Stopping at "nothing advertises streaming" would imply an impossibility the +// evidence does not support. What is true, and what the caller needs, is that +// this backend declines to dress an offline call up as a streaming one. +const char *const kTransformStreamReason = + "no audio.cpp family advertises streaming for any task AudioTransform routes " + "to (sep, vc, svc, s2s); upstream advertises RunMode::Streaming for tts, asr " + "and vad only, and no conversion or separation family even implements its " + "IStreamingVoiceTaskSession interface, so a streaming transform here would be " + "a buffered offline call in disguise, which this backend does not pretend to " + "offer"; +// "clip-to-clip processing against a target voice", NOT "voice conversion". The +// latter is true of miocodec and FALSE of vevo2, whose s2s route is `editing` +// and only `editing`: src/models/vevo2/session.cpp's default_route_for_task maps +// SpeechToSpeech to Editing and route_matches_task accepts nothing else, and +// docs/models/vevo2.md defines that route as "Edit source speech into new target +// text while using the target voice", requiring --target-text. It rewrites what +// was said. vevo2's actual voice conversion is its separate vc task, which is +// why upstream's README tags the family "TTS, Music, VC, Edit". The conclusion +// is unaffected: neither family converses. +const char *const kAudioToAudioReason = + "LocalAI's contract here is OpenAI-Realtime shaped, an audio conversation " + "emitting audio, transcript and tool-call deltas from a system prompt and a " + "tool list; audio.cpp's s2s is offline clip-to-clip processing against a " + "target voice, declared only by miocodec (voice conversion) and vevo2 " + "(speech editing), with no conversation, system prompt or tool loop"; +const char *const kVoiceEmbedReason = + "no audio.cpp family advertises the spk (SpeakerRecognition) task, so " + "nothing in the engine can produce a speaker embedding; the task kind " + "itself exists upstream, and TitaNet and ECAPA-TDNN exist as internal " + "conditioning encoders, but neither is registered as a loadable family"; + +// The tasks an RPC is ever willing to route to, independent of request shape. +// +// DERIVED from task_candidates rather than restated, so a task added to an +// RPC's candidate list cannot become inadmissible as a pin by omission. Setting +// every shape flag yields each RPC's widest list: the per-flag branches only +// reorder the same three tasks for Tts, and only ADD Alignment for +// transcription, so the union is what comes back. +std::vector admissible_tasks(Rpc rpc) { + RequestShape widest; + widest.has_voice_reference = true; + widest.has_instructions = true; + widest.has_prompt_text = true; + return task_candidates(rpc, widest); +} + +std::string join_task_names(const std::vector &tasks) { + std::string out; + for (const Task task : tasks) { + if (!out.empty()) { + out += ", "; + } + out += task_name(task); + } + if (out.empty()) { + out = "nothing"; + } + return out; +} + +std::string join_attempts(const std::vector &tasks, + const std::vector &modes) { + std::string out; + for (const Task task : tasks) { + for (const Mode mode : modes) { + if (!out.empty()) { + out += ", "; + } + out += task_name(task); + out += "/"; + out += mode_name(mode); + } + } + return out; +} + +} // namespace + +const char *task_name(Task task) { + for (const auto &entry : kTaskNames) { + if (entry.task == task) { + return entry.name; + } + } + return "unknown"; +} + +const char *mode_name(Mode mode) { + return mode == Mode::Streaming ? "streaming" : "offline"; +} + +const char *rpc_name(Rpc rpc) { + switch (rpc) { + case Rpc::Tts: + return "TTS"; + case Rpc::TtsStream: + return "TTSStream"; + case Rpc::AudioTranscription: + return "AudioTranscription"; + case Rpc::AudioTranscriptionStream: + return "AudioTranscriptionStream"; + case Rpc::AudioTranscriptionLive: + return "AudioTranscriptionLive"; + case Rpc::Vad: + return "VAD"; + case Rpc::Diarize: + return "Diarize"; + case Rpc::SoundGeneration: + return "SoundGeneration"; + case Rpc::AudioTransform: + return "AudioTransform"; + } + return "unknown"; +} + +bool parse_task_name(const std::string &value, Task &out) { + for (const auto &entry : kTaskNames) { + if (value == entry.name) { + out = entry.task; + return true; + } + } + for (const auto &entry : kTaskAliases) { + if (value == entry.name) { + out = entry.task; + return true; + } + } + return false; +} + +std::string describe_capabilities(const Capabilities &caps) { + std::string out; + for (const auto &capability : caps.tasks) { + for (const Mode mode : capability.modes) { + if (!out.empty()) { + out += ", "; + } + out += task_name(capability.task); + out += "/"; + out += mode_name(mode); + } + } + if (out.empty()) { + out = "nothing"; + } + return out; +} + +const std::vector &unsupported_surfaces() { + // Ordered as UnsupportedRpc declares them. unsupported_surface() names each + // index in a switch rather than casting the enum, so the order is checked at + // compile time rather than trusted. + static const std::vector kSurfaces = { + {"AudioEncode", kCodecReason}, + {"AudioDecode", kCodecReason}, + {"AudioTransformStream", kTransformStreamReason}, + {"AudioToAudioStream", kAudioToAudioReason}, + {"VoiceEmbed", kVoiceEmbedReason}, + }; + return kSurfaces; +} + +// A switch with NO default label, deliberately. -Wswitch is on under -Wall, so a +// sixth UnsupportedRpc added without a case here is a BUILD diagnostic, which is +// the only place this class of mistake can be caught for free: a positional +// static_cast(rpc) would compile fine and read past the end of the table +// at run time, on the one code path whose entire job is to be diagnosable. The +// table stays a table because the tests iterate it. +// +// The trailing return is unreachable through the enum and exists only for a +// caller that hands over a value outside it, which is already undefined +// behaviour by the time it arrives. +const UnsupportedSurface &unsupported_surface(UnsupportedRpc rpc) { + const std::vector &surfaces = unsupported_surfaces(); + switch (rpc) { + case UnsupportedRpc::AudioEncode: + return surfaces[0]; + case UnsupportedRpc::AudioDecode: + return surfaces[1]; + case UnsupportedRpc::AudioTransformStream: + return surfaces[2]; + case UnsupportedRpc::AudioToAudioStream: + return surfaces[3]; + case UnsupportedRpc::VoiceEmbed: + return surfaces[4]; + } + return surfaces[0]; +} + +std::string unsupported_surface_message(const Capabilities &caps, const char *rpc, + const char *reason) { + return std::string("audio-cpp: the ") + rpc + + " RPC is not available through this backend because " + reason + + ". Loaded family '" + caps.family + + "' supports: " + describe_capabilities(caps); +} + +std::string unsupported_surface_message(const char *rpc, const char *reason) { + return std::string("audio-cpp: the ") + rpc + + " RPC is not available through this backend because " + reason + + ". No model is loaded, so there is no family to list; loading one " + "would not change this answer"; +} + +Route resolve_route(Rpc rpc, const RequestShape &shape, + const Capabilities &caps) { + Route route; + + std::vector tasks; + if (!shape.pinned_task.empty()) { + Task pinned = Task::Tts; + if (!parse_task_name(shape.pinned_task, pinned)) { + route.error = "audio-cpp: unknown task option '" + shape.pinned_task + + "'. Known tasks: gen, tts, clon, vc, svc, s2s, asr, " + "align, vad, diar, sep, vdes, spk"; + return route; + } + // A pin is honoured exactly, but ONLY on an RPC that could have routed + // to it anyway. It used to replace the candidate list wholesale for + // every RPC, and because the model's `task:` option is copied into the + // shape by all nine handlers, one pin bled across all nine surfaces and + // produced wrong 200s rather than errors: nemotron with task:asr made + // Vad return 200 with zero segments after a full ASR decode, so 14 + // seconds of speech was reported as silence, and silero_vad with + // task:vad made AudioTranscription return 200 with empty text and four + // segments whose spans were VAD segments, which the srt/vtt/lrc writers + // then rendered as a well formed subtitle file of four timed EMPTY + // cues. Refusing is what the docs already promise: "if the family + // cannot serve it, the request is refused rather than rerouted". + // + // Every legitimate pin survives, because a pin only ever names the task + // its own RPC already routes to: svc is in AudioTransform's candidates, + // tts/clon/vdes in TTS's, asr in transcription's, vad and diar in + // theirs. + const std::vector admissible = admissible_tasks(rpc); + if (std::find(admissible.begin(), admissible.end(), pinned) == + admissible.end()) { + route.error = std::string("audio-cpp: this model pins task '") + + task_name(pinned) + "', which the " + rpc_name(rpc) + + " RPC never routes to (it routes to " + + join_task_names(admissible) + + "). Remove the task option to reach this RPC, or call " + "the RPC the pinned task serves"; + return route; + } + tasks = {pinned}; + } else { + tasks = task_candidates(rpc, shape); + } + + const std::vector modes = mode_candidates(rpc); + + // Task-major: prefer the right task in a fallback mode over the wrong task + // in the preferred mode. + for (const Task task : tasks) { + for (const Mode mode : modes) { + if (family_supports(caps, task, mode)) { + route.ok = true; + route.task = task; + route.mode = mode; + return route; + } + } + } + + route.error = std::string("audio-cpp: family '") + caps.family + + "' cannot serve the " + rpc_name(rpc) + " RPC (tried " + + join_attempts(tasks, modes) + "); it supports: " + + describe_capabilities(caps); + return route; +} + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/capability_routing.h b/backend/cpp/audio-cpp/capability_routing.h new file mode 100644 index 000000000..9ba9c9aca --- /dev/null +++ b/backend/cpp/audio-cpp/capability_routing.h @@ -0,0 +1,131 @@ +#pragma once + +// Decides which audio.cpp (task, mode) pair serves a given LocalAI RPC, or +// produces the capability error when none can. Standard library only, so this +// unit is tested without an audio.cpp checkout; loaded_model.cpp converts +// to and from engine::runtime types at the boundary. + +#include +#include + +namespace audiocpp_backend { + +// Mirrors engine::runtime::VoiceTaskKind, same members and same order. +enum class Task { + Vad, + Asr, + Diarization, + SourceSeparation, + AudioGeneration, + Tts, + VoiceCloning, + VoiceConversion, + SpeechToSpeech, + Alignment, + VoiceDesign, + SpeakerRecognition, + Svc, +}; + +// Mirrors engine::runtime::RunMode. +enum class Mode { Offline, Streaming }; + +struct TaskCapability { + Task task = Task::Vad; + std::vector modes; +}; + +struct Capabilities { + std::string family; + std::vector tasks; +}; + +// The LocalAI RPCs this backend serves. The ones it cannot serve at all are in +// UnsupportedRpc below rather than here: they never reach routing, because no +// family could satisfy them. +enum class Rpc { + Tts, + TtsStream, + AudioTranscription, + AudioTranscriptionStream, + AudioTranscriptionLive, + Vad, + Diarize, + SoundGeneration, + AudioTransform, +}; + +struct RequestShape { + // A speaker reference clip was supplied (TTSRequest.voice resolved to audio). + bool has_voice_reference = false; + // TTSRequest.instructions is set. + bool has_instructions = false; + // TranscriptRequest.prompt is set. + bool has_prompt_text = false; + // The model's `task:` option, empty when unset. Overrides routing. + std::string pinned_task; +}; + +struct Route { + bool ok = false; + Task task = Task::Tts; + Mode mode = Mode::Offline; + // Set when ok is false. Suitable verbatim as an UNIMPLEMENTED message. + std::string error; +}; + +Route resolve_route(Rpc rpc, const RequestShape &shape, const Capabilities &caps); + +// Canonical audio.cpp short names: gen, tts, clon, vc, svc, s2s, asr, align, +// vad, diar, sep, vdes, spk. parse_task_name additionally accepts "spkrec" as +// a legacy alias; task_name only ever emits "spk". +const char *task_name(Task task); +const char *mode_name(Mode mode); +const char *rpc_name(Rpc rpc); +bool parse_task_name(const std::string &value, Task &out); + +// "asr/offline, asr/streaming", for error messages. +std::string describe_capabilities(const Capabilities &caps); + +// The RPCs in LocalAI's backend contract that audio.cpp has no counterpart for, +// as opposed to the ones in Rpc above, which a particular family may or may not +// be able to serve. Nothing routes to these: the refusal is a property of the +// engine, not of the loaded model, so loading a different family cannot change +// it. +// +// This is deliberately NOT deferred work. Each entry names the upstream +// limitation that keeps it out of Rpc, and each becomes an ordinary routing +// entry the day upstream lifts that limitation. +enum class UnsupportedRpc { + AudioEncode, + AudioDecode, + AudioTransformStream, + AudioToAudioStream, + VoiceEmbed, +}; + +struct UnsupportedSurface { + // The RPC's name as backend.proto spells it, for the message. + const char *rpc; + // Why audio.cpp cannot serve it, in terms of what upstream does and does + // not have. Stated so a caller can tell "not built yet" from "not possible". + const char *reason; +}; + +// The table behind the five refusals. Exposed whole so a test can assert every +// entry rather than the one somebody remembered to cover, and so the reasons +// are data in one place instead of string literals hand-copied into handlers. +const std::vector &unsupported_surfaces(); +const UnsupportedSurface &unsupported_surface(UnsupportedRpc rpc); + +// Message for an RPC this backend cannot serve at all, as opposed to one this +// particular family cannot serve. `reason` states the upstream limitation. +std::string unsupported_surface_message(const Capabilities &caps, const char *rpc, + const char *reason); + +// Same, for the no-model-loaded case. Says so explicitly rather than naming an +// empty family, and says that loading one would not help, because the caller's +// obvious next move otherwise is to load a model and try again. +std::string unsupported_surface_message(const char *rpc, const char *reason); + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/capability_routing_test.cpp b/backend/cpp/audio-cpp/capability_routing_test.cpp new file mode 100644 index 000000000..b4f44e289 --- /dev/null +++ b/backend/cpp/audio-cpp/capability_routing_test.cpp @@ -0,0 +1,600 @@ +// Unit tests for capability_routing. Standard library only. The harness +// compiles this as a single translation unit, so the implementation is +// included directly rather than linked. + +#include "capability_routing.cpp" + +#include +#include +#include + +static int failures = 0; + +static void check(bool ok, const std::string &name) { + if (!ok) { + failures++; + fprintf(stderr, "FAIL: %s\n", name.c_str()); + } else { + fprintf(stderr, "ok: %s\n", name.c_str()); + } +} + +using namespace audiocpp_backend; + +// Mirrors what supertonic advertises: TTS offline and streaming. +static Capabilities supertonic() { + return Capabilities{"supertonic", + {{Task::Tts, {Mode::Offline, Mode::Streaming}}}}; +} + +// Mirrors chatterbox: TTS, cloning and voice conversion, offline only. +static Capabilities chatterbox() { + return Capabilities{"chatterbox", + {{Task::Tts, {Mode::Offline}}, + {Task::VoiceCloning, {Mode::Offline}}, + {Task::VoiceConversion, {Mode::Offline}}}}; +} + +// Mirrors nemotron_asr: ASR offline and streaming. +static Capabilities nemotron() { + return Capabilities{"nemotron_asr", + {{Task::Asr, {Mode::Offline, Mode::Streaming}}}}; +} + +// Mirrors qwen3_asr: ASR offline only. +static Capabilities qwen3_asr() { + return Capabilities{"qwen3_asr", {{Task::Asr, {Mode::Offline}}}}; +} + +// Mirrors qwen3_forced_aligner: alignment only. +static Capabilities aligner() { + return Capabilities{"qwen3_forced_aligner", + {{Task::Alignment, {Mode::Offline}}}}; +} + +// Mirrors htdemucs: separation only. +static Capabilities htdemucs() { + return Capabilities{"htdemucs", {{Task::SourceSeparation, {Mode::Offline}}}}; +} + +static void test_plain_tts() { + const auto r = resolve_route(Rpc::Tts, RequestShape{}, chatterbox()); + check(r.ok, "plain TTS routes"); + check(r.task == Task::Tts, "plain TTS picks Tts, not VoiceCloning"); + check(r.mode == Mode::Offline, "TTS runs offline"); +} + +static void test_tts_with_voice_reference_prefers_cloning() { + RequestShape shape; + shape.has_voice_reference = true; + const auto r = resolve_route(Rpc::Tts, shape, chatterbox()); + check(r.ok, "TTS with a voice reference routes"); + check(r.task == Task::VoiceCloning, "voice reference prefers VoiceCloning"); +} + +// supertonic has no VoiceCloning: a voice reference must fall back to Tts +// rather than failing the request. +static void test_tts_voice_reference_falls_back_to_tts() { + RequestShape shape; + shape.has_voice_reference = true; + const auto r = resolve_route(Rpc::Tts, shape, supertonic()); + check(r.ok, "voice reference on a clone-less family still routes"); + check(r.task == Task::Tts, "falls back to Tts"); +} + +static void test_tts_instructions_prefer_voice_design() { + RequestShape shape; + shape.has_instructions = true; + Capabilities caps{"qwen3_tts", + {{Task::Tts, {Mode::Offline}}, + {Task::VoiceDesign, {Mode::Offline}}}}; + const auto r = resolve_route(Rpc::Tts, shape, caps); + check(r.ok, "TTS with instructions routes"); + check(r.task == Task::VoiceDesign, "instructions prefer VoiceDesign"); +} + +// A voice reference is a stronger signal than free-form instructions: cloning +// a specific voice is what the user asked for. +static void test_voice_reference_beats_instructions() { + RequestShape shape; + shape.has_voice_reference = true; + shape.has_instructions = true; + Capabilities caps{"omnivoice", + {{Task::Tts, {Mode::Offline}}, + {Task::VoiceCloning, {Mode::Offline}}, + {Task::VoiceDesign, {Mode::Offline}}}}; + const auto r = resolve_route(Rpc::Tts, shape, caps); + check(r.ok, "both signals present routes"); + check(r.task == Task::VoiceCloning, "voice reference outranks instructions"); +} + +static void test_tts_stream_requires_streaming() { + const auto ok = resolve_route(Rpc::TtsStream, RequestShape{}, supertonic()); + check(ok.ok, "streaming TTS routes on supertonic"); + check(ok.mode == Mode::Streaming, "TTSStream runs in streaming mode"); + + const auto bad = resolve_route(Rpc::TtsStream, RequestShape{}, chatterbox()); + check(!bad.ok, "streaming TTS is refused on an offline-only family"); + check(bad.error.find("chatterbox") != std::string::npos, + "error names the family"); + check(bad.error.find("tts/offline") != std::string::npos, + "error lists what the family does support"); + check(bad.error.find("TTSStream") != std::string::npos, + "error names the RPC that was refused"); + check(bad.error.find("tts/streaming") != std::string::npos, + "error lists the (task, mode) pairs that were tried"); +} + +static void test_transcription_stream_falls_back_to_offline() { + const auto streaming = + resolve_route(Rpc::AudioTranscriptionStream, RequestShape{}, nemotron()); + check(streaming.ok && streaming.mode == Mode::Streaming, + "streaming ASR uses streaming mode when offered"); + + const auto offline = + resolve_route(Rpc::AudioTranscriptionStream, RequestShape{}, qwen3_asr()); + check(offline.ok, "streaming ASR falls back on an offline-only family"); + check(offline.mode == Mode::Offline, "fallback mode is offline"); + check(offline.task == Task::Asr, "fallback task is still Asr"); +} + +// Task preference dominates mode preference: it is better to run the right +// task in a fallback mode than the wrong task in the preferred mode. This is +// the only RPC where the two orderings can disagree, because it is the only +// one with more than one acceptable mode. +static void test_task_preference_beats_mode_preference() { + RequestShape shape; + shape.has_prompt_text = true; + Capabilities mixed{"mixed_asr_aligner", + {{Task::Asr, {Mode::Offline}}, + {Task::Alignment, {Mode::Streaming}}}}; + const auto r = resolve_route(Rpc::AudioTranscriptionStream, shape, mixed); + check(r.ok, "mixed family routes"); + check(r.task == Task::Asr, + "the preferred task wins even in its fallback mode"); + check(r.mode == Mode::Offline, + "the fallback mode is accepted to keep the preferred task"); +} + +// Live transcription is bidirectional and cannot be faked from an offline run. +static void test_live_transcription_has_no_offline_fallback() { + const auto r = + resolve_route(Rpc::AudioTranscriptionLive, RequestShape{}, qwen3_asr()); + check(!r.ok, "live transcription is refused on an offline-only family"); +} + +static void test_alignment_needs_prompt_text() { + const auto without = + resolve_route(Rpc::AudioTranscription, RequestShape{}, aligner()); + check(!without.ok, "aligner without a transcript is refused"); + + RequestShape shape; + shape.has_prompt_text = true; + const auto with = resolve_route(Rpc::AudioTranscription, shape, aligner()); + check(with.ok, "aligner with a transcript routes"); + check(with.task == Task::Alignment, "routes to Alignment"); +} + +// A real ASR family must not be hijacked to Alignment just because the caller +// passed a prompt: `prompt` is also whisper-style decoding context. +static void test_prompt_does_not_hijack_asr() { + RequestShape shape; + shape.has_prompt_text = true; + const auto r = resolve_route(Rpc::AudioTranscription, shape, nemotron()); + check(r.ok, "ASR with a prompt routes"); + + // nemotron advertises Asr alone, so the assertion has to be made against a + // family that advertises both: otherwise "Asr is preferred" only restates + // that Asr is the only option, and reversing the preference order passes. + Capabilities both{"asr_with_aligner", + {{Task::Asr, {Mode::Offline}}, + {Task::Alignment, {Mode::Offline}}}}; + const auto pref = resolve_route(Rpc::AudioTranscription, shape, both); + check(pref.ok, "a family offering both routes"); + check(pref.task == Task::Asr, "Asr is preferred over Alignment"); +} + +static void test_audio_transform_prefers_separation() { + const auto sep = + resolve_route(Rpc::AudioTransform, RequestShape{}, htdemucs()); + check(sep.ok && sep.task == Task::SourceSeparation, "separation routes"); + + // htdemucs advertises separation alone, so the check above cannot fail on + // ordering. This family advertises both, which is what pins the preference. + Capabilities sep_and_vc{"sep_and_vc", + {{Task::SourceSeparation, {Mode::Offline}}, + {Task::VoiceConversion, {Mode::Offline}}}}; + const auto pref = + resolve_route(Rpc::AudioTransform, RequestShape{}, sep_and_vc); + check(pref.ok && pref.task == Task::SourceSeparation, + "separation is preferred over voice conversion"); + + Capabilities miocodec{"miocodec", + {{Task::VoiceConversion, {Mode::Offline}}, + {Task::SpeechToSpeech, {Mode::Offline}}}}; + const auto vc = resolve_route(Rpc::AudioTransform, RequestShape{}, miocodec); + check(vc.ok && vc.task == Task::VoiceConversion, + "voice conversion is preferred over speech-to-speech"); +} + +static void test_pinned_task_overrides_routing() { + RequestShape shape; + shape.pinned_task = "s2s"; + Capabilities miocodec{"miocodec", + {{Task::VoiceConversion, {Mode::Offline}}, + {Task::SpeechToSpeech, {Mode::Offline}}}}; + const auto r = resolve_route(Rpc::AudioTransform, shape, miocodec); + check(r.ok && r.task == Task::SpeechToSpeech, "pinned task wins"); + + RequestShape bad; + bad.pinned_task = "not-a-task"; + const auto e = resolve_route(Rpc::AudioTransform, bad, miocodec); + check(!e.ok, "an unknown pinned task is an error"); + check(e.error.find("not-a-task") != std::string::npos, + "error names the bad task"); + + // A pinned task the family does not offer must fail, not silently reroute. + RequestShape unsupported; + unsupported.pinned_task = "sep"; + const auto u = resolve_route(Rpc::AudioTransform, unsupported, miocodec); + check(!u.ok, "a pinned but unsupported task is refused"); +} + +// A pin lives on the MODEL, and every one of the nine handlers copies it into +// the shape, so a pin set for one RPC arrives at all of them. It used to +// replace the candidate list wholesale, which turned the other eight into wrong +// 200s rather than errors: nemotron pinned to asr made Vad answer with zero +// segments after a full ASR decode, and silero_vad pinned to vad made +// AudioTranscription answer with empty text and four segments whose spans were +// VAD segments, which the srt/vtt/lrc writers rendered as timed EMPTY cues. +static void test_pin_must_be_admissible_for_the_rpc() { + Capabilities nemotron_asr{"nemotron_asr", + {{Task::Asr, {Mode::Offline, Mode::Streaming}}, + {Task::Vad, {Mode::Offline}}}}; + + // The pin is legitimate on the RPC it was meant for. + RequestShape asr_pin; + asr_pin.pinned_task = "asr"; + const auto transcription = + resolve_route(Rpc::AudioTranscription, asr_pin, nemotron_asr); + check(transcription.ok && transcription.task == Task::Asr, + "an admissible pin is still honoured exactly"); + + // ...and refused on one that never routes to it, EVEN THOUGH the family + // advertises the pinned task. That is the whole point: family support is + // not the question, RPC admissibility is. + const auto vad = resolve_route(Rpc::Vad, asr_pin, nemotron_asr); + check(!vad.ok, "an inadmissible pin is refused rather than served"); + check(vad.error.find("asr") != std::string::npos, + "the refusal names the pinned task"); + check(vad.error.find(rpc_name(Rpc::Vad)) != std::string::npos, + "the refusal names the RPC that cannot serve it"); + + Capabilities silero{"silero_vad", {{Task::Vad, {Mode::Offline}}}}; + RequestShape vad_pin; + vad_pin.pinned_task = "vad"; + const auto vad_ok = resolve_route(Rpc::Vad, vad_pin, silero); + check(vad_ok.ok && vad_ok.task == Task::Vad, "vad is admissible on Vad"); + const auto transcribe_vad = + resolve_route(Rpc::AudioTranscription, vad_pin, silero); + check(!transcribe_vad.ok, + "a vad pin cannot make a transcription request return empty cues"); + + // Every pin a shipped configuration could sensibly set stays reachable on + // the RPC that serves it. This is the list the fix was checked against. + struct AdmissibleCase { + Rpc rpc; + const char *task; + }; + const AdmissibleCase kAdmissible[] = { + {Rpc::AudioTransform, "svc"}, {Rpc::AudioTransform, "sep"}, + {Rpc::AudioTransform, "vc"}, {Rpc::AudioTransform, "s2s"}, + {Rpc::Tts, "tts"}, {Rpc::Tts, "clon"}, + {Rpc::Tts, "vdes"}, {Rpc::TtsStream, "tts"}, + {Rpc::AudioTranscription, "asr"}, + {Rpc::AudioTranscription, "align"}, + {Rpc::AudioTranscriptionStream, "asr"}, + {Rpc::AudioTranscriptionLive, "asr"}, + {Rpc::Vad, "vad"}, {Rpc::Diarize, "diar"}, + {Rpc::SoundGeneration, "gen"}, + }; + for (const auto &entry : kAdmissible) { + Task task = Task::Tts; + check(parse_task_name(entry.task, task), + std::string("known task name: ") + entry.task); + // A family that advertises the pinned task offline and nothing else, so + // the ONLY thing that can refuse the route is the admissibility check. + Capabilities only{"probe", + {{task, {Mode::Offline, Mode::Streaming}}}}; + RequestShape pin; + pin.pinned_task = entry.task; + const auto route = resolve_route(entry.rpc, pin, only); + check(route.ok && route.task == task, + std::string("pin '") + entry.task + "' stays admissible on " + + rpc_name(entry.rpc)); + } + + // And the pins that must NOT cross over, one per RPC pair that was + // observed producing a wrong 200. + const AdmissibleCase kInadmissible[] = { + {Rpc::Vad, "asr"}, {Rpc::Diarize, "asr"}, + {Rpc::AudioTranscription, "vad"}, {Rpc::AudioTranscription, "diar"}, + {Rpc::Tts, "asr"}, {Rpc::Vad, "tts"}, + {Rpc::SoundGeneration, "tts"}, {Rpc::AudioTransform, "asr"}, + }; + for (const auto &entry : kInadmissible) { + Task task = Task::Tts; + check(parse_task_name(entry.task, task), + std::string("known task name: ") + entry.task); + Capabilities only{"probe", + {{task, {Mode::Offline, Mode::Streaming}}}}; + RequestShape pin; + pin.pinned_task = entry.task; + const auto route = resolve_route(entry.rpc, pin, only); + check(!route.ok, + std::string("pin '") + entry.task + "' is refused on " + + rpc_name(entry.rpc)); + } +} + +static void test_vad_and_diarize() { + Capabilities silero{"silero_vad", {{Task::Vad, {Mode::Offline, Mode::Streaming}}}}; + const auto v = resolve_route(Rpc::Vad, RequestShape{}, silero); + check(v.ok && v.task == Task::Vad && v.mode == Mode::Offline, "VAD routes offline"); + + const auto d = resolve_route(Rpc::Diarize, RequestShape{}, silero); + check(!d.ok, "diarization is refused on a VAD-only family"); + + Capabilities sortformer{"sortformer_diar", {{Task::Diarization, {Mode::Offline}}}}; + const auto ok = resolve_route(Rpc::Diarize, RequestShape{}, sortformer); + check(ok.ok && ok.task == Task::Diarization, "diarization routes"); +} + +static void test_sound_generation() { + Capabilities stable{"stable_audio", {{Task::AudioGeneration, {Mode::Offline}}}}; + const auto r = resolve_route(Rpc::SoundGeneration, RequestShape{}, stable); + check(r.ok && r.task == Task::AudioGeneration, "sound generation routes"); +} + +static void test_names_round_trip() { + const Task all[] = {Task::Vad, Task::Asr, Task::Diarization, + Task::SourceSeparation, Task::AudioGeneration, Task::Tts, + Task::VoiceCloning, Task::VoiceConversion, + Task::SpeechToSpeech, Task::Alignment, Task::VoiceDesign, + Task::SpeakerRecognition, Task::Svc}; + for (const Task t : all) { + Task parsed = Task::Vad; + const bool ok = parse_task_name(task_name(t), parsed); + check(ok && parsed == t, + std::string("task name round-trips: ") + task_name(t)); + } + check(std::string(mode_name(Mode::Offline)) == "offline", "offline name"); + check(std::string(mode_name(Mode::Streaming)) == "streaming", "streaming name"); + + // The emitted name must be the one audio.cpp itself prints and parses + // (framework/runtime/session.cpp), because `task:` is user-facing: a name + // copied out of audio.cpp has to be accepted here, and a name pinned here + // has to survive conversion at the engine boundary. + check(std::string(task_name(Task::SpeakerRecognition)) == "spk", + "speaker recognition emits upstream's name 'spk'"); + + Task pinned = Task::Vad; + check(parse_task_name("spk", pinned) && pinned == Task::SpeakerRecognition, + "'spk' parses to SpeakerRecognition"); + + // Accepted as a legacy alias so configs written against the earlier name + // keep working, but never emitted. + Task alias = Task::Vad; + check(parse_task_name("spkrec", alias) && alias == Task::SpeakerRecognition, + "'spkrec' is still accepted as an alias"); +} + +static void test_describe_capabilities() { + const std::string described = describe_capabilities(nemotron()); + check(described.find("asr/offline") != std::string::npos, + "description lists asr/offline"); + check(described.find("asr/streaming") != std::string::npos, + "description lists asr/streaming"); +} + +static void test_empty_capabilities() { + const auto r = resolve_route(Rpc::Tts, RequestShape{}, Capabilities{"mystery", {}}); + check(!r.ok, "a family advertising nothing is refused"); + check(r.error.find("mystery") != std::string::npos, "error names the family"); +} + +// The table is indexed by UnsupportedRpc's underlying value, so a reordering of +// either list silently pairs an RPC with another's reason. Nothing else would +// catch that: both sides still compile and every message still reads plausibly. +static void test_unsupported_surface_table_matches_the_enum() { + check(unsupported_surfaces().size() == 5, + "all five unsupported surfaces are tabulated"); + check(std::string(unsupported_surface(UnsupportedRpc::AudioEncode).rpc) == + "AudioEncode", + "UnsupportedRpc::AudioEncode indexes AudioEncode"); + check(std::string(unsupported_surface(UnsupportedRpc::AudioDecode).rpc) == + "AudioDecode", + "UnsupportedRpc::AudioDecode indexes AudioDecode"); + check(std::string( + unsupported_surface(UnsupportedRpc::AudioTransformStream).rpc) == + "AudioTransformStream", + "UnsupportedRpc::AudioTransformStream indexes AudioTransformStream"); + check(std::string( + unsupported_surface(UnsupportedRpc::AudioToAudioStream).rpc) == + "AudioToAudioStream", + "UnsupportedRpc::AudioToAudioStream indexes AudioToAudioStream"); + check(std::string(unsupported_surface(UnsupportedRpc::VoiceEmbed).rpc) == + "VoiceEmbed", + "UnsupportedRpc::VoiceEmbed indexes VoiceEmbed"); + + // Two entries may share a reason (the codec pair does), but two entries + // naming the same RPC would mean one of the five is unreachable. + for (size_t i = 0; i < unsupported_surfaces().size(); ++i) { + for (size_t j = i + 1; j < unsupported_surfaces().size(); ++j) { + check(std::string(unsupported_surfaces()[i].rpc) != + unsupported_surfaces()[j].rpc, + std::string("no duplicate RPC name at ") + std::to_string(i) + + "/" + std::to_string(j)); + } + } +} + +// There is no out-of-range test for unsupported_surface(). It switches over the +// enumerators with no default label, so a sixth UnsupportedRpc without a case is +// a -Wswitch diagnostic at build time and cannot reach a run-time check at all. + +// Every entry, not just the one somebody remembered to cover. A refusal that +// drops the family, the RPC or the reason is a refusal the caller cannot act +// on, which is the entire point of this surface existing. +static void test_every_unsupported_surface_message_is_diagnosable() { + for (const auto &surface : unsupported_surfaces()) { + const std::string label = std::string(" [") + surface.rpc + "]"; + const std::string loaded = + unsupported_surface_message(nemotron(), surface.rpc, surface.reason); + + check(loaded.find(surface.rpc) != std::string::npos, + "message names the RPC" + label); + check(std::string(surface.reason).size() > 20 && + loaded.find(surface.reason) != std::string::npos, + "message gives a substantive upstream reason" + label); + check(loaded.find("nemotron_asr") != std::string::npos, + "message names the loaded family" + label); + check(loaded.find("asr/offline") != std::string::npos && + loaded.find("asr/streaming") != std::string::npos, + "message lists what the family does support" + label); + + // The no-model form keeps the two facts that do not depend on a model + // and drops only the one that does, so the caller still learns why. + const std::string unloaded = + unsupported_surface_message(surface.rpc, surface.reason); + check(unloaded.find(surface.rpc) != std::string::npos, + "no-model message names the RPC" + label); + check(unloaded.find(surface.reason) != std::string::npos, + "no-model message gives the upstream reason" + label); + check(unloaded.find("nemotron_asr") == std::string::npos, + "no-model message names no family" + label); + // Without this the caller's obvious next move is to load a model and + // retry, which cannot work: the refusal is a property of the engine. + check(unloaded.find("would not change this answer") != std::string::npos, + "no-model message says loading a model would not help" + label); + } +} + +// The reasons are the load-bearing half of this feature and each was checked +// against the pinned upstream checkout. Pinning the distinguishing phrase here +// means a later edit that guts one into a generic "not supported" fails rather +// than passes quietly. +static void test_unsupported_reasons_name_the_upstream_limitation() { + const auto &encode = unsupported_surface(UnsupportedRpc::AudioEncode); + const auto &decode = unsupported_surface(UnsupportedRpc::AudioDecode); + check(std::string(encode.reason).find("VoiceTaskKind") != std::string::npos && + std::string(encode.reason).find("codec") != std::string::npos, + "the AudioEncode reason names the missing VoiceTaskKind entry"); + // miocodec is the family a reader will reach for first, because upstream's + // README tags it Codec. Naming it and its actual advertised tasks is what + // stops the next person re-deriving the same dead end. + check(std::string(encode.reason).find("miocodec") != std::string::npos, + "the AudioEncode reason disposes of miocodec's README Codec tag"); + check(std::string(encode.reason) == decode.reason, + "AudioEncode and AudioDecode refuse for the same reason"); + + const auto &transform = + unsupported_surface(UnsupportedRpc::AudioTransformStream); + // The reason must be scoped to the tasks AudioTransform routes to. The + // broader claim, "upstream streams tts and asr only", is FALSE: silero_vad + // advertises vad with RunMode::Streaming. A refusal resting on a false + // premise is worse than a bare UNIMPLEMENTED, because it will be believed. + check(std::string(transform.reason).find("sep, vc, svc, s2s") != + std::string::npos, + "the AudioTransformStream reason is scoped to the routed tasks"); + check(std::string(transform.reason).find("tts, asr and vad") != + std::string::npos, + "the AudioTransformStream reason counts vad among the streaming tasks"); + check(std::string(transform.reason).find("tts and asr only") == + std::string::npos, + "the AudioTransformStream reason does not repeat the refuted claim"); + // An absence is not an impossibility. A sep family could be buffered and + // emitted as a stream, so the reason has to say this backend declines to + // rather than cannot, or it overreaches on a true premise. + check(std::string(transform.reason).find("buffered offline call in disguise") != + std::string::npos, + "the AudioTransformStream reason does not overclaim impossibility"); + + const auto &s2s = unsupported_surface(UnsupportedRpc::AudioToAudioStream); + check(std::string(s2s.reason).find("Realtime") != std::string::npos && + std::string(s2s.reason).find("clip-to-clip") != std::string::npos, + "the AudioToAudioStream reason contrasts the two contracts"); + // Naming both s2s families, and what each of them actually does, makes the + // claim checkable. It must NOT say s2s is voice conversion full stop: that + // is true of miocodec and false of vevo2, whose s2s route is `editing` and + // rewrites the spoken content against a target voice. + check(std::string(s2s.reason).find("miocodec (voice conversion)") != + std::string::npos && + std::string(s2s.reason).find("vevo2 (speech editing)") != + std::string::npos, + "the AudioToAudioStream reason names each s2s family's actual task"); + check(std::string(s2s.reason).find("s2s is offline voice conversion") == + std::string::npos, + "the AudioToAudioStream reason does not miscast vevo2 as conversion"); + + const auto &embed = unsupported_surface(UnsupportedRpc::VoiceEmbed); + check(std::string(embed.reason).find("spk") != std::string::npos, + "the VoiceEmbed reason names the task no family advertises"); + // spk IS a VoiceTaskKind upstream; what is missing is any family that + // advertises it. Saying the kind does not exist would be false, and would + // send a reader looking in the wrong place. + check(std::string(embed.reason).find("no audio.cpp family") != + std::string::npos, + "the VoiceEmbed reason blames the families, not the enum"); + // The speaker encoders DO exist upstream, as conditioning modules inside + // TTS and VC families. Not saying so invites "but audio.cpp ships TitaNet". + check(std::string(embed.reason).find("TitaNet") != std::string::npos, + "the VoiceEmbed reason disposes of the internal speaker encoders"); + Task parsed = Task::Vad; + check(parse_task_name("spk", parsed) && parsed == Task::SpeakerRecognition, + "spk is a real task kind, so the reason must not claim otherwise"); +} + +// A family advertising nothing still gets a message that reads, rather than one +// trailing off after "supports: ". +static void test_unsupported_surface_message_with_empty_capabilities() { + const auto &embed = unsupported_surface(UnsupportedRpc::VoiceEmbed); + const std::string message = unsupported_surface_message( + Capabilities{"mystery", {}}, embed.rpc, embed.reason); + check(message.find("mystery") != std::string::npos, + "empty-capability message still names the family"); + check(message.find("supports: nothing") != std::string::npos, + "empty-capability message says the family supports nothing"); +} + +int main() { + test_plain_tts(); + test_tts_with_voice_reference_prefers_cloning(); + test_tts_voice_reference_falls_back_to_tts(); + test_tts_instructions_prefer_voice_design(); + test_voice_reference_beats_instructions(); + test_tts_stream_requires_streaming(); + test_transcription_stream_falls_back_to_offline(); + test_task_preference_beats_mode_preference(); + test_live_transcription_has_no_offline_fallback(); + test_alignment_needs_prompt_text(); + test_prompt_does_not_hijack_asr(); + test_audio_transform_prefers_separation(); + test_pinned_task_overrides_routing(); + test_pin_must_be_admissible_for_the_rpc(); + test_vad_and_diarize(); + test_sound_generation(); + test_names_round_trip(); + test_describe_capabilities(); + test_empty_capabilities(); + test_unsupported_surface_table_matches_the_enum(); + test_every_unsupported_surface_message_is_diagnosable(); + test_unsupported_reasons_name_the_upstream_limitation(); + test_unsupported_surface_message_with_empty_capabilities(); + if (failures) { + fprintf(stderr, "%d check(s) failed\n", failures); + return 1; + } + fprintf(stderr, "all capability_routing checks passed\n"); + return 0; +} diff --git a/backend/cpp/audio-cpp/family_gate.cpp b/backend/cpp/audio-cpp/family_gate.cpp new file mode 100644 index 000000000..2515c7abb --- /dev/null +++ b/backend/cpp/audio-cpp/family_gate.cpp @@ -0,0 +1,163 @@ +#include "family_gate.h" + +#include +#include + +namespace audiocpp_backend { +namespace { + +// PRECONDITION: `suffix` must already be lowercase. Both sides are folded, so +// this reads as symmetric, but only `value` can carry case in practice and a +// caller passing ".GGUF" would still work today for that reason alone. Do not +// rely on it: the fold on the suffix side is the only thing standing between +// this and a helper that answers false for every input, and it is not covered +// by any test, because with a lowercase suffix no input can distinguish it. +bool ends_with_ci(const std::string &value, const std::string &suffix) { + if (value.size() <= suffix.size()) { + return false; // a bare ".gguf" is an extension, not a model file + } + const size_t offset = value.size() - suffix.size(); + for (size_t i = 0; i < suffix.size(); ++i) { + const auto lhs = static_cast(value[offset + i]); + const auto rhs = static_cast(suffix[i]); + if (std::tolower(lhs) != std::tolower(rhs)) { + return false; + } + } + return true; +} + +} // namespace + +bool path_looks_like_gguf(const std::string &path) { + return ends_with_ci(path, ".gguf"); +} + +FamilyDecision decide_family(bool path_is_gguf, const std::string &embedded_family, + const std::string &configured_family) { + FamilyDecision decision; + + if (!configured_family.empty()) { + decision.ok = true; + decision.family = configured_family; + return decision; + } + + if (path_is_gguf) { + if (!embedded_family.empty()) { + decision.ok = true; + decision.family = embedded_family; + return decision; + } + decision.error = + "audio-cpp: this GGUF carries no 'audiocpp.model_spec.family' " + "metadata key, so it is not an audio.cpp model. Convert it with " + "audiocpp_gguf, or name the family explicitly with the model option " + "'family:'"; + return decision; + } + + decision.error = + "audio-cpp: a model path that is not a standalone audio.cpp GGUF needs " + "an explicit 'family:' model option, because the audio.cpp family " + "cannot be inferred from a safetensors or package directory"; + return decision; +} + +namespace { + +// Families that ABORT THE PROCESS on a weight dtype they cannot handle, and the +// dtypes they can. See the header for why this is a list of crashes rather than +// a list of preferences. +// +// supertonic: upstream's docs/gguf.md:90 records its 16-bit GGUF column as +// "---", i.e. NOT TESTED, and its q8_0 as "No (unsupported weight dtype)". Only +// the `orig` package is marked Pass, and its 698 weight tensors are f32 while +// its 72 index and shape constants are i64. The f16 abort is a LOCAL +// OBSERVATION rather than an upstream claim, and it is attributed rather than +// assumed: it is identical through the unary TTS RPC and through TTSStream, so +// it is the packaging and not the streaming path. q8_0 was never run here and is +// refused on upstream's "unsupported weight dtype" alone, which is the weaker of +// the two claims. See the header for why keeping them apart matters. +// +// TO REMOVE AN ENTRY: bump AUDIO_CPP_VERSION past a fix, load a package in the +// refused dtype, and synthesise. If audio comes out, delete the entry. No test +// can do that for you, which is exactly why it is written here: the test beside +// this file pins WHAT the table says, not whether upstream has moved on. Do not +// widen an entry without running that, because what it prevents is a process +// death rather than a wrong answer. +struct DtypeAllowList { + // NULL TERMINATED, and the terminator occupies one of these slots: both + // loops below stop at the first nullptr and have no other bound, so an entry + // that named three dtypes would leave them reading past the end of the + // array. That is undefined behaviour rather than a wrong answer, and it is + // one keystroke away from any edit that widens an entry, so the terminator + // is asserted at compile time below rather than trusted. + static constexpr std::size_t kSlots = 3; + + const char *family; + const char *allowed[kSlots]; +}; + +constexpr DtypeAllowList kDtypeAllowLists[] = { + {"supertonic", {"f32", "i64", nullptr}}, +}; + +constexpr bool allow_lists_are_terminated() { + for (const auto &entry : kDtypeAllowLists) { + if (entry.allowed[DtypeAllowList::kSlots - 1] != nullptr) { + return false; + } + } + return true; +} + +static_assert(allow_lists_are_terminated(), + "every DtypeAllowList must leave its last slot null: the lookups " + "below stop at the first nullptr and would otherwise read past " + "the end of the array"); + +const DtypeAllowList *find_allow_list(const std::string &family) { + for (const auto &entry : kDtypeAllowLists) { + if (family == entry.family) { + return &entry; + } + } + return nullptr; +} + +} // namespace + +bool family_has_weight_dtype_allow_list(const std::string &family) { + return find_allow_list(family) != nullptr; +} + +bool weight_dtype_is_supported(const std::string &family, const std::string &dtype) { + const DtypeAllowList *list = find_allow_list(family); + if (list == nullptr) { + return true; + } + for (const char *const *name = list->allowed; *name != nullptr; ++name) { + if (dtype == *name) { + return true; + } + } + return false; +} + +std::string supported_weight_dtypes(const std::string &family) { + const DtypeAllowList *list = find_allow_list(family); + if (list == nullptr) { + return {}; + } + std::string out; + for (const char *const *name = list->allowed; *name != nullptr; ++name) { + if (!out.empty()) { + out += ", "; + } + out += *name; + } + return out; +} + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/family_gate.h b/backend/cpp/audio-cpp/family_gate.h new file mode 100644 index 000000000..e516ca54b --- /dev/null +++ b/backend/cpp/audio-cpp/family_gate.h @@ -0,0 +1,86 @@ +#pragma once + +// Decides which audio.cpp family a model path belongs to, and refuses paths +// this backend must not claim. Standard library only. +// +// This is the guard against issue #9287. A model config with no explicit +// backend makes LocalAI probe every installed backend and bind to the first +// Load that succeeds, so accepting an arbitrary GGUF here would capture +// unrelated LLMs. audio.cpp GGUFs carry an audiocpp.model_spec.family metadata +// key; llama.cpp GGUFs do not. + +#include + +namespace audiocpp_backend { + +// True when the path ends in ".gguf", case insensitively, and has a stem. +bool path_looks_like_gguf(const std::string &path); + +struct FamilyDecision { + bool ok = false; + std::string family; + // Set when ok is false. Suitable verbatim as an INVALID_ARGUMENT message. + std::string error; +}; + +// Precedence: +// 1. an explicit `family:` option, so a user can override wrong metadata; +// 2. for a GGUF, the family embedded in audiocpp.model_spec.family; +// 3. otherwise refuse. +// A directory path never consults embedded metadata: there is no single GGUF +// to read it from. +FamilyDecision decide_family(bool path_is_gguf, const std::string &embedded_family, + const std::string &configured_family); + +// True when `family` can run weights stored as `dtype`, where dtype is the +// string a TensorMetadata carries ("f32", "f16", "q8_0", "i64", ...). +// +// This is a LIST OF FAMILIES THAT CRASH THE PROCESS, not a list of families that +// perform badly. It exists because the failure is not an exception: loading the +// supertonic f16 GGUF package reaches ggml_concat with one f16 operand and one +// f32 one, GGML_ASSERT(a->type == b->type) fails (external/ggml/src/ggml.c:2595) +// and ggml_abort takes the backend down with SIGABRT on the FIRST request. +// Nothing upstream of the load can catch that, so an operator sees a model that +// loaded successfully and a backend that dies on every request with no status +// and no message. +// +// EVIDENCE, per dtype, because the two are not equally attested: +// - f16 was OBSERVED to abort here, identically through the unary TTS RPC and +// through TTSStream, so it is the packaging and not the streaming path. +// Upstream's docs/gguf.md:90 has supertonic's 16-bit column as "---", which +// its own legend (:53) defines as not tested, so upstream neither confirms +// nor contradicts it. +// - q8_0 was NOT run here. Upstream records it as "No (unsupported weight +// dtype)" in the same row, which is a weaker claim than the f16 abort: it +// says the format is unusable, not that it takes the process down. +// Both are refused, because the allow list is what the family CAN run (f32 for +// weights, i64 for the shape and index constants) rather than a list of the +// dtypes that fail, and a format upstream calls unusable has no business being +// loaded either way. +// +// A family with no entry is unrestricted, which is every family but one. +// +// Split out of loaded_model.cpp, where the caller lives, so that the policy is +// stdlib-only and can be held by a test: the caller needs a real GGUF on disk +// and an engine, and neither is available to a unit test. What the test pins is +// that the table says what it is meant to say, so widening it is a deliberate +// act rather than a typo. It CANNOT pin the removal criterion, which is +// "upstream fixed it": no test can know that without downloading the package and +// synthesising, so that step stays a documented manual one at the table itself. +bool weight_dtype_is_supported(const std::string &family, const std::string &dtype); + +// True when `family` has an entry in the table at all, which is the question a +// caller deciding whether to OPEN THE FILE has to ask. Distinct from +// "supported_weight_dtypes(family) is empty": that string is also empty for an +// entry with an empty allow list, and such an entry means "this family can run +// nothing", which weight_dtype_is_supported already answers by refusing every +// dtype. Deciding from the string would skip the check on precisely the entry +// that most needs it. +bool family_has_weight_dtype_allow_list(const std::string &family); + +// The dtypes `family` is restricted to, as "f32, i64", or empty when it is not +// restricted at all. For the refusal message, so the operator is told what to +// look for rather than only what is wrong. +std::string supported_weight_dtypes(const std::string &family); + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/family_gate_test.cpp b/backend/cpp/audio-cpp/family_gate_test.cpp new file mode 100644 index 000000000..e57995973 --- /dev/null +++ b/backend/cpp/audio-cpp/family_gate_test.cpp @@ -0,0 +1,179 @@ +// Unit tests for family_gate. Standard library only. The harness compiles this +// as a single translation unit, so the implementation is included directly. +// +// This unit is the guard against issue #9287: when a model config has no +// explicit backend, LocalAI probes every installed backend and binds to the +// first Load that succeeds. Accepting an arbitrary GGUF here would capture +// unrelated LLMs. + +#include "family_gate.cpp" + +#include +#include + +static int failures = 0; + +static void check(bool ok, const std::string &name) { + if (!ok) { + failures++; + fprintf(stderr, "FAIL: %s\n", name.c_str()); + } else { + fprintf(stderr, "ok: %s\n", name.c_str()); + } +} + +static void check_eq(const std::string &got, const std::string &want, + const std::string &name) { + check(got == want, name + " (got \"" + got + "\" want \"" + want + "\")"); +} + +using namespace audiocpp_backend; + +static void test_gguf_suffix_detection() { + check(path_looks_like_gguf("/models/chatterbox-q8_0.gguf"), "plain .gguf"); + check(path_looks_like_gguf("/models/CHATTERBOX.GGUF"), "uppercase .GGUF"); + check(path_looks_like_gguf("/models/x.GgUf"), "mixed case .GgUf"); + check(path_looks_like_gguf("a.gguf"), "a one character stem is still a stem"); + check(!path_looks_like_gguf("/models/chatterbox"), "extensionless directory"); + check(!path_looks_like_gguf("/models/model.safetensors"), "safetensors"); + check(!path_looks_like_gguf("/models/gguf"), "a name that is merely 'gguf'"); + check(!path_looks_like_gguf("/models/GGUF"), "an uppercase name that is merely 'GGUF'"); + check(!path_looks_like_gguf(""), "empty path"); + check(!path_looks_like_gguf(".gguf"), "a bare extension is not a model file"); + // The suffix has to be at the end. A prefix or infix match would let + // ".gguf.tmp" download artefacts and ".ggufx" siblings through. + check(!path_looks_like_gguf("/models/model.gguf.tmp"), ".gguf in the middle"); + check(!path_looks_like_gguf("/models/model.ggufx"), "a longer extension"); + // Every character of the suffix has to match, including the last one. + check(!path_looks_like_gguf("/models/model.ggug"), "a near miss in the final character"); + check(!path_looks_like_gguf("/models/model_gguf"), "a near miss in the first character"); + check(!path_looks_like_gguf("/models/.gguf-notes"), "a leading .gguf"); +} + +static void test_explicit_family_always_wins() { + // Explicit configuration beats metadata, so a user can force a family when + // upstream metadata is wrong or absent. + const auto gguf = decide_family(true, "chatterbox", "omnivoice"); + check(gguf.ok && gguf.family == "omnivoice", "explicit family overrides GGUF metadata"); + check(gguf.error.empty(), "an accepted decision carries no error text"); + + const auto dir = decide_family(false, "", "qwen3_tts"); + check(dir.ok && dir.family == "qwen3_tts", "explicit family satisfies a directory path"); + + // A GGUF with no embedded spec is still loadable when the user names the + // family: the option is an override, not a tie-break that needs metadata to + // break against. + const auto bare = decide_family(true, "", "supertonic"); + check(bare.ok && bare.family == "supertonic", + "explicit family rescues a GGUF that carries no spec"); +} + +static void test_gguf_metadata_supplies_the_family() { + const auto d = decide_family(true, "nemotron_asr", ""); + check(d.ok, "an audio.cpp GGUF loads with no family option"); + check(d.family == "nemotron_asr", "family comes from the embedded spec"); + check(d.error.empty(), "an accepted GGUF carries no error text"); +} + +// THE GATE. A llama.cpp GGUF has no audiocpp.model_spec.family key. +static void test_foreign_gguf_is_refused() { + const auto d = decide_family(true, "", ""); + check(!d.ok, "a GGUF with no audio.cpp spec is refused"); + check(d.family.empty(), "no family is guessed"); + check(!d.error.empty(), "a refusal always says why"); + check(d.error.find("audiocpp.model_spec.family") != std::string::npos, + "error names the missing metadata key so the cause is diagnosable"); + check(d.error.find("family:") != std::string::npos, + "error names the option that would override it"); +} + +static void test_directory_without_family_is_refused() { + const auto d = decide_family(false, "", ""); + check(!d.ok, "a non-GGUF path with no family option is refused"); + check(d.family.empty(), "a refused directory guesses no family"); + check(d.error.find("family:") != std::string::npos, + "error names the required option"); +} + +// A directory path never consults embedded metadata, because there is no single +// GGUF to read it from. +static void test_directory_ignores_embedded_family() { + const auto d = decide_family(false, "chatterbox", ""); + check(!d.ok, "a directory is refused even when an embedded family is supplied"); + check(d.family.empty(), "a refused directory does not adopt the embedded family"); + // If the GGUF branch ever leaked into the directory branch this message + // would start blaming a metadata key that a directory has no place to carry. + check(d.error.find("audiocpp.model_spec.family") == std::string::npos, + "a directory refusal does not blame GGUF metadata it could not have"); +} + +// Pins the weight-dtype allow list. Not a style preference: an entry here is a +// family that ABORTS THE PROCESS on the first request when handed the wrong +// dtype, so the model loads and then every request kills the backend with no +// status and no message. +// +// What this test can and cannot do, stated so the next reader does not expect +// more of it: it pins WHAT THE TABLE SAYS, so widening an entry is a deliberate +// act rather than a typo, and it pins that an unlisted family is unrestricted. +// It CANNOT pin the removal criterion, which is "upstream fixed it": knowing +// that needs the package downloaded and a synthesis run, so it stays a manual +// step documented at the table in family_gate.cpp. +static void test_weight_dtype_allow_list() { + // The entry that exists, and the exact reason it exists. + check(!weight_dtype_is_supported("supertonic", "f16"), + "supertonic refuses f16, the package that aborts the process"); + check(!weight_dtype_is_supported("supertonic", "q8_0"), + "supertonic refuses q8_0, which upstream records as unsupported"); + check(!weight_dtype_is_supported("supertonic", "bf16"), + "supertonic refuses bf16, which is untested rather than known good"); + check(weight_dtype_is_supported("supertonic", "f32"), + "supertonic accepts f32, which is what the orig package stores"); + check(weight_dtype_is_supported("supertonic", "i64"), + "supertonic accepts i64: the orig package carries 72 such tensors and " + "refusing them would refuse the artifact that works"); + + // Every other family is unrestricted, and must stay that way: this guard is + // for process death, not for quality. + check(weight_dtype_is_supported("nemotron_asr", "q8_0"), + "an unlisted family is not restricted"); + check(weight_dtype_is_supported("citrinet_asr", "f16"), + "an unlisted family is not restricted by another family's entry"); + check(weight_dtype_is_supported("", "anything"), + "an empty family name is not restricted"); + + // The message the operator reads has to name the remedy, so the refusal is + // actionable rather than only correct. + check_eq(supported_weight_dtypes("supertonic"), "f32, i64", + "the refusal can name what to look for"); + check_eq(supported_weight_dtypes("nemotron_asr"), "", + "an unlisted family reports no restriction"); + + // What the caller actually decides on, and it is a DIFFERENT question from + // "is the description empty": an entry with an empty allow list would + // describe itself as "" while refusing every dtype, so a caller that skipped + // the file read on the empty string would skip the check on the one entry + // that refuses everything. + check(family_has_weight_dtype_allow_list("supertonic"), + "a listed family has an allow list"); + check(!family_has_weight_dtype_allow_list("nemotron_asr"), + "an unlisted family has none, which is what lets the caller skip " + "opening the file at all"); + check(!family_has_weight_dtype_allow_list(""), + "an empty family name has no allow list"); +} + +int main() { + test_gguf_suffix_detection(); + test_explicit_family_always_wins(); + test_gguf_metadata_supplies_the_family(); + test_foreign_gguf_is_refused(); + test_directory_without_family_is_refused(); + test_directory_ignores_embedded_family(); + test_weight_dtype_allow_list(); + if (failures) { + fprintf(stderr, "%d check(s) failed\n", failures); + return 1; + } + fprintf(stderr, "all family_gate checks passed\n"); + return 0; +} diff --git a/backend/cpp/audio-cpp/generation_request.cpp b/backend/cpp/audio-cpp/generation_request.cpp new file mode 100644 index 000000000..f74cb1a21 --- /dev/null +++ b/backend/cpp/audio-cpp/generation_request.cpp @@ -0,0 +1,280 @@ +#include "generation_request.h" + +#include +#include +#include +#include + +namespace audiocpp_backend { +namespace { + +const char *bool_option(bool value) { return value ? "true" : "false"; } + +} // namespace + +bool voice_is_reference_file(const std::string &voice) { + if (voice.empty()) { + return false; + } + std::error_code ec; + return std::filesystem::is_regular_file(std::filesystem::path(voice), ec); +} + +RequestShape build_tts_shape(const backend::TTSRequest &request) { + RequestShape shape; + shape.has_voice_reference = voice_is_reference_file(request.voice()); + // !empty() as well as has_instructions(), and it must match the guard in + // build_tts_request: a request whose instructions are an empty string + // carries no style condition, so telling routing to prefer VoiceDesign for + // it would route to a task with nothing to design from. + shape.has_instructions = + request.has_instructions() && !request.instructions().empty(); + return shape; +} + +engine::runtime::TaskRequest +build_tts_request(const backend::TTSRequest &request, + std::optional reference_audio) { + engine::runtime::TaskRequest task; + + // The Transcript, not an option, is where every TTS family reads its + // language: chatterbox normalises request.text_input->language into its + // voice-clone config, qwen3_tts reads it as out.language, ace_step turns it + // into vocal_language. Upstream's own HTTP server does the same and sets no + // language option at all (app/server/runtime.cpp build_speech_request). + engine::runtime::Transcript transcript; + transcript.text = request.text(); + if (request.has_language()) { + transcript.language = request.language(); + } + task.text_input = std::move(transcript); + + engine::runtime::VoiceCondition condition; + bool condition_used = false; + + if (reference_audio.has_value()) { + // A clip: VoiceReference::audio, at the file's own rate and channel + // count. See kVoiceReferenceSampleRate in grpc-server.cpp for why it is + // not folded first. + engine::runtime::VoiceReference reference; + reference.audio = std::move(*reference_audio); + condition.speaker = std::move(reference); + condition_used = true; + } else if (!request.voice().empty()) { + // A named preset. cached_voice_id is the channel that actually lands: + // supertonic (options.voice), pocket_tts (voice_config.preset_name), + // voxcpm2, vibevoice, fish_audio (its saved-reference lookup) and + // qwen3_tts CustomVoice all read request.voice->speaker->cached_voice_id, + // and upstream's own server puts a non-preset `voice` body field in + // exactly this slot (app/server/runtime.cpp build_speech_request). + engine::runtime::VoiceReference reference; + reference.cached_voice_id = request.voice(); + condition.speaker = std::move(reference); + condition_used = true; + // Forward-tolerant alias only. A bare "voice" REQUEST OPTION is read by + // no family in the pinned upstream: grepping find_option for it returns + // nothing. It is sent so a family adopting the name later works with no + // change here, not because it does anything today. + task.options["voice"] = request.voice(); + } + + if (request.has_instructions() && !request.instructions().empty()) { + // "instruct" is the key upstream itself maps the OpenAI `instructions` + // body field onto (app/server/runtime.cpp: request.options["instruct"] + // = value->as_string()), and it is read: qwen3_tts VoiceDesign and + // CustomVoice both take find_option(options, {"instruct"}) first, and + // omnivoice reads it in resolve_instruct. + task.options["instruct"] = request.instructions(); + // "caption" is irodori_tts's name for the same thing, read in its + // make_request and documented as the voice-design caption for the 600M + // VoiceDesign model (docs/tts.md). Without this, that family's voice + // design cannot be driven from this RPC at all. + task.options["caption"] = request.instructions(); + // The proto field's own name, forwarded for the same forward-tolerant + // reason as "voice" above and with the same honest accounting: NO family + // in the pinned upstream reads a request option called "instructions". + task.options["instructions"] = request.instructions(); + + engine::runtime::StyleCondition style; + // "instruct", not "instructions". This tag IS read, and only under that + // spelling: omnivoice and qwen3_tts both fall back to + // request.voice->style->tags.find("instruct") when the option is absent. + // Spelling it "instructions" here would have made the whole + // StyleCondition dead weight. + style.tags["instruct"] = request.instructions(); + // !empty(), matching the option emission below, and load-bearing rather + // than tidiness. core/backend/tts.go's newTTSRequest sets + // `Language: &language` UNCONDITIONALLY, so has_language() is true on + // every request LocalAI sends and carries "" whenever the caller named + // no language. An engaged-but-empty style language is WORSE than an + // absent one: supertonic reads text_input->language behind its own + // !empty() guard and then OVERRIDES it from style->language with no + // guard at all (supertonic/session.cpp), so "" would replace its "en" + // default and tokenizer_text.cpp would throw + // "invalid Supertonic language: " on every request that set + // instructions and no language. + if (request.has_language() && !request.language().empty()) { + style.language = request.language(); + } + condition.style = std::move(style); + condition_used = true; + } + + if (condition_used) { + task.voice = std::move(condition); + } + + if (request.has_language() && !request.language().empty()) { + // Forward-tolerant alias, exactly as in build_transcription_request. The + // families that read a "language" request option are the ASR ones + // (nemotron_asr, hviske_asr, vibevoice_asr, higgs_audio_stt), none of + // which this RPC can route to; pocket_tts reads one but from its + // ModelLoadRequest at load time, not from here. The Transcript above is + // what actually carries the language to a TTS family. + task.options["language"] = request.language(); + } + + // LAST, so an explicit params entry wins over anything derived above. That + // matters for "caption": a caller who sets params[caption] has named the + // exact string they want, and it must not be overwritten by `instructions`. + for (const auto ¶m : request.params()) { + task.options[param.first] = param.second; + } + return task; +} + +engine::runtime::TaskRequest +build_sound_generation_request(const backend::SoundGenerationRequest &request, + std::optional source_audio) { + engine::runtime::TaskRequest task; + + engine::runtime::Transcript transcript; + transcript.text = request.text(); + if (request.has_language()) { + transcript.language = request.language(); + } + task.text_input = std::move(transcript); + + // src is the input clip for the editing routes. ace_step's repaint, cover + // and edit routes need it; stable_audio uses it as init_audio or + // inpaint_audio; heartmula refuses it outright. + if (source_audio.has_value()) { + task.audio_input = std::move(*source_audio); + } + + // WHAT LANDS AND WHAT DOES NOT. Three families advertise AudioGeneration in + // the pinned upstream: ace_step, heartmula and stable_audio. Every key below + // was grepped against find_option/parse_*_option in src/ and include/ rather + // than assumed, because a key nobody reads is not a feature and shipping one + // while implying it works is the mistake this comment exists to prevent. + // + // Unknown REQUEST options cannot turn a valid request into an error: + // families look theirs up by name and ignore the rest, and the unknown-key + // refusals upstream does have are on SESSION options, which arrive at load + // time. So a forward-tolerant alias is free; it is just not a feature. + + if (request.has_duration()) { + // duration_seconds is the key that works, and it works everywhere: + // ace_step (request_parser.cpp), heartmula (session.cpp, which also + // refuses a non-positive value) and stable_audio (request.cpp) all read + // it. This is the SoundGeneration analogue of Task 9's return_timestamps. + task.options["duration_seconds"] = std::to_string(request.duration()); + // The proto field's own name. Read by exactly one family, omnivoice, and + // omnivoice advertises Tts rather than AudioGeneration, so this RPC can + // never route to it: DEAD here, kept only as a forward-tolerant alias. + task.options["duration"] = std::to_string(request.duration()); + } + if (request.has_temperature()) { + // Read by heartmula. ace_step's sampling temperature is a different, + // narrower knob it calls lm_temperature (it drives the caption/thinking + // LM, not the audio diffusion), so this is deliberately NOT mapped onto + // it; a caller who wants it sets it through the request options that + // reach ace_step by name. stable_audio has no temperature at all. + task.options["temperature"] = std::to_string(request.temperature()); + } + if (request.has_sample()) { + // do_sample is read widely upstream, but only by TTS and ASR families + // (chatterbox, index_tts2, miotts, moss, qwen3_tts, vibevoice, + // hviske_asr, voxtral_realtime). NO AudioGeneration family reads it, so + // it is dead on this route. + task.options["do_sample"] = bool_option(request.sample()); + } + if (request.has_src_divisor()) { + // Read by nobody, anywhere in the pinned upstream. Forwarded because the + // proto documents it as part of this request and a family adopting it + // then works unchanged. + task.options["src_divisor"] = std::to_string(request.src_divisor()); + } + if (request.has_think()) { + // "thinking" is the key ace_step actually reads (request_parser.cpp), + // which is why it is sent alongside the proto's own "think". "think" on + // its own is read by nobody. + task.options["thinking"] = bool_option(request.think()); + task.options["think"] = bool_option(request.think()); + } + if (request.has_caption()) { + // Read only by irodori_tts, which advertises Tts/VoiceCloning/VoiceDesign + // and not AudioGeneration, so it is unreachable from this RPC: DEAD here. + task.options["caption"] = request.caption(); + } + if (request.has_lyrics()) { + // Read by ace_step and heartmula. + task.options["lyrics"] = request.lyrics(); + } + if (request.has_bpm()) { + // Read by ace_step. + task.options["bpm"] = std::to_string(request.bpm()); + } + if (request.has_keyscale()) { + // Read by ace_step. + task.options["keyscale"] = request.keyscale(); + } + if (request.has_timesignature()) { + // Read by ace_step. + task.options["timesignature"] = request.timesignature(); + } + if (request.has_instrumental()) { + // Read by nobody: "instrumental" appears in the pinned upstream only as + // a roformer STEM NAME, never as a request option. A forward-tolerant + // alias and nothing more. + task.options["instrumental"] = bool_option(request.instrumental()); + } + if (request.has_language() && !request.language().empty()) { + // Alias again: the Transcript above is what ace_step reads as + // vocal_language. No AudioGeneration family reads a "language" option. + task.options["language"] = request.language(); + } + return task; +} + +bool apply_transform_text_input(engine::runtime::TaskRequest &task) { + // Canonical first, alias second, and an empty value falls through to the + // next candidate rather than ending the search: a caller who sent + // target_text="" and text="the real one" meant the second one. + static const char *const kTextKeys[] = {"target_text", "text"}; + + std::string text; + for (const char *key : kTextKeys) { + const auto found = task.options.find(key); + if (found != task.options.end() && !found->second.empty()) { + text = found->second; + break; + } + } + if (text.empty()) { + return false; + } + + engine::runtime::Transcript transcript; + transcript.text = std::move(text); + // Inside the has-text branch on purpose. See the header: a language on its + // own conditions nothing and must not manufacture a text_input. + const auto language = task.options.find("language"); + if (language != task.options.end()) { + transcript.language = language->second; + } + task.text_input = std::move(transcript); + return true; +} + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/generation_request.h b/backend/cpp/audio-cpp/generation_request.h new file mode 100644 index 000000000..c497fff0c --- /dev/null +++ b/backend/cpp/audio-cpp/generation_request.h @@ -0,0 +1,108 @@ +#pragma once + +// Builds the engine::runtime::TaskRequest for the two audio-PRODUCING offline +// RPCs, TTS and SoundGeneration, and answers the one filesystem question TTS +// routing depends on. +// +// It is a unit of its own rather than a pair of statics in grpc-server.cpp so +// that it can be tested: grpc-server.cpp has a main() and cannot be linked into +// a test binary, and everything here is a pure function of its arguments once +// the file read has been lifted out (which is why the reference clip arrives as +// an already-read buffer rather than a path). TTSStream reuses build_tts_request +// unchanged. +// +// Only the plain structs in engine/framework/runtime/session.h are touched, so +// this compiles against the header without linking engine_runtime, the same way +// result_map does. + +#include "backend.pb.h" +#include "capability_routing.h" + +#include "engine/framework/runtime/session.h" + +#include +#include + +namespace audiocpp_backend { + +// True when TTSRequest.voice names an existing regular file, in which case it +// is a speaker reference clip and routing prefers VoiceCloning; false when it is +// a named preset (or empty). +// +// The overload is LocalAI's, not this backend's: `voice` is the OpenAI speech +// field and different LocalAI backends have always read it both ways. Deciding +// it from the filesystem needs no new option and matches how somebody actually +// configures a cloning family, which is by pointing at a clip. +// +// A DIRECTORY is deliberately not a reference: is_regular_file, not exists. A +// directory named as a voice cannot be read as a WAV, and treating it as a +// reference would turn a preset typo into "cannot read /x as WAV" instead of +// letting it travel as the preset name it looks like. +// +// The error_code overload is used so an unreadable parent directory answers +// false rather than throwing. That is the right answer here: the name is then +// passed on as a preset, and if it really was meant to be a clip the family +// refuses a request it cannot serve, which is a better message than a +// filesystem exception thrown while classifying a string. +bool voice_is_reference_file(const std::string &voice); + +// Everything routing needs to know about a TTSRequest, in one place, so that TTS +// and TTSStream cannot describe the same request differently. +// +// `pinned_task` is deliberately NOT filled here: it comes off the LoadedModel, +// not off the request, and this unit links no engine. The caller must still +// write `shape.pinned_task = model->pinned_task();` or the model's `task:` +// option is dead. That is the one field a new handler can forget, so it is the +// one field left visible at the call site rather than hidden behind this +// helper. +RequestShape build_tts_shape(const backend::TTSRequest &request); + +// `reference_audio` is the already-read speaker clip, present exactly when +// voice_is_reference_file(request.voice()) was true. Passing it in rather than a +// path keeps this function pure and lets the caller do the read where the +// ordering rules (capability refusal first, then the lane) are enforced. +// +// It is taken BY VALUE and moved in: a reference clip is seconds of audio and +// the caller has no use for it afterwards. +engine::runtime::TaskRequest +build_tts_request(const backend::TTSRequest &request, + std::optional reference_audio); + +// `source_audio` is SoundGenerationRequest.src already read, present exactly +// when the field was set and non-empty. Same reasoning as above. +engine::runtime::TaskRequest +build_sound_generation_request(const backend::SoundGenerationRequest &request, + std::optional source_audio); + +// Lifts a text-conditioned transform route's text out of the request params +// into TaskRequest.text_input, and reports whether it set one. +// +// WHY THIS EXISTS. AudioTransform is an audio-in / audio-out RPC and its proto +// message has no text field, but not every task it routes to is audio-only. +// vevo2's speech-to-speech and prosody routes read their text from +// request.text_input (src/models/vevo2/session.cpp fills refs.target_text from +// exactly there and nowhere else) and refuse the run without one: "Vevo2 +// text/prosody route requires text_input or target_text". The params map is the +// only channel AudioTransform has that reaches the engine, so the text travels +// through it and is unpacked here. Without this, s2s is not merely awkward to +// reach through this RPC, it is unreachable. +// +// CALL IT AFTER the params have been copied into task.options, and note that it +// does NOT erase the keys it reads. vevo2's loader advertises "target_text" in +// its own documented request-option table, so a family that looks there keeps +// finding it; the copy in text_input is what the session actually reads today. +// +// "target_text" is canonical and "text" is its alias, the same order vevo2's +// option table declares them in. A request setting both gets target_text, so +// the canonical spelling wins rather than whichever the map happened to store +// first. An empty value is not a text: it means the caller sent the key with +// nothing in it, and a family asked to vocalise "" should say so itself rather +// than be handed an empty Transcript that looks deliberate. +// +// "language" rides along when a text was found, and only then. On its own it +// conditions nothing, and setting text_input for it alone would turn a plain +// separation request that happened to carry a language hint into a text-routed +// one. +bool apply_transform_text_input(engine::runtime::TaskRequest &task); + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/generation_request_ctest.cpp b/backend/cpp/audio-cpp/generation_request_ctest.cpp new file mode 100644 index 000000000..507b2054d --- /dev/null +++ b/backend/cpp/audio-cpp/generation_request_ctest.cpp @@ -0,0 +1,575 @@ +// Tests for the TTS and SoundGeneration request builders, and for the +// filesystem rule that decides whether TTSRequest.voice is a speaker reference +// clip or a named preset. +// +// NAMED _ctest AND NOT _test ON PURPOSE: see the note at the top of +// result_map_ctest.cpp. This file needs the generated protobuf messages and the +// audio.cpp include path, neither of which backend/cpp/run-unit-tests.sh +// provides, so it is built and run by ctest: +// +// make -C backend/cpp/audio-cpp test-engine +// +// The assertions on OPTION KEYS are the point of this file, not decoration. +// Every one of them names a key that was grepped against the pinned upstream: +// "instruct" is read, "instructions" is not; "duration_seconds" is read, +// "duration" is not. A rename that looks harmless is exactly the change that +// silently stops a family honouring the request, so the spellings are pinned +// here rather than left to a comment. + +#include "generation_request.h" + +#include +#include +#include +#include +#include + +static int failures = 0; + +static void check(bool ok, const std::string &name) { + if (!ok) { + failures++; + fprintf(stderr, "FAIL: %s\n", name.c_str()); + } else { + fprintf(stderr, "ok: %s\n", name.c_str()); + } +} + +using namespace audiocpp_backend; + +static bool has_key(const std::unordered_map &options, + const std::string &key) { + return options.find(key) != options.end(); +} + +static std::string option_or( + const std::unordered_map &options, + const std::string &key, const std::string &fallback) { + const auto it = options.find(key); + return it == options.end() ? fallback : it->second; +} + +static engine::runtime::AudioBuffer clip(int sample_rate, int channels) { + engine::runtime::AudioBuffer buffer; + buffer.sample_rate = sample_rate; + buffer.channels = channels; + // Distinguishable content, so a builder that swapped one buffer for another + // (or default-constructed one) is visible rather than merely a size change. + buffer.samples = {0.25f, -0.5f, 0.75f, -1.0f}; + return buffer; +} + +static void test_voice_is_reference_file() { + const auto dir = std::filesystem::temp_directory_path() / + "audiocpp-generation-request-ctest"; + std::filesystem::remove_all(dir); + std::filesystem::create_directories(dir); + + const auto file = dir / "reference.wav"; + { + std::ofstream out(file, std::ios::binary); + out << "not really a wav, but a regular file"; + } + const auto subdir = dir / "a-directory"; + std::filesystem::create_directories(subdir); + + check(!voice_is_reference_file(""), "voice_is_reference_file: empty"); + check(!voice_is_reference_file("alloy"), + "voice_is_reference_file: bare preset name"); + check(!voice_is_reference_file((dir / "absent.wav").string()), + "voice_is_reference_file: missing path"); + check(voice_is_reference_file(file.string()), + "voice_is_reference_file: existing regular file"); + // A directory is NOT a reference. exists() would say yes here and the read + // would then fail with "cannot read as WAV", which sends the operator + // after a file problem instead of a preset typo. + check(!voice_is_reference_file(subdir.string()), + "voice_is_reference_file: directory is not a reference"); + + std::filesystem::remove_all(dir); +} + +// build_tts_shape is what TTS and TTSStream both hand to routing, so a wrong +// answer here silently changes which task a request runs as, with a 200 and no +// diagnostic. Every field is asserted in both directions. +static void test_tts_shape() { + const auto dir = std::filesystem::temp_directory_path() / + "audiocpp-generation-request-ctest-shape"; + std::filesystem::remove_all(dir); + std::filesystem::create_directories(dir); + const auto file = dir / "reference.wav"; + { + std::ofstream out(file, std::ios::binary); + out << "a regular file"; + } + + { + backend::TTSRequest request; + request.set_text("hello"); + const auto shape = build_tts_shape(request); + check(!shape.has_voice_reference && !shape.has_instructions, + "shape: bare request has neither signal"); + // Never filled here: it comes off the LoadedModel, and leaving it empty + // is what makes the caller's assignment visible at the call site. + check(shape.pinned_task.empty(), "shape: pinned_task is left to the caller"); + } + { + backend::TTSRequest request; + request.set_voice(file.string()); + const auto shape = build_tts_shape(request); + check(shape.has_voice_reference, + "shape: an existing file is a voice reference"); + check(!shape.has_instructions, "shape: a clip is not an instruction"); + } + { + backend::TTSRequest request; + request.set_voice("alloy"); + const auto shape = build_tts_shape(request); + check(!shape.has_voice_reference, + "shape: a preset name is not a voice reference"); + } + { + backend::TTSRequest request; + request.set_voice(dir.string()); + const auto shape = build_tts_shape(request); + check(!shape.has_voice_reference, + "shape: a directory is not a voice reference"); + } + { + backend::TTSRequest request; + request.set_instructions("a calm older man"); + const auto shape = build_tts_shape(request); + check(shape.has_instructions, "shape: instructions are seen"); + check(!shape.has_voice_reference, + "shape: instructions do not imply a reference"); + } + { + // The guard that has to match build_tts_request's. An empty + // instructions string builds no style condition, so telling routing to + // prefer VoiceDesign for it would route to a task with nothing to + // design from. + backend::TTSRequest request; + request.set_instructions(""); + const auto shape = build_tts_shape(request); + check(!shape.has_instructions, + "shape: an empty instructions string is not an instruction"); + } + { + backend::TTSRequest request; + request.set_voice(file.string()); + request.set_instructions("a calm older man"); + const auto shape = build_tts_shape(request); + check(shape.has_voice_reference && shape.has_instructions, + "shape: both signals are reported when both are set"); + } + + std::filesystem::remove_all(dir); +} + +static void test_tts_plain() { + backend::TTSRequest request; + request.set_text("hello there"); + + const auto task = build_tts_request(request, std::nullopt); + + check(task.text_input.has_value() && task.text_input->text == "hello there", + "tts: text reaches the transcript"); + // No voice and no instructions means NO voice condition at all. A builder + // that always emitted one would make every family think a speaker was + // named, and chatterbox in particular refuses a prepare whose voice + // condition carries neither audio nor anything else it can use. + check(!task.voice.has_value(), "tts: no voice condition when nothing is set"); + check(task.options.empty(), "tts: no options when nothing is set"); + check(!task.audio_input.has_value(), "tts: no audio input"); +} + +static void test_tts_named_preset() { + backend::TTSRequest request; + request.set_text("hello"); + request.set_voice("alloy"); + + const auto task = build_tts_request(request, std::nullopt); + + check(task.voice.has_value() && task.voice->speaker.has_value(), + "tts preset: speaker condition present"); + check(task.voice->speaker->cached_voice_id.has_value() && + *task.voice->speaker->cached_voice_id == "alloy", + "tts preset: lands in cached_voice_id"); + // The clip slot must stay empty, or a cloning family would try to prepare + // conditionals from a default-constructed buffer. + check(!task.voice->speaker->audio.has_value(), + "tts preset: no reference audio"); + check(!task.voice->style.has_value(), "tts preset: no style condition"); + check(option_or(task.options, "voice", "") == "alloy", + "tts preset: forwarded as the voice option too"); +} + +static void test_tts_reference_clip() { + backend::TTSRequest request; + request.set_text("hello"); + request.set_voice("/tmp/reference.wav"); + + const auto task = build_tts_request(request, clip(44100, 2)); + + check(task.voice.has_value() && task.voice->speaker.has_value(), + "tts clip: speaker condition present"); + check(task.voice->speaker->audio.has_value(), + "tts clip: reference audio present"); + // Rate and channels survive untouched. This is the assertion that fails if + // anybody decides to fold the clip to 16 kHz mono on the way in. + check(task.voice->speaker->audio->sample_rate == 44100 && + task.voice->speaker->audio->channels == 2 && + task.voice->speaker->audio->samples.size() == 4, + "tts clip: rate, channels and samples pass through unchanged"); + // A clip is NOT also a cached voice id, and the path must not travel as a + // preset name: a family reading cached_voice_id would then look up a voice + // called "/tmp/reference.wav". + check(!task.voice->speaker->cached_voice_id.has_value(), + "tts clip: no cached_voice_id"); + check(!has_key(task.options, "voice"), "tts clip: no voice option"); +} + +static void test_tts_instructions() { + backend::TTSRequest request; + request.set_text("hello"); + request.set_instructions("a calm older man, speaking slowly"); + + const auto task = build_tts_request(request, std::nullopt); + + check(option_or(task.options, "instruct", "") == + "a calm older man, speaking slowly", + "tts instructions: instruct option is the one qwen3_tts reads"); + check(option_or(task.options, "caption", "") == + "a calm older man, speaking slowly", + "tts instructions: caption option is the one irodori_tts reads"); + check(option_or(task.options, "instructions", "") == + "a calm older man, speaking slowly", + "tts instructions: proto field name forwarded as an alias"); + check(task.voice.has_value() && task.voice->style.has_value(), + "tts instructions: style condition present"); + // "instruct", not "instructions". omnivoice and qwen3_tts both look this tag + // up by that exact name and by no other. + check(option_or(task.voice->style->tags, "instruct", "") == + "a calm older man, speaking slowly", + "tts instructions: style tag is spelled instruct"); + check(!has_key(task.voice->style->tags, "instructions"), + "tts instructions: style tag is NOT spelled instructions"); + // Instructions alone must not invent a speaker: has_voice_reference is what + // routing keys VoiceCloning off, and a speaker here would make a + // clone-capable family expect a clip it never received. + check(!task.voice->speaker.has_value(), + "tts instructions: no speaker without a voice"); +} + +static void test_tts_empty_instructions_are_not_instructions() { + backend::TTSRequest request; + request.set_text("hello"); + request.set_instructions(""); + + const auto task = build_tts_request(request, std::nullopt); + + // has_instructions() is true here, because the field was set. An empty + // string is still no instruction, and forwarding it would set an empty + // instruct option that qwen3_tts would prefer over its style tag fallback. + check(!has_key(task.options, "instruct"), + "tts: an empty instructions string sets no instruct option"); + check(!task.voice.has_value(), + "tts: an empty instructions string sets no voice condition"); +} + +// THE EXACT SHAPE LocalAI PUTS ON THE WIRE. core/backend/tts.go's newTTSRequest +// sets Language: &language UNCONDITIONALLY, so has_language() is true on every +// request that ever reaches this backend, carrying an empty string whenever the +// caller named no language. +// +// An empty StyleCondition::language is not a harmless default. supertonic reads +// text_input->language behind a !empty() guard and then OVERRIDES it from +// style->language whenever that optional is engaged, with no guard at all +// (supertonic/session.cpp generation_options_from_request), so an empty style +// language replaces its "en" default (session.h) with "" and +// tokenizer_text.cpp's preprocess throws "invalid Supertonic language: ". +// Every /v1/audio/speech request carrying instructions and no language would be +// an INTERNAL. A plain request never sees it, because the style condition only +// exists when instructions are non-empty. +static void test_tts_empty_language_is_not_a_language() { + backend::TTSRequest request; + request.set_text("hello"); + request.set_instructions("a calm older man"); + request.set_language(""); + + const auto task = build_tts_request(request, std::nullopt); + + check(task.voice.has_value() && task.voice->style.has_value(), + "tts empty language: the style condition still exists"); + check(!task.voice->style->language.has_value(), + "tts empty language: style language is left unset, not set to empty"); + // The option emission has always guarded on !empty(); this pins the two to + // the same rule so they cannot drift apart again. + check(!has_key(task.options, "language"), + "tts empty language: no language option"); + check(task.text_input.has_value() && task.text_input->language.empty(), + "tts empty language: transcript language stays empty"); +} + +// The one shape routing treats specially: a clip outranks instructions, and the +// VoiceCondition then has to carry BOTH, because the family that wins is chosen +// on the clip but may still read the style tag. +static void test_tts_clip_and_instructions() { + backend::TTSRequest request; + request.set_text("hello"); + request.set_voice("/tmp/reference.wav"); + request.set_instructions("bright and fast"); + request.set_language("en"); + + const auto task = build_tts_request(request, clip(22050, 1)); + + check(task.voice.has_value(), "tts clip+instructions: voice condition present"); + check(task.voice->speaker.has_value() && + task.voice->speaker->audio.has_value() && + task.voice->speaker->audio->sample_rate == 22050, + "tts clip+instructions: speaker carries the clip"); + check(!task.voice->speaker->cached_voice_id.has_value(), + "tts clip+instructions: the clip path is not also a preset id"); + check(task.voice->style.has_value() && + option_or(task.voice->style->tags, "instruct", "") == "bright and fast", + "tts clip+instructions: style carries the instruct tag"); + check(task.voice->style->language.has_value() && + *task.voice->style->language == "en", + "tts clip+instructions: a real language does reach the style condition"); + check(option_or(task.options, "instruct", "") == "bright and fast", + "tts clip+instructions: instruct option still emitted"); + check(!has_key(task.options, "voice"), + "tts clip+instructions: still no voice option for a clip"); +} + +static void test_tts_language_and_params() { + backend::TTSRequest request; + request.set_text("ciao"); + request.set_language("it"); + request.set_instructions("warm"); + (*request.mutable_params())["exaggeration"] = "0.7"; + // An explicit param must win over the value derived from instructions. + (*request.mutable_params())["caption"] = "explicitly chosen caption"; + + const auto task = build_tts_request(request, std::nullopt); + + check(task.text_input.has_value() && task.text_input->language == "it", + "tts language: reaches the transcript"); + check(option_or(task.options, "language", "") == "it", + "tts language: forwarded as an option alias"); + check(task.voice.has_value() && task.voice->style.has_value() && + task.voice->style->language.has_value() && + *task.voice->style->language == "it", + "tts language: reaches the style condition"); + check(option_or(task.options, "exaggeration", "") == "0.7", + "tts params: passed through verbatim"); + check(option_or(task.options, "caption", "") == "explicitly chosen caption", + "tts params: an explicit param overrides the derived caption"); +} + +static void test_sound_generation_minimal() { + backend::SoundGenerationRequest request; + request.set_text("a distant thunderstorm"); + + const auto task = build_sound_generation_request(request, std::nullopt); + + check(task.text_input.has_value() && + task.text_input->text == "a distant thunderstorm", + "sound: text reaches the transcript"); + // Unset optionals must emit NOTHING. Emitting a zero for an unset duration + // would make heartmula refuse the request ("duration_seconds must be + // positive") on a request that never mentioned a duration. + check(task.options.empty(), "sound: unset optionals emit no options"); + check(!task.audio_input.has_value(), "sound: no audio input without src"); + check(!task.voice.has_value(), "sound: no voice condition"); +} + +static void test_sound_generation_full() { + backend::SoundGenerationRequest request; + request.set_text("a slow blues in E"); + request.set_duration(30.0f); + request.set_temperature(0.8f); + request.set_sample(false); + request.set_src_divisor(4); + request.set_think(true); + request.set_caption("smoky bar recording"); + request.set_lyrics("first line\nsecond line"); + request.set_bpm(72); + request.set_keyscale("E minor"); + request.set_language("en"); + request.set_timesignature("4/4"); + request.set_instrumental(true); + + const auto task = build_sound_generation_request(request, clip(48000, 2)); + + // duration_seconds is the key every AudioGeneration family actually reads; + // "duration" rides along as an alias. Both spellings are pinned so a + // "cleanup" that keeps only the proto's own name is a test failure and not + // a silent loss of the duration. + check(option_or(task.options, "duration_seconds", "").rfind("30.", 0) == 0, + "sound: duration lands as duration_seconds"); + check(option_or(task.options, "duration", "").rfind("30.", 0) == 0, + "sound: duration also forwarded under its own name"); + check(option_or(task.options, "temperature", "").rfind("0.8", 0) == 0, + "sound: temperature forwarded"); + // Set to FALSE, so this also proves the key is written whenever the field is + // present rather than only when the value is truthy. + check(option_or(task.options, "do_sample", "") == "false", + "sound: sample=false is forwarded as do_sample=false"); + check(option_or(task.options, "src_divisor", "") == "4", + "sound: src_divisor forwarded"); + check(option_or(task.options, "thinking", "") == "true", + "sound: think lands as thinking, the key ace_step reads"); + check(option_or(task.options, "think", "") == "true", + "sound: think also forwarded under its own name"); + check(option_or(task.options, "caption", "") == "smoky bar recording", + "sound: caption forwarded"); + check(option_or(task.options, "lyrics", "") == "first line\nsecond line", + "sound: lyrics forwarded"); + check(option_or(task.options, "bpm", "") == "72", "sound: bpm forwarded"); + check(option_or(task.options, "keyscale", "") == "E minor", + "sound: keyscale forwarded"); + check(option_or(task.options, "timesignature", "") == "4/4", + "sound: timesignature forwarded"); + check(option_or(task.options, "instrumental", "") == "true", + "sound: instrumental forwarded"); + check(option_or(task.options, "language", "") == "en", + "sound: language forwarded as an option alias"); + check(task.text_input->language == "en", + "sound: language reaches the transcript, which is what ace_step reads"); + check(task.audio_input.has_value() && + task.audio_input->sample_rate == 48000 && + task.audio_input->channels == 2 && + task.audio_input->samples.size() == 4, + "sound: src passes through at its own rate and channel count"); +} + + +// --------------------------------------------------------------------------- +// apply_transform_text_input +// +// AudioTransform has no text field on the wire, so a text-conditioned route +// (vevo2's speech-to-speech) can only be reached if the text travels as a +// param and is unpacked into text_input. Every assertion below pins a spelling +// or a precedence that a family actually depends on, not a shape that merely +// looks tidy. + +static void test_transform_text_absent() { + engine::runtime::TaskRequest task; + task.options["stem"] = "vocals"; + check(!apply_transform_text_input(task), + "transform text: reports false when no text key is present"); + check(!task.text_input.has_value(), + "transform text: a request with no text keeps text_input unset"); +} + +static void test_transform_text_canonical_key() { + engine::runtime::TaskRequest task; + task.options["target_text"] = "sing this line"; + check(apply_transform_text_input(task), "transform text: target_text reports true"); + check(task.text_input.has_value() && task.text_input->text == "sing this line", + "transform text: target_text becomes text_input.text"); + check(has_key(task.options, "target_text"), + "transform text: target_text survives in options for families that read it there"); +} + +static void test_transform_text_alias_key() { + engine::runtime::TaskRequest task; + task.options["text"] = "say this instead"; + check(apply_transform_text_input(task), "transform text: text alias reports true"); + check(task.text_input.has_value() && task.text_input->text == "say this instead", + "transform text: the text alias becomes text_input.text"); +} + +static void test_transform_text_canonical_wins() { + engine::runtime::TaskRequest task; + task.options["target_text"] = "canonical"; + task.options["text"] = "alias"; + check(apply_transform_text_input(task), "transform text: both keys reports true"); + check(task.text_input.has_value() && task.text_input->text == "canonical", + "transform text: target_text wins over text, not whichever hashed first"); +} + +static void test_transform_text_empty_is_not_a_text() { + engine::runtime::TaskRequest task; + task.options["target_text"] = ""; + check(!apply_transform_text_input(task), + "transform text: an empty target_text reports false"); + check(!task.text_input.has_value(), + "transform text: an empty target_text leaves text_input unset"); +} + +static void test_transform_text_empty_canonical_falls_through_to_alias() { + engine::runtime::TaskRequest task; + task.options["target_text"] = ""; + task.options["text"] = "the real one"; + check(apply_transform_text_input(task), + "transform text: an empty canonical key does not mask a usable alias"); + check(task.text_input.has_value() && task.text_input->text == "the real one", + "transform text: the alias is used when the canonical key is empty"); +} + +static void test_transform_text_language_rides_along() { + engine::runtime::TaskRequest task; + task.options["target_text"] = "vocalise me"; + task.options["language"] = "ja"; + check(apply_transform_text_input(task), "transform text: text plus language reports true"); + check(task.text_input.has_value() && task.text_input->language == "ja", + "transform text: language lands on the Transcript alongside the text"); + check(has_key(task.options, "language"), + "transform text: language survives in options too"); +} + +static void test_transform_language_alone_is_not_a_text() { + engine::runtime::TaskRequest task; + task.options["language"] = "ja"; + check(!apply_transform_text_input(task), + "transform text: a language with no text reports false"); + check(!task.text_input.has_value(), + "transform text: a language alone must not route a separation request through text"); +} + +static void test_transform_text_preserves_other_inputs() { + engine::runtime::TaskRequest task; + engine::runtime::AudioBuffer audio; + audio.sample_rate = 44100; + audio.channels = 2; + audio.samples = {0.1f, 0.2f, 0.3f, 0.4f}; + task.audio_input = audio; + task.options["target_text"] = "keep the audio"; + check(apply_transform_text_input(task), "transform text: with audio present reports true"); + check(task.audio_input.has_value() && task.audio_input->samples.size() == 4 && + task.audio_input->sample_rate == 44100, + "transform text: the source audio is untouched"); +} + +int main() { + test_voice_is_reference_file(); + test_tts_shape(); + test_tts_plain(); + test_tts_named_preset(); + test_tts_reference_clip(); + test_tts_instructions(); + test_tts_empty_instructions_are_not_instructions(); + test_tts_empty_language_is_not_a_language(); + test_tts_clip_and_instructions(); + test_tts_language_and_params(); + test_sound_generation_minimal(); + test_sound_generation_full(); + test_transform_text_absent(); + test_transform_text_canonical_key(); + test_transform_text_alias_key(); + test_transform_text_canonical_wins(); + test_transform_text_empty_is_not_a_text(); + test_transform_text_empty_canonical_falls_through_to_alias(); + test_transform_text_language_rides_along(); + test_transform_language_alone_is_not_a_text(); + test_transform_text_preserves_other_inputs(); + + if (failures != 0) { + fprintf(stderr, "%d check(s) failed\n", failures); + return 1; + } + fprintf(stderr, "all checks passed\n"); + return 0; +} diff --git a/backend/cpp/audio-cpp/grpc-server.cpp b/backend/cpp/audio-cpp/grpc-server.cpp new file mode 100644 index 000000000..53fa173c5 --- /dev/null +++ b/backend/cpp/audio-cpp/grpc-server.cpp @@ -0,0 +1,2088 @@ +// audio.cpp LocalAI gRPC backend. +// +// Links 0xShug0/audio.cpp's engine_runtime through its public +// include/engine/framework/** headers only. Nothing under upstream's app/, +// src/ or tests/ is used: those are application internals, they are where +// upstream expects churn, and upstream is Apache-2.0 while LocalAI is MIT. +// +// This commit adds LoadModel/Free/Status plus the AudioTranscription, VAD, +// Diarize, AudioTransform, TTS, SoundGeneration, TTSStream, +// AudioTranscriptionStream and AudioTranscriptionLive RPCs, and explicit +// refusals for the five audio RPCs audio.cpp has no counterpart for: +// AudioEncode, AudioDecode, AudioTransformStream, AudioToAudioStream and +// VoiceEmbed. + +#include "backend.pb.h" +#include "backend.grpc.pb.h" + +#include "audio_io.h" +#include "audio_units.h" +#include "capability_routing.h" +#include "generation_request.h" +#include "inference_lane.h" +#include "live_watchdog.h" +#include "loaded_model.h" +#include "model_options.h" +#include "result_map.h" +#include "stem_selection.h" +#include "stream_delta.h" +#include "wav_header.h" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using grpc::Server; +using grpc::ServerBuilder; +using grpc::ServerContext; +// Do NOT alias grpc::Status as Status: the Status RPC method would shadow the +// type and break every other method that names it as a return type. +using GStatus = ::grpc::Status; + +namespace { + +// Set by the signal handler, read by the thread that owns the server. The whole +// point of the indirection is that the handler itself does nothing else: see +// signal_handler below. +std::atomic g_shutdown_requested{false}; + +// One model per process, matching every other LocalAI backend. The mutex guards +// the pointer itself, not the model: swapping or dropping it races with the +// handlers that read it, whereas the model's own concurrency is the inference +// lane's job. +// +// shared_ptr, not unique_ptr. An audio RPC runs for seconds and cannot hold +// g_model_mu for its duration, so it has to work from a reference taken under +// the lock and used after releasing it. Under a unique_ptr, a Free or a reload +// arriving mid-request destroys the model out from under that reference. Every +// handler instead takes a counted reference through snapshot_for(), so Free +// drops the global's reference and whichever request finishes last destroys the +// model. +std::mutex g_model_mu; +std::shared_ptr g_model; + +// The raw reference-taking primitive. Returns by value, so the caller owns a +// reference for as long as its local lives, and null when no model is loaded. +// Never return LoadedModel& from here: that is the shape that reintroduces the +// use-after-free. +// +// _unchecked because it performs NO request-level validation. There are exactly +// three legitimate classes of caller, and none of them is "the check was +// skipped": +// +// 1. Status, which takes a HealthMessage. It carries no ModelIdentity field. +// 2. An RPC whose request message has no ModelIdentity field either, so +// snapshot_for does not even instantiate for it. AudioTranscriptionLive is +// the routed one: TranscriptLiveRequest is a oneof of TranscriptLiveConfig +// and TranscriptLiveAudio and neither carries an identity. The four +// unroutable surfaces refuse_surface answers are the rest of the class: +// AudioEncodeRequest, AudioDecodeRequest, AudioTransformFrameRequest and +// AudioToAudioRequest carry no identity either. The stale-route hazard +// #10952 describes is therefore unguarded on these RPCs, and the fix has +// to be in backend.proto rather than here: nothing this process can read +// off the request says which model the caller thought it was reaching. +// When that field lands, the routed handlers move to snapshot_for and this +// clause shrinks to the refusals. It costs nothing on the four refusals, +// which answer the same UNIMPLEMENTED whatever model is loaded. +// 3. VoiceEmbed, which is the one refusal whose request DOES carry a +// ModelIdentity. It cannot use snapshot_for, because snapshot_for's +// no-model branch is FAILED_PRECONDITION "call LoadModel first" and that +// is the wrong instruction here: no model can make this backend serve +// VoiceEmbed, so the caller has to be told the reason instead. It +// therefore takes the reference here and runs check_model_identity itself, +// which keeps the ordering rule snapshot_for exists to enforce: identity +// before UNIMPLEMENTED, or a stale route gets "audio.cpp cannot embed +// speakers" when the model it actually asked for lives on a backend that +// can. +// +// Every OTHER RPC must call snapshot_for; the name is deliberately unpleasant so +// that reaching for it is a visible decision rather than an omission. +std::shared_ptr snapshot_unchecked() { + std::lock_guard lock(g_model_mu); + return g_model; +} + +// LocalAI's VADRequest carries raw floats and no sample rate at all. Every +// LocalAI VAD backend treats them as 16 kHz mono (backend/go/silero-vad/vad.go +// and backend/go/sherpa-onnx/backend.go both hardcode 16000) and core/backend +// feeds them from a 16 kHz pipeline, so the same assumption is made here rather +// than left implicit. If the proto ever grows a sample rate field, this is the +// constant to delete. +constexpr int kVadSampleRate = 16000; + +// The rate every file-fed speech route reads its input at, and the rate the +// spans that come back are therefore interpreted in. Passed to read_audio_file, +// which resamples only when the file differs. +// +// Two independent reasons, and the second is the one that is easy to miss: +// +// 1. Some families refuse anything else outright. silero_vad throws +// "Silero VAD 16k model only supports sample_rate=16000" and +// sortformer_diar throws "Sortformer diar currently requires 16 kHz input +// audio". That is what made a 44.1 kHz upload return INTERNAL with an +// engine-internal message instead of an answer. +// 2. The families that do NOT refuse still do not all express their result +// spans in the input's domain. nemotron_asr builds every word timestamp as +// token_frame * hop_length * subsampling_factor, which is its own 16 kHz +// feature domain whatever the input was; the handler then converts those +// spans with the buffer's rate. Feed it 44.1 kHz and every emitted +// timestamp is 2.76x too small, with a 200 and no diagnostic. Resampling +// the input to 16 kHz makes the buffer domain and the span domain the same +// one, which is the only reason the nanoseconds are right. +// +// The cost, stated plainly: vibevoice_asr resamples internally to 24 kHz, so a +// 48 kHz upload now travels 48 -> 16 -> 24 rather than 48 -> 24. That is not +// merely band limiting. resample_mono_linear is linear interpolation with no +// anti-alias filter (framework/audio/conversion.cpp), so the content above +// 8 kHz is ALIASED down on the way in, and the 16 -> 24 step aliases whatever +// the first step left. It is accepted because a silently wrong timestamp is +// worse than an aliased band, and because every other LocalAI ASR path already +// feeds 16 kHz. If the framework ever publishes a per-model preferred input +// rate, this constant is what should become that lookup. +constexpr int kSpeechSampleRate = 16000; + +// The rate AudioTransform reads BOTH its input and its reference clip at. Zero +// means "the file's own rate and its own channel count", i.e. this handler does +// no resampling and no downmix at all, and it is the only value that is correct +// for all four tasks this RPC can route to. +// +// Separation forces it. htdemucs and mel_band_roformer refuse anything but +// their own rate outright, from prepare(): "HTDemucs prepare() sample rate +// mismatch: expected N" (src/models/demucs/session.cpp) and the same shape in +// src/models/roformer/session.cpp. That N is the CHECKPOINT'S declared rate, +// read from the packaged config ("samplerate" in demucs/assets.cpp, +// "sample_rate" in roformer/assets.cpp), not a constant: it is 44100 for every +// published checkpoint of both, which is why the messages below say 44100, but +// a checkpoint declaring something else would demand that instead, and only +// passing the file through unchanged can satisfy either. Passing +// kSpeechSampleRate here would +// therefore turn every separation request into an INTERNAL, which is a loud +// failure. The downmix half is the quiet one: both models declare two channels +// and ACCEPT mono, by duplicating it across both, so a mono read would be taken +// and would merely delete the stereo image, which is the cue that separates a +// centred vocal from a wide mix. Verified rather than reasoned: a mono input +// comes back as mono stems, a stereo input as stereo ones. +// +// The three conversion tasks do not force it, and were checked one family at a +// time rather than assumed, because "it happens to work" and "it is documented +// to work" are different claims: +// +// seed_vc (vc, svc) seed_vc_prepare_audio_for_sample_rate takes the +// buffer's rate and channel count and resamples to its own mel +// rate: soxr when it is available, falling back to +// resample_mono_torchaudio_sinc_hann (seed_vc/audio_features.cpp). +// vevo2 (vc, s2s, svc) normalize_audio_to_24k_mono converts whatever +// it is given to 24 kHz mono before anything else runs, linearly. +// miocodec (vc, s2s) prepare_miocodec_mono_audio mixes down, then +// resamples to the model's rate with sinc-hann. +// chatterbox (vc) ChatterboxVcComponent::convert normalizes the source to +// 16 kHz and the reference to 24 kHz itself, linearly. +// +// So every one of them resamples internally, and passing the file through +// unchanged is better than folding it first. The reason is the DOUBLE +// conversion, not resampler quality: two of the four resample linearly, exactly +// as read_audio_file would. Their outputs run at 22.05 to 44.1 kHz, so +// pre-folding to 16 kHz mono with read_audio_file's LINEAR, unfiltered +// resampler would band-limit at 8 kHz and alias, and then the family would +// resample that damaged signal a second time on its way up. One conversion is +// the floor; this constant is what keeps it at one. +// +// The cost, stated plainly: a family that cannot handle the file's rate reports +// it from inside the engine as a plain runtime_error, which to_status maps to +// INTERNAL rather than INVALID_ARGUMENT. That is the accepted trade, since the +// only families that refuse are the separators and what they want is an +// ordinary music file at its own rate. +constexpr int kTransformSampleRate = 0; + +// The rate TTS reads its speaker reference clip at, and the rate +// SoundGeneration reads its `src` editing clip at. Zero, i.e. the file's own +// rate and its own channel count, no resample and no downmix. +// +// Settled from upstream rather than reasoned about, and upstream settles it +// twice over: +// +// 1. Upstream does exactly this itself. Both its CLI and its HTTP server load +// a voice reference with minitts::cli::read_audio_buffer (app/cli/request.cpp), +// which is read_wav_f32 and then {wav.sample_rate, wav.channels, +// wav.samples} verbatim. app/server/runtime.cpp's build_speech_request +// puts that buffer straight into voice.speaker->audio, and its +// audio_input path (line 518) does the same for the editing clip. Every +// family that consumes a reference has therefore only ever been tested +// against native-rate, native-channel input. +// 2. Every consuming family converts it itself, and most of them with a +// BETTER resampler than read_audio_file's. Checked one at a time: +// chatterbox to_mono_audio, then 24 kHz, then 16 kHz derived from +// the 24 kHz path (conditionals.cpp), matching Python's +// prepare_conditionals ordering. +// index_tts2 mixdown_interleaved_to_mono_average, then a soxr or +// torchaudio sinc-hann resample to its mel rate and to +// 16 kHz (audio_features.cpp, waveform_22k). +// higgs_audio_tts mixdown, then sinc-hann to 24 kHz and 16 kHz. +// irodori_tts mixdown, then sinc-hann (codec.cpp). +// voxcpm2 mixdown, then soxr-or-linear (audiovae.cpp). +// omnivoice mixdown, then sinc-hann (audio_tokenizer.cpp). +// qwen3_tts convert_interleaved_audio_to_mono_linear_resampled. +// For the SoundGeneration side, ace_step resamples the LEFT and RIGHT +// channels SEPARATELY (pre_dit.cpp), and stable_audio resamples per +// channel too, so a mono downmix here would destroy input those two are +// built to consume as stereo. +// +// So folding to 16 kHz mono first would band-limit at 8 kHz through +// read_audio_file's LINEAR, unfiltered resampler, alias what is above it, and +// then hand that damaged signal to a good resampler for a second conversion. +// One conversion is the floor, and passing the file through unchanged is what +// keeps it at one. Same argument as kTransformSampleRate, same value, kept as a +// separate constant because the routes are different and a future per-model +// preferred-rate lookup would have to answer them separately. +constexpr int kVoiceReferenceSampleRate = 0; + +// Parses ModelOptions.MainGPU into a device index. +// +// Not std::atoi: it returns 0 for anything unparseable, so "gpu1" or a device +// UUID would silently become device 0 and the model would load on the wrong +// device with no diagnostic anywhere. A refusal the operator can read beats a +// wrong answer they cannot see. +int parse_device_index(const std::string &value) { + size_t consumed = 0; + long parsed = 0; + try { + parsed = std::stol(value, &consumed); + } catch (const std::exception &) { + consumed = 0; + } + if (consumed != value.size() || parsed < 0 || + parsed > std::numeric_limits::max()) { + throw audiocpp_backend::ConfigError( + "audio-cpp: main_gpu must be a non-negative device index, got '" + + value + "'"); + } + return static_cast(parsed); +} + +// Mirrors pkg/grpc/server.go's checkModelIdentity, backend/cpp/llama-cpp, +// backend/cpp/ds4, backend/cpp/privacy-filter and +// backend/python/common/model_identity.py. +// +// Why every backend has to carry this: in distributed mode a worker can recycle +// a stopped backend's gRPC port for a different model's backend, and the +// controller's liveness-only health probe cannot tell a stale cached route from +// a live one. Only the backend knows which model it actually loaded, so only +// the backend can catch it (#10952). Without this, an audio-cpp process reached +// through a stale route answers with a DIFFERENT model's VAD or diarization +// result and a 200, which is a wrong answer nobody can see. +// +// Either side empty means "skip". The request side is empty for a controller +// that predates the field; the loaded side is empty when such a controller +// performed the load. Neither can judge the other, and a false rejection is +// worse than the miss it prevents. +// +// Templated over the request type because every guarded message exposes +// modelidentity(), and one body keeps the rule identical across RPCs rather +// than letting it drift per handler. +// +// The loaded identity is read off the LoadedModel rather than a separate +// global, which is where this differs from llama-cpp. A handler holding the +// model through snapshot_for() is then necessarily judging against the identity +// THAT model was loaded with, and a concurrent reload cannot swap one without +// the other. +template +GStatus check_model_identity(const audiocpp_backend::LoadedModel &model, + const Request *request) { + if (request == nullptr || request->modelidentity().empty()) { + return GStatus::OK; + } + const std::string &loaded = model.identity(); + if (loaded.empty() || loaded == request->modelidentity()) { + return GStatus::OK; + } + // NOT_FOUND plus this exact sentinel is the cross-language wire contract + // the router matches on (grpcerrors.ModelMismatchSentinel). The code alone + // is not enough, since NOT_FOUND is returned for unrelated reasons + // elsewhere, so the substring "model identity mismatch" must survive + // verbatim through any edit to this message. + return GStatus(grpc::StatusCode::NOT_FOUND, + "audio-cpp: model identity mismatch: loaded \"" + loaded + + "\", requested \"" + request->modelidentity() + "\""); +} + +// The ONE way an RPC handler should reach the model. Takes the counted +// reference, refuses when nothing is loaded, and runs the identity check, in +// that order. Returns null with `out` set to the status to return; on success +// returns the model and leaves `out` OK. +// +// This exists as a structural guarantee rather than a convenience. The identity +// check used to be two lines every handler had to remember, and nothing failed +// if a new handler forgot them: there is no C++ equivalent of +// pkg/grpc/model_identity_modalities_test.go, so the convention could only rot. +// Every handler already has to call something to obtain the model, so making +// the guarded call the shortest path means the default is correct and skipping +// it requires deliberately typing snapshot_unchecked. +// +// The order is forced: the identity being compared against belongs to the model +// this call is about to use, so it cannot precede taking the reference. And it +// must precede routing, or a mismatched request to a model that cannot serve +// the RPC leaks UNIMPLEMENTED instead of the NOT_FOUND the router matches on. +template +std::shared_ptr +snapshot_for(const Request *request, GStatus &out) { + auto model = snapshot_unchecked(); + if (model == nullptr) { + out = GStatus(grpc::StatusCode::FAILED_PRECONDITION, + "audio-cpp: no model is loaded; call LoadModel first"); + return nullptr; + } + out = check_model_identity(*model, request); + if (!out.ok()) { + return nullptr; + } + return model; +} + +// The single body behind the five refusals in the class below. +// +// UNIMPLEMENTED is the code by contract, not by taste: +// pkg/grpc/grpcerrors/errors.go degrades to an alternative path on +// UNIMPLEMENTED and on nothing else, so a caller that has a fallback keeps it. +// +// The loaded family is read only to ENRICH the message. The refusal itself is a +// property of the engine rather than of the model, which is why the no-model +// case is the same UNIMPLEMENTED and not the FAILED_PRECONDITION every routed +// handler returns: "call LoadModel first" would send the operator to do work +// that cannot help. The message says so in as many words. +// +// It does not read the stream, and the two bidirectional surfaces must not +// start: returning before the first Read means a client that opened the stream +// and is writing frames gets the status on its next operation instead of +// blocking on a response that would never come. +GStatus refuse_surface(audiocpp_backend::UnsupportedRpc rpc, + const audiocpp_backend::LoadedModel *model) { + const auto &surface = audiocpp_backend::unsupported_surface(rpc); + if (model == nullptr) { + return GStatus(grpc::StatusCode::UNIMPLEMENTED, + audiocpp_backend::unsupported_surface_message( + surface.rpc, surface.reason)); + } + return GStatus(grpc::StatusCode::UNIMPLEMENTED, + audiocpp_backend::unsupported_surface_message( + model->capabilities(), surface.rpc, surface.reason)); +} + +// Overload for the four surfaces whose request carries no ModelIdentity: there +// is nothing to check, so the reference is taken here. See case 2 at +// snapshot_unchecked. +GStatus refuse_surface(audiocpp_backend::UnsupportedRpc rpc) { + return refuse_surface(rpc, snapshot_unchecked().get()); +} + +// Builds the TaskRequest for a transcription-shaped RPC. `task` is the ROUTED +// audio.cpp task, not a guess from the request: it decides how +// TranscriptRequest.prompt is used, and routing has already decided it. +// +// The audio is taken by value and moved in. A long recording runs to tens of +// megabytes and the caller has no use for it afterwards; the previous shape, +// a const reference, copied it. +engine::runtime::TaskRequest +build_transcription_request(const backend::TranscriptRequest &request, + audiocpp_backend::Task task, + engine::runtime::AudioBuffer audio) { + engine::runtime::TaskRequest task_request; + task_request.audio_input = std::move(audio); + + if (task == audiocpp_backend::Task::Alignment) { + // For forced alignment the prompt IS the transcript to align, so it + // becomes the text input rather than a decoding hint. Set even when + // empty: an aligner given no text should say so itself rather than be + // handed an audio-only request it cannot describe. + engine::runtime::Transcript transcript; + transcript.text = request.prompt(); + transcript.language = request.language(); + task_request.text_input = transcript; + } else if (!request.prompt().empty()) { + // For ASR the prompt is decoding context, the whisper meaning. + task_request.options["prompt"] = request.prompt(); + } + + if (!request.language().empty()) { + task_request.options["language"] = request.language(); + } + if (request.translate()) { + task_request.options["translate"] = "true"; + } + if (request.temperature() > 0.0f) { + task_request.options["temperature"] = std::to_string(request.temperature()); + } + for (const auto &granularity : request.timestamp_granularities()) { + if (granularity == "word") { + // return_timestamps is the key upstream actually reads, and it is + // the one that does something: qwen3_asr defaults it to false + // (include/engine/models/qwen3_asr/types.h) and, when set, both + // runs its forced aligner and shortens its chunk window from 30 s + // to 15 s (src/models/qwen3_asr/session.cpp). Without it, asking + // for word granularity silently came back with no word timing. + // + // word_timestamps rides along as a forward-tolerant alias. NO + // family in the pinned upstream reads that key, a grep over src/ + // and include/ returns nothing, so it is sent only so that a family + // adopting the name later works with no change here. + task_request.options["return_timestamps"] = "true"; + task_request.options["word_timestamps"] = "true"; + } + } + // What lands and what does not. Families look their request options up by + // name (runtime::find_option) and ignore every key they do not know; the + // unknown-key refusals upstream does have are on SESSION options, which + // arrive at load time rather than here. So an unread key cannot turn a + // valid request into an error, but it is also not a feature, and the + // honest accounting against the pinned upstream is: + // + // language read, by nemotron_asr and vibevoice_asr among others. + // return_timestamps read, by qwen3_asr. + // prompt read by NO ASR family. Forwarded because it is the + // whisper meaning of the field and a family adopting + // it then works unchanged, not because it does + // anything today. + // translate read by nobody anywhere in upstream. + // temperature read only by TTS and voice conversion families, none + // of which this RPC can route to. + // + // Two request fields are deliberately not forwarded at all: + // + // threads thread count is a SessionOptions field fixed when the session + // was built, so a per-request value has nowhere to go. + // diarize this RPC routes to Asr or Alignment. A family that diarizes + // is reached through Diarize, which has its own handler and its + // own response shape, so forwarding this would imply that + // setting it turns speaker labels on here, and nothing would. + return task_request; +} + +// Duration of a possibly multi-channel buffer, in seconds. Frames, not floats: +// a stereo buffer holds two floats per position and would otherwise report +// twice its real length. +float audio_duration_seconds(const engine::runtime::AudioBuffer &audio) { + const std::int64_t frames = audiocpp_backend::interleaved_frame_count( + audio.samples.size(), audio.channels); + return audiocpp_backend::samples_to_seconds(frames, audio.sample_rate); +} + +// Refuses a request that routing sent to voice cloning with no speaker +// reference clip. Shared by TTS and TTSStream so the two RPCs cannot drift into +// refusing the same request differently. +// +// Refused FROM THE ROUTE, the same trick AudioTransform uses for params[stem]. +// Voice cloning without a clip is a request the family cannot answer, and every +// cloning family says so only from inside its own prepare(): chatterbox throws +// "Chatterbox prepare requires speaker reference audio", which to_status maps to +// INTERNAL and which names neither the RPC nor the field the caller has to set. +// chatterbox advertises clon and no tts at all, so before this every preset-only +// or voice-less request to it got that engine-internal message. +// +// The case it catches is the FALLBACK one, which is worth stating the right way +// round. A clip is what makes routing put VoiceCloning first, but VoiceCloning +// is also the last candidate a clip-less TTS request falls back to +// (task_candidates returns {Tts, VoiceCloning, VoiceDesign} when no clip and no +// instructions are supplied), so a family that advertises clon and no tts at +// all - chatterbox, which is what ships in the gallery - routes EVERY voice-less +// request here. A task:clon pin arrives here the same way. +// +// It cannot misfire on the legitimate case for the opposite reason: a request +// that did supply a clip has voice_is_file set and returns on the line above +// before anything is thrown. +// +// Deliberately NOT generalised to "every task whose family declares +// supports_speaker_reference". That flag lives on the engine's CapabilitySet, is +// not carried through this backend's Capabilities mirror, and would need its own +// answer to whether it is advisory or binding before routing could act on it. +// Tracked as a follow-up. +void refuse_cloning_without_a_clip(const audiocpp_backend::LoadedModel &model, + const audiocpp_backend::Route &route, + bool voice_is_file) { + if (route.task != audiocpp_backend::Task::VoiceCloning || voice_is_file) { + return; + } + throw audiocpp_backend::ConfigError( + "audio-cpp: family '" + model.family() + "' routes this request to " + + audiocpp_backend::task_name(route.task) + + ", which needs a speaker reference clip; set TTSRequest.voice to the " + "path of a WAV file (a voice that is not a file on disk is treated as a " + "named preset)"); +} + +// A live stream whose peer stopped sending without closing, cancelled by the +// idle watchdog. Its own type rather than a flag, so the read loop can unwind +// without the driver going on to finalize a decode nobody is waiting for. +// +// NOT handled by to_status: it is thrown and caught inside one handler, and +// giving it a clause there keeps to_status a statement about exceptions that +// cross unit boundaries. +class LiveIdleTimeout : public std::runtime_error { +public: + using std::runtime_error::runtime_error; +}; + +// Maps a thrown exception onto the gRPC status the client should see. +GStatus to_status(const std::exception &err) { + if (dynamic_cast(&err) != nullptr) { + return GStatus(grpc::StatusCode::INVALID_ARGUMENT, err.what()); + } + if (dynamic_cast(&err) != nullptr) { + return GStatus(grpc::StatusCode::UNIMPLEMENTED, err.what()); + } + // The lane was busy and this caller ran out of patience. UNAVAILABLE, not + // RESOURCE_EXHAUSTED: the request is fine and retrying later is the right + // response, which is what UNAVAILABLE tells a client. + if (dynamic_cast(&err) != nullptr) { + return GStatus(grpc::StatusCode::UNAVAILABLE, err.what()); + } + return GStatus(grpc::StatusCode::INTERNAL, err.what()); +} + +class AudioCppBackend final : public backend::Backend::Service { +public: + GStatus Health(ServerContext *, const backend::HealthMessage *, + backend::Reply *reply) override { + reply->set_message("OK"); + return GStatus::OK; + } + + GStatus Status(ServerContext *, const backend::HealthMessage *, + backend::StatusResponse *response) override { + // snapshot_unchecked is right here: HealthMessage carries no + // ModelIdentity, so there is nothing to check. See case 1 at the + // function. + response->set_state(snapshot_unchecked() + ? backend::StatusResponse::READY + : backend::StatusResponse::UNINITIALIZED); + return GStatus::OK; + } + + GStatus LoadModel(ServerContext *, const backend::ModelOptions *request, + backend::Result *result) override { + try { + std::vector entries(request->options().begin(), + request->options().end()); + auto parsed = audiocpp_backend::parse_model_options(entries); + if (!parsed.error.empty()) { + throw audiocpp_backend::ConfigError(parsed.error); + } + // ModelOptions.Threads is the model YAML's threads setting; the + // explicit threads: option wins when both are set. + if (parsed.options.threads == 0 && request->threads() > 0) { + parsed.options.threads = static_cast(request->threads()); + } + // MainGPU carries the device index for GPU backends, and is the + // fallback: an explicit device: option wins, matching threads above. + // device_set rather than `device == 0`, because 0 is a real device + // index and the value alone cannot say whether anyone chose it. + if (!parsed.options.device_set && !request->maingpu().empty()) { + parsed.options.device = parse_device_index(request->maingpu()); + } + + const std::string path = audiocpp_backend::resolve_model_path( + request->modelpath(), request->modelfile(), request->model()); + + // ModelOptions.Model is the UNTRANSLATED controller-side name, and + // is what every ModelIdentity field on a request is compared + // against. Passed at construction so the model and the identity it + // was loaded with are published together. + auto loaded = std::make_shared( + path, parsed.options, request->model()); + + // Read everything the reply and the log line need while this scope + // still owns the model. After the move, `loaded` is null and the + // global is only safe to touch under the mutex. + const std::string family = loaded->family(); + const std::string variant = loaded->variant(); + const std::string capabilities = + audiocpp_backend::describe_capabilities(loaded->capabilities()); + std::shared_ptr replaced; + { + std::lock_guard lock(g_model_mu); + replaced = std::move(g_model); + g_model = std::move(loaded); + } + // The previous model, if any, is dropped outside the lock, for the + // same reason Free does it: teardown must not hold up Status. + replaced.reset(); + + // Logged once at load time so an operator can see what the model + // can do without making a request. ModelMetadata is deliberately + // not implemented: its response message is chat-template oriented + // and has no field for family, variant, languages or capabilities, + // so there is nothing truthful to put in it. + std::cerr << "audio-cpp: loaded family '" << family << "' variant '" + << variant << "' capabilities: " << capabilities << "\n"; + + result->set_success(true); + result->set_message("loaded audio.cpp family " + family); + return GStatus::OK; + } catch (const std::exception &err) { + result->set_success(false); + result->set_message(err.what()); + // A failed load must also be a gRPC error, or the model loader's + // backend probe treats this backend as having accepted the model. + return to_status(err); + } + } + + GStatus Free(ServerContext *, const backend::HealthMessage *, + backend::Result *result) override { + // Drops the global's reference. Any handler still running holds its own + // from snapshot_for(), so the model is destroyed by whichever of them + // finishes last rather than underneath one of them. Destroying + // LoadedModel releases its sessions first, then the model. + std::shared_ptr released; + { + std::lock_guard lock(g_model_mu); + released = std::move(g_model); + } + // Released outside the lock: if this is the last reference, the model + // and its sessions are torn down here, and that must not block Status + // or a fresh LoadModel behind g_model_mu. + released.reset(); + result->set_success(true); + return GStatus::OK; + } + + GStatus AudioTranscription(ServerContext *, + const backend::TranscriptRequest *request, + backend::TranscriptResult *response) override { + try { + GStatus refusal = GStatus::OK; + const auto model = snapshot_for(request, refusal); + if (model == nullptr) { + return refusal; + } + + audiocpp_backend::RequestShape shape; + // has_prompt_text is what lets a family that can only align be + // reached through this RPC at all, so it is not decoration. + shape.has_prompt_text = !request->prompt().empty(); + shape.pinned_task = model->pinned_task(); + + // Capability refusal before the lane and before the file read, for + // the reasons spelled out in Diarize. + model->check_can_serve(audiocpp_backend::Rpc::AudioTranscription, shape); + + // TranscriptRequest.dst is the INPUT audio path, despite the field + // name. The HTTP layer materialises the upload to a temp file and + // passes the path; nothing is written back. + auto audio = audiocpp_backend::read_audio_file(request->dst(), + kSpeechSampleRate); + // Both read before the buffer is moved into the request below. The + // rate is the BUFFER's, which after read_audio_file is + // kSpeechSampleRate and not necessarily the file's, and it is the + // domain the result spans come back in. + const int sample_rate = audio.sample_rate; + const float duration = audio_duration_seconds(audio); + + audiocpp_backend::LaneEntry lane = model->acquire(0); + const auto session = model->session_for( + audiocpp_backend::Rpc::AudioTranscription, shape, lane); + + // session.task, not the request: for Asr the prompt is decoding + // context, for Alignment it is the transcript to align, and routing + // has already decided which of those this is. + const auto task_request = + build_transcription_request(*request, session.task, std::move(audio)); + const auto result = + audiocpp_backend::run_offline(session, task_request, lane); + + // text comes from result.text_output verbatim, never from the + // segments. See THE RULE in result_map.h. + audiocpp_backend::fill_transcript_result(result, sample_rate, duration, + response); + // eou stays false. It marks a decode that ended on the model's + // end-of-utterance token, which is a cache-aware STREAMING concept; + // an offline run over a whole file has no turn to yield. + return GStatus::OK; + } catch (const std::exception &err) { + return to_status(err); + } + } + + // Verifying this by hand: silero_vad is a SPEECH detector, so a synthetic + // stimulus does not exercise it. A 220 Hz sine, broadband noise and a + // harmonic buzz all return zero segments, correctly. Use real speech, which + // upstream bundles and which therefore needs no download: + // audio.cpp/assets/resources/sample_16k.wav (16 kHz mono, 14.07 s) + // Taking 1 s from offset 8000 samples (0.5 s in), padded with 1 s of + // silence either side, yields exactly one segment at start=0.9940 + // end=2.0460. A different excerpt of the same file shifts the start, + // because it changes where speech actually begins inside the window. + GStatus VAD(ServerContext *, const backend::VADRequest *request, + backend::VADResponse *response) override { + try { + // snapshot_for, and the result is held for the whole call. Free may + // arrive mid-request; it drops the global's reference only, so this + // one keeps the model alive until the handler returns. It also + // performs the not-loaded and identity checks, so a stale route is + // refused before any work happens. + GStatus refusal = GStatus::OK; + const auto model = snapshot_for(request, refusal); + if (model == nullptr) { + return refusal; + } + + audiocpp_backend::RequestShape shape; + // Without this the `task:` model option is dead: routing is + // otherwise derived from the RPC alone. + shape.pinned_task = model->pinned_task(); + + // Before the lane. A model that cannot do VAD at all answers + // immediately instead of queueing behind somebody else's run only + // to be refused; routing is pure and needs no lane. session_for + // routes again below and reaches the same answer. + model->check_can_serve(audiocpp_backend::Rpc::Vad, shape); + + engine::runtime::TaskRequest task; + task.audio_input = audiocpp_backend::buffer_from_mono( + std::vector(request->audio().begin(), request->audio().end()), + kVadSampleRate); + + // The lane is taken BEFORE session_for, not after. session_for + // reads and writes an unsynchronised session cache and prepare() + // mutates the session itself, so both belong inside the lane; see + // the note on session_for in loaded_model.h. Bound to a named local + // because LaneEntry is immovable and C++17 elides the return. + audiocpp_backend::LaneEntry lane = model->acquire(0); + const auto session = + model->session_for(audiocpp_backend::Rpc::Vad, shape, lane); + const auto result = + audiocpp_backend::run_offline(session, task, lane); + + // VADSegment.start/end are float SECONDS, not sample indices. + for (const auto &segment : result.speech_segments) { + auto *out = response->add_segments(); + out->set_start(audiocpp_backend::samples_to_seconds( + segment.span.start_sample, kVadSampleRate)); + out->set_end(audiocpp_backend::samples_to_seconds( + segment.span.end_sample, kVadSampleRate)); + } + return GStatus::OK; + } catch (const std::exception &err) { + return to_status(err); + } + } + + GStatus Diarize(ServerContext *, const backend::DiarizeRequest *request, + backend::DiarizeResponse *response) override { + try { + GStatus refusal = GStatus::OK; + const auto model = snapshot_for(request, refusal); + if (model == nullptr) { + return refusal; + } + + audiocpp_backend::RequestShape shape; + shape.pinned_task = model->pinned_task(); + + // Route, then read, then lane, in that order, and every step of it + // is deliberate. + // + // The capability refusal comes first because a family that cannot + // diarize at all should say so, not complain about the input file: + // on a VAD-only model, reading the audio first would surface + // "cannot read /tmp/x.wav" and send the operator hunting a file + // problem instead of a model choice. + // + // It also comes before the lane, which is the part that used to be + // impossible. Routing needs no lane, so the refusal no longer waits + // out somebody else's thirty second run to be told no, and neither + // does the file read. + model->check_can_serve(audiocpp_backend::Rpc::Diarize, shape); + + // DiarizeRequest.dst is the INPUT path, despite the field name: the + // HTTP layer materialises the upload to a temp file and passes the + // path here. Nothing is written back. + // + // Read at 16 kHz mono: sortformer refuses any other rate, and the + // HTTP layer copies the upload byte for byte, so whatever the user + // posted is what arrives. See kSpeechSampleRate. + auto audio = audiocpp_backend::read_audio_file(request->dst(), + kSpeechSampleRate); + const int sample_rate = audio.sample_rate; + // Read off the buffer before it is moved into the request below. + // Frames, not floats: a stereo input holds two floats per position + // and would otherwise report twice its real duration. + const std::int64_t frames = audiocpp_backend::interleaved_frame_count( + audio.samples.size(), audio.channels); + + audiocpp_backend::LaneEntry lane = model->acquire(0); + const auto session = + model->session_for(audiocpp_backend::Rpc::Diarize, shape, lane); + + engine::runtime::TaskRequest task; + // Moved, not copied: a long recording runs to tens of megabytes and + // this is its only owner. + task.audio_input = std::move(audio); + // READ THIS BEFORE WIRING THE FIRST DIARIZATION FAMILY. This + // forwarding is currently DEAD for sortformer, the only diarizer + // audio.cpp ships, and from the caller's side that is a silent + // wrong answer, not a soft hint: + // + // - backend.proto documents num_speakers as "exact speaker count + // if known (>0 forces)". A caller who asks for 2 and gets the + // model's 4 labels back with a 200 and no diagnostic has been + // given a wrong answer, so core/backend/diarization.go's + // "backends ignore what they don't act on" does not cover it. + // - sortformer takes its speaker count from the model package + // (assets.model_config.modules.num_speakers; the bundled + // variant is 4-speaker), and its postprocess config is parsed + // from runtime::SessionOptions, i.e. load time, not from + // TaskRequest.options at all. + // - worse, its unknown-option check only inspects keys prefixed + // "sortformer_diar.", so a bare "num_speakers" here is not even + // a recognised key: it is dropped without a trace. + // + // The keys are still forwarded, because a family that does read + // them then works with no change here. But the family that lands + // MUST either honour num_speakers or refuse it with + // INVALID_ARGUMENT. No refusal path is added now because there is + // no diarization family wired yet to test one against, and an + // untested refusal is its own hazard. + // + // Also dropped today, and worth revisiting at the same time: + // clustering_threshold, min_duration_on and min_duration_off. + // The two min_duration_* fields are NOT family-specific: they are + // generic post-filters over the returned turns (drop a turn shorter + // than min_duration_on, merge two turns of the same speaker + // separated by less than min_duration_off), so the handler could + // honour them in about six lines whatever the family does. + // + // The knobs that DO reach sortformer (speaker_threshold, + // speaker_min_frames, speaker_pad_frames, session_len_sec) are + // session options and arrive through the model YAML's + // `session.:` entries at load time, not per request. + if (request->num_speakers() > 0) { + task.options["num_speakers"] = + std::to_string(request->num_speakers()); + } + if (request->min_speakers() > 0) { + task.options["min_speakers"] = + std::to_string(request->min_speakers()); + } + if (request->max_speakers() > 0) { + task.options["max_speakers"] = + std::to_string(request->max_speakers()); + } + + const auto result = + audiocpp_backend::run_offline(session, task, lane); + + // Emitted verbatim, in the order the model produced them. NOT + // sorted, merged or de-overlapped: sortformer binarizes each speaker + // independently, so a turn nested inside another speaker's turn is + // correct output for overlapped speech. LocalAI is overlap-tolerant + // downstream (core/backend/diarization.go passes segments through + // and the RTTM renderer handles overlap), so smoothing here would + // destroy real information. + std::set speakers; + int id = 0; + for (const auto &turn : result.speaker_turns) { + auto *out = response->add_segments(); + out->set_id(id++); + // Seconds, like VADSegment. Only TranscriptSegment/Word are ns. + // + // Converted against the INPUT rate, which is correct and has + // been checked against upstream rather than assumed: every + // TimeSpan sortformer emits is built as + // llround(seconds * 16000.0) (postprocess.cpp). The read above + // guarantees the buffer is 16 kHz, so the span domain and the + // input domain are the same one, and this is not a place where + // a resampling family could silently scale every timestamp by a + // constant. + out->set_start(audiocpp_backend::samples_to_seconds( + turn.span.start_sample, sample_rate)); + out->set_end(audiocpp_backend::samples_to_seconds( + turn.span.end_sample, sample_rate)); + out->set_speaker(turn.speaker_id); + // text stays EMPTY, including when include_text is set. + // audio.cpp's SpeakerTurn carries a span and a speaker label + // only, and TaskResult has no per-segment text anywhere, so + // there is nothing truthful to put here. + speakers.insert(turn.speaker_id); + } + response->set_num_speakers(static_cast(speakers.size())); + response->set_duration( + audiocpp_backend::samples_to_seconds(frames, sample_rate)); + + // Only set when the family bundles transcription, which no + // diarization-only family does; kept because a combined family + // would fill it in and the field is documented as optional. + if (result.text_output.has_value()) { + response->set_language(result.text_output->language); + } + return GStatus::OK; + } catch (const std::exception &err) { + return to_status(err); + } + } + + // Serves the four tasks LocalAI's AudioTransform can represent: voice + // conversion, singing voice conversion, speech to speech and source + // separation. Which one a request becomes is routing's decision, not this + // handler's; see task_candidates for Rpc::AudioTransform, and note that svc + // is reachable only through an explicit task: pin because no request signal + // means "this input is singing". + // + // THE STEM COMPROMISE. AudioTransformResult carries a single dst, while + // htdemucs and mel_band_roformer produce four and two named stems from one + // run. Running inference once per stem would cost four full separations of + // the same file, so this runs ONCE, writes every stem to a sibling file + // .., and puts the selected one in dst. Those siblings + // are real files in the caller's output directory that LocalAI's caller + // does not know about and will not clean up: a documented trade, not an + // oversight, and the reason params["stem"] exists to say which one dst gets. + GStatus AudioTransform(ServerContext *, + const backend::AudioTransformRequest *request, + backend::AudioTransformResult *response) override { + try { + GStatus refusal = GStatus::OK; + const auto model = snapshot_for(request, refusal); + if (model == nullptr) { + return refusal; + } + + const bool has_reference = !request->reference_path().empty(); + audiocpp_backend::RequestShape shape; + // Unused by AudioTransform's own routing today, since none of its + // four candidate tasks is chosen by the presence of a reference + // clip. Set anyway, because leaving a shape field stale is how a + // later routing rule silently reads the wrong thing. + shape.has_voice_reference = has_reference; + shape.pinned_task = model->pinned_task(); + + // First, before the lane, before the file reads, and before the + // argument check below. A family that cannot transform at all is + // answering a question about ITSELF, so it must not first queue + // behind somebody else's thirty second run, and it must not blame + // the caller's paths for a decision that had nothing to do with + // them. Same ordering as Diarize and AudioTranscription. + const audiocpp_backend::Route route = + model->check_can_serve(audiocpp_backend::Rpc::AudioTransform, shape); + + if (request->audio_path().empty() || request->dst().empty()) { + throw audiocpp_backend::ConfigError( + "audio-cpp: AudioTransform needs both audio_path and dst"); + } + + engine::runtime::TaskRequest task; + std::string requested_stem; + for (const auto ¶m : request->params()) { + if (param.first == "stem") { + // CONSUMED here and deliberately NOT forwarded into + // task.options: it selects which output the caller + // receives, it does not tune the model. Forwarding it would + // put a key no family reads into every request and imply + // the model had been asked to produce only that stem. + requested_stem = param.second; + continue; + } + task.options[param.first] = param.second; + } + + // AFTER the loop, so it reads exactly what the caller sent. This is + // what makes a text-conditioned route reachable through an RPC whose + // message has no text field; see apply_transform_text_input's header + // for why vevo2's speech-to-speech route is unreachable without it. + // A request carrying no text key is untouched, so separation and + // voice conversion pay nothing for this. + audiocpp_backend::apply_transform_text_input(task); + + // Refused from the ROUTE, before the file reads and before the run. + // Only source separation produces named stems, and the route says + // whether this is separation without running anything: the identical + // refusal below, taken from the empty named_audio_outputs list, + // cannot fire until a full conversion has been paid for (measured at + // 1.6 s on miocodec, far worse on seed_vc or vevo2). The one below + // stays as the backstop for a separation-routed family that returns + // no stems anyway. + if (!requested_stem.empty() && + route.task != audiocpp_backend::Task::SourceSeparation) { + throw audiocpp_backend::ConfigError( + "audio-cpp: family '" + model->family() + "' routes this " + + "request to " + audiocpp_backend::task_name(route.task) + + ", which produces a single output with no named stems, so " + "params[stem]='" + requested_stem + "' cannot be honoured"); + } + + // Native rate, native channels, for both files. See + // kTransformSampleRate: separation is destroyed by a downmix and + // every conversion family resamples internally anyway. + task.audio_input = audiocpp_backend::read_audio_file( + request->audio_path(), kTransformSampleRate); + if (has_reference) { + engine::runtime::VoiceReference reference; + reference.audio = audiocpp_backend::read_audio_file( + request->reference_path(), kTransformSampleRate); + engine::runtime::VoiceCondition condition; + condition.speaker = std::move(reference); + task.voice = std::move(condition); + } + + // Named local: LaneEntry is immovable and must span both + // session_for and run_offline, which is the constraint the proof of + // holding parameter exists to enforce. + audiocpp_backend::LaneEntry lane = model->acquire(0); + const auto session = + model->session_for(audiocpp_backend::Rpc::AudioTransform, shape, lane); + const auto result = audiocpp_backend::run_offline(session, task, lane); + + const engine::runtime::AudioBuffer *chosen = nullptr; + if (!result.named_audio_outputs.empty()) { + std::vector names; + names.reserve(result.named_audio_outputs.size()); + for (const auto &named : result.named_audio_outputs) { + names.push_back(named.id); + } + // Selected BEFORE the first write, not after the loop. An + // unknown stem name is a refused request, and a refused request + // must not leave four files in the caller's output directory + // that it then reports nothing about. + // + // The guarantee is exactly that and no more: a request REFUSED + // ON ITS ARGUMENTS writes nothing. A write that FAILS partway + // through the loop below, on a full disk say, still leaves the + // siblings written before it, with no dst. Nothing is rolled + // back, because deleting files after a disk error is its own + // way to lose data, and the caller sees the failure. + const auto choice = + audiocpp_backend::select_named_output(names, requested_stem); + if (!choice.error.empty()) { + throw audiocpp_backend::ConfigError(choice.error); + } + + for (size_t i = 0; i < result.named_audio_outputs.size(); ++i) { + const std::string sibling = + audiocpp_backend::sibling_stem_path(request->dst(), names[i]); + audiocpp_backend::write_audio_file( + sibling, result.named_audio_outputs[i].audio); + // Named in the response, in the model's own order. Without + // this the siblings are files nobody can find, and a caller + // that wants the drums as well as the vocals has to run the + // whole separation again per stem, which is the cost the + // single run exists to avoid. + auto *stem = response->add_stems(); + stem->set_name(names[i]); + stem->set_dst(sibling); + } + chosen = &result.named_audio_outputs[static_cast(choice.index)] + .audio; + // dst last, so it is the file that exists only once every stem + // beside it does. Its content duplicates the selected sibling + // on purpose: the caller reads dst, an operator reads the + // siblings, and neither should have to know about the other. + audiocpp_backend::write_audio_file(request->dst(), *chosen); + } else if (result.audio_output.has_value()) { + if (!requested_stem.empty()) { + // A conversion family produces one unnamed output, so there + // is no stem to choose. Refused rather than ignored: a + // caller who asked for "vocals" and received the whole + // converted signal has been answered with something else. + throw audiocpp_backend::ConfigError( + "audio-cpp: family '" + model->family() + + "' produces a single output with no named stems, so " + "params[stem]='" + requested_stem + "' cannot be honoured"); + } + chosen = &*result.audio_output; + audiocpp_backend::write_audio_file(request->dst(), *chosen); + } else { + // Reached only if a family routed successfully and then + // returned no audio at all, which is a broken family rather + // than a wrong request. CapabilityError so it reads as "this + // model does not do that" instead of as an internal fault. + throw audiocpp_backend::CapabilityError( + "audio-cpp: family '" + model->family() + + "' produced no audio for the AudioTransform RPC"); + } + + response->set_dst(request->dst()); + response->set_sample_rate(chosen->sample_rate); + // FRAMES, not floats. Separation output is stereo, so reporting + // samples.size() would tell the caller a 3 second stem is 6 seconds + // long. + response->set_samples(static_cast( + audiocpp_backend::interleaved_frame_count(chosen->samples.size(), + chosen->channels))); + response->set_reference_provided(has_reference); + return GStatus::OK; + } catch (const std::exception &err) { + return to_status(err); + } + } + + // The first RPC that GENERATES audio from text, together with + // SoundGeneration below. AudioTransform already wrote files, but it + // transformed audio it was given; these two have no required input audio at + // all, which is why both carry a Result with a success flag rather than a + // response message describing what was found. + // + // TTSRequest.voice is overloaded across LocalAI backends, some reading it as + // a named preset and some as a path to a clip. The rule here is decided from + // the filesystem: an existing regular file is a speaker reference and + // routing prefers VoiceCloning, anything else is a preset. Instructions with + // no clip prefer VoiceDesign, and a clip outranks instructions when both are + // set. None of that is re-derived here; it is capability_routing's + // task_candidates, and this handler's job is only to describe the request + // truthfully in the RequestShape. + GStatus TTS(ServerContext *, const backend::TTSRequest *request, + backend::Result *result) override { + try { + GStatus refusal = GStatus::OK; + const auto model = snapshot_for(request, refusal); + if (model == nullptr) { + // Answered with the STATUS ALONE, leaving Result at its default. + // core/backend/tts.go checks the transport error before it looks + // at res.Success and returns on it, so the Result is never read + // on this path; and the router matches a stale route on the + // NOT_FOUND code plus the sentinel in the status message, which + // is where check_model_identity already put it. + return refusal; + } + + // One description of the request, shared with TTSStream, so the two + // handlers cannot describe the same request differently. + // pinned_task stays here on purpose: it comes off the model rather + // than the request, and it is the field a new handler forgets. + audiocpp_backend::RequestShape shape = + audiocpp_backend::build_tts_shape(*request); + const bool voice_is_file = shape.has_voice_reference; + shape.pinned_task = model->pinned_task(); + + // FIRST, before the lane and before any file read. A family that + // cannot synthesise at all is answering a question about itself, so + // it must not queue behind somebody else's run to be told no, and it + // must not blame the caller's reference clip for a decision that had + // nothing to do with it. Same ordering as Diarize and + // AudioTransform. + const audiocpp_backend::Route route = + model->check_can_serve(audiocpp_backend::Rpc::Tts, shape); + + refuse_cloning_without_a_clip(*model, route, voice_is_file); + + if (request->dst().empty()) { + throw audiocpp_backend::ConfigError( + "audio-cpp: TTS needs a dst output path"); + } + + std::optional reference; + if (voice_is_file) { + // Native rate, native channels: see kVoiceReferenceSampleRate. + reference = audiocpp_backend::read_audio_file( + request->voice(), kVoiceReferenceSampleRate); + } + // Read before the lane, like every other file-fed handler here, so + // a slow or large clip is not decoded while holding it. + const auto task = + audiocpp_backend::build_tts_request(*request, std::move(reference)); + + // Named local: LaneEntry is immovable and has to span both + // session_for and run_offline. + audiocpp_backend::LaneEntry lane = model->acquire(0); + const auto session = + model->session_for(audiocpp_backend::Rpc::Tts, shape, lane); + const auto task_result = + audiocpp_backend::run_offline(session, task, lane); + + if (!task_result.audio_output.has_value()) { + // A family that routed successfully and then produced no audio + // is a broken family rather than a wrong request, but from the + // caller's side the actionable fact is that this model does not + // do this, so CapabilityError. Same judgement as AudioTransform. + throw audiocpp_backend::CapabilityError( + "audio-cpp: family '" + model->family() + + "' produced no audio for the TTS RPC"); + } + // Throws a plain runtime_error, i.e. INTERNAL, on a write failure: + // dst is LocalAI's own generated-content path and not anything the + // caller named, so a full disk there is a server fault and is worth + // retrying. Only an empty dst is INVALID_ARGUMENT, and that is + // already refused above with a message naming the RPC. + audiocpp_backend::write_audio_file(request->dst(), + *task_result.audio_output); + result->set_success(true); + // The path, because that is what core/backend/tts.go's caller reads + // back and what every other LocalAI TTS backend puts here. + result->set_message(request->dst()); + return GStatus::OK; + } catch (const std::exception &err) { + // DEFENSIVE, not load-bearing, and worth saying so plainly. gRPC + // discards the response message entirely when the status is not OK, + // so a client that checks the status, which core/backend/tts.go + // does before it ever looks at res.Success, receives a nil Result + // and never sees these two fields. They are set for a client that + // ignores the status, and because a half-filled Result is a worse + // thing to leave behind than a filled one. + result->set_success(false); + result->set_message(err.what()); + return to_status(err); + } + } + + GStatus SoundGeneration(ServerContext *, + const backend::SoundGenerationRequest *request, + backend::Result *result) override { + try { + GStatus refusal = GStatus::OK; + const auto model = snapshot_for(request, refusal); + if (model == nullptr) { + return refusal; + } + + audiocpp_backend::RequestShape shape; + // No request signal chooses between generation tasks: this RPC has + // exactly one candidate, AudioGeneration. pinned_task is still + // copied, because without it the model's `task:` option is dead + // here, which is how a gen family pinned to something else would + // silently be routed to generation anyway. + shape.pinned_task = model->pinned_task(); + + model->check_can_serve(audiocpp_backend::Rpc::SoundGeneration, shape); + + if (request->dst().empty()) { + throw audiocpp_backend::ConfigError( + "audio-cpp: SoundGeneration needs a dst output path"); + } + + // DO NOT WIRE src INTO THE HTTP LAYER UNTIL THE PIN IS BUMPED PAST + // THE FIX. Setting src on a stable_audio model CORRUPTS THE HEAP AND + // ABORTS THE PROCESS in the pinned upstream: "free(): invalid size" + // / "munmap_chunk(): invalid pointer", SIGABRT, backend gone. It is + // upstream's, not this handler's, and it was attributed rather than + // assumed: upstream's own audiocpp_cli, built from this same + // checkout, aborts identically with --audio at exit 134, at both + // 44.1 kHz stereo and 24 kHz mono, and completes cleanly with no + // --audio at all. It is therefore neither caused nor worsened by + // reading the clip at its native rate below. + // + // The only thing keeping that off the network today is an OMISSION: + // core/http/endpoints/elevenlabs/soundgeneration.go passes nil for + // sourceFile, and schema.ElevenLabsSoundGenerationRequest has no + // field for it, so the sole caller that can set src is + // core/cli/soundgeneration.go. Adding the field to that schema turns + // a local CLI crash into a remotely reachable heap corruption with + // fully attacker-influenced input. Nobody reading that Go schema + // would know why the field is missing, which is why this is written + // here as well as in the task report. + // + // No family blocklist is applied: ace_step's editing routes + // legitimately need src, and a family-specific guard here would rot. + std::optional source; + if (request->has_src() && !request->src().empty()) { + // Native rate and channels. ace_step resamples left and right + // separately and stable_audio resamples per channel, so a + // downmix here would delete the stereo image of the very clip + // they are editing. + source = audiocpp_backend::read_audio_file( + request->src(), kVoiceReferenceSampleRate); + } + const auto task = audiocpp_backend::build_sound_generation_request( + *request, std::move(source)); + + audiocpp_backend::LaneEntry lane = model->acquire(0); + const auto session = model->session_for( + audiocpp_backend::Rpc::SoundGeneration, shape, lane); + const auto task_result = + audiocpp_backend::run_offline(session, task, lane); + + if (!task_result.audio_output.has_value()) { + throw audiocpp_backend::CapabilityError( + "audio-cpp: family '" + model->family() + + "' produced no audio for the SoundGeneration RPC"); + } + audiocpp_backend::write_audio_file(request->dst(), + *task_result.audio_output); + result->set_success(true); + result->set_message(request->dst()); + return GStatus::OK; + } catch (const std::exception &err) { + result->set_success(false); + result->set_message(err.what()); + return to_status(err); + } + } + + // The streaming counterpart of TTS, and the first RPC here that writes to + // the wire while the model is still generating. + // + // THE WIRE CONTRACT, which is the thing most likely to go wrong. + // pkg/grpc/server.go's TTSStream wrapper puts every chunk in Reply.audio, + // and core/backend/tts.go's ModelTTSStream forwards those bytes to the HTTP + // response verbatim. There is no framing and no format negotiation, so THE + // FIRST CHUNK MUST BE A WAV HEADER or the client receives raw PCM it has no + // way to interpret; browsers simply refuse the stream. Because the total + // length is unknown while generating, both size fields carry 0xFFFFFFFF, + // which is the convention backend/go/vibevoice-cpp established. + // + // ModelTTSStream will synthesise a header of its own, but ONLY when the + // first Reply carries a non-empty `message` holding JSON with a sample_rate. + // Nothing here ever sets Reply.message, so that branch never fires and there + // is exactly one header on the wire. Setting `message` on this RPC without + // deleting the header below would put a second header 44 bytes into the PCM. + // + // There is no dst: streaming TTS writes to the response, not to a file, and + // core/backend/tts.go sends an empty dst on purpose. That is the one place + // this RPC deliberately diverges from TTS. + GStatus TTSStream(ServerContext *, const backend::TTSRequest *request, + grpc::ServerWriter *writer) override { + try { + GStatus refusal = GStatus::OK; + const auto model = snapshot_for(request, refusal); + if (model == nullptr) { + return refusal; + } + + // The SAME shape builder TTS uses, so the two RPCs cannot describe + // one request differently. pinned_task stays at the call site: it + // comes off the model rather than the request. + audiocpp_backend::RequestShape shape = + audiocpp_backend::build_tts_shape(*request); + const bool voice_is_file = shape.has_voice_reference; + shape.pinned_task = model->pinned_task(); + + // FIRST, before the lane and before any file read, exactly as in + // TTS. Note that this RPC's mode candidates are streaming ONLY: a + // family that can synthesise but cannot stream is refused here + // rather than quietly served a whole buffer at the end, because a + // caller that asked to stream is asking for time to first audio. + const audiocpp_backend::Route route = + model->check_can_serve(audiocpp_backend::Rpc::TtsStream, shape); + refuse_cloning_without_a_clip(*model, route, voice_is_file); + + std::optional reference; + if (voice_is_file) { + // Native rate, native channels: see kVoiceReferenceSampleRate. + reference = audiocpp_backend::read_audio_file( + request->voice(), kVoiceReferenceSampleRate); + } + const auto task = + audiocpp_backend::build_tts_request(*request, std::move(reference)); + + audiocpp_backend::LaneEntry lane = model->acquire(0); + const auto session = + model->session_for(audiocpp_backend::Rpc::TtsStream, shape, lane); + + bool header_sent = false; + bool client_gone = false; + int stream_rate = 0; + int stream_channels = 0; + const auto write_audio = + [&](const engine::runtime::AudioBuffer &audio) { + if (client_gone || audio.samples.empty()) { + return; + } + const int channels = audio.channels > 0 ? audio.channels : 1; + if (!header_sent) { + backend::Reply head; + head.set_audio(audiocpp_backend::streaming_wav_header( + audio.sample_rate, channels)); + if (!writer->Write(head)) { + client_gone = true; + return; + } + header_sent = true; + stream_rate = audio.sample_rate; + stream_channels = channels; + } else if (audio.sample_rate != stream_rate || + channels != stream_channels) { + // The header already went out declaring the first + // chunk's format, and it cannot be taken back. Every + // later byte would be decoded at the wrong rate or the + // wrong interleaving, which plays as a speed or channel + // fault nobody can see in a 200. Fail loudly instead. + throw std::runtime_error( + "audio-cpp: family '" + model->family() + + "' changed its output format mid-stream (" + + std::to_string(stream_rate) + " Hz x" + + std::to_string(stream_channels) + " to " + + std::to_string(audio.sample_rate) + " Hz x" + + std::to_string(channels) + ")"); + } + backend::Reply reply; + reply.set_audio(audiocpp_backend::f32_to_s16le(audio.samples)); + if (!writer->Write(reply)) { + client_gone = true; + } + }; + + // BOTH fields, and named_audio_outputs is the one that carries the + // audio in practice. All three streaming TTS families put their + // chunks there ("chunk_0", "chunk_1", ...) and leave audio_output + // empty until the very end: supertonic and omnivoice build the event + // in next_stream_event, voxcpm2 replays the ones its start_stream + // produced. Reading only audio_output, which is the obvious field, + // yields a stream with no audio in it at all. + const auto emit = [&](const engine::runtime::StreamEvent &event) { + if (event.audio_output.has_value()) { + write_audio(*event.audio_output); + } + for (const auto &named : event.named_audio_outputs) { + write_audio(named.audio); + } + }; + + const auto result = + audiocpp_backend::run_streaming_pull(session, task, emit, lane); + + // The finish_stream result is the session's own MERGED WHOLE, not a + // tail the pull loop missed: supertonic returns its accumulated + // buffer, omnivoice post-processes the merge, voxcpm2 hands back the + // stored result. Emitting it after the chunks would send the entire + // utterance twice. It is therefore used only when the family + // streamed nothing, which is what a family that held everything back + // to finalize looks like. + // + // The cost of that choice, stated plainly: what a client hears is + // the concatenation of the chunks, and for omnivoice that is not + // byte-identical to what the unary TTS RPC returns, because its + // postprocessor runs over the merged buffer at the end. Audio + // already on the wire cannot be post-processed, so any streaming + // implementation has to accept this. + // + // This guard also depends on the pull loop having DRAINED the + // session. It does today: no family sets is_final on a pulled + // event, so the loop runs to nullopt. If one ever does, the loop + // breaks early, and omnivoice's finish_stream drains and merges the + // chunks that were never emitted (session.cpp:576) into the result + // this branch then discards, silently truncating the audio. Whoever + // teaches a family to end a pull stream early has to revisit the + // pair together. + if (!header_sent) { + if (result.audio_output.has_value()) { + write_audio(*result.audio_output); + } else { + for (const auto &named : result.named_audio_outputs) { + write_audio(named.audio); + } + } + } + + if (!header_sent && !client_gone) { + // Routed successfully and produced nothing. Same judgement as + // TTS: from the caller's side the actionable fact is that this + // model does not do this. + throw audiocpp_backend::CapabilityError( + "audio-cpp: family '" + model->family() + + "' produced no audio for the TTSStream RPC"); + } + // client_gone is NOT an error: the client hung up, gRPC already + // knows, and there is nobody left to tell. + return GStatus::OK; + } catch (const std::exception &err) { + // No Result message to fill in on this RPC. A status mid-stream is + // what the client sees, and gRPC delivers it after whatever chunks + // already went out. + return to_status(err); + } + } + + // Server-streaming transcription: zero or more TranscriptStreamResponse + // messages carrying `delta`, then exactly one carrying `final_result`. + // + // THE DELTA CONTRACT. core/backend/transcript.go documents Delta as "an + // incremental text fragment" and the HTTP layer appends them, so a client + // that concatenates every delta must end up with the final text and must + // never see a prefix repeated. The families do not agree on what they report + // (nemotron_asr and vibevoice_asr send fragments, voxtral_realtime sends the + // whole hypothesis and repeats the last event of each batch), and a delta + // must additionally be valid UTF-8 or the Go client cannot unmarshal it at + // all, so the reconciliation lives in TranscriptDeltaTracker rather than + // here. See stream_delta.h. + // + // THE OFFLINE FALLBACK. mode_candidates lets this RPC fall back to an + // offline route, so a family with no streaming ASR still answers rather than + // returning UNIMPLEMENTED. It then runs once and the reconciliation below + // emits the whole transcript as a single delta: the same message sequence a + // streaming family produces, with one delta instead of many. That is a + // truthful degradation, and it needs no branch of its own, because a tracker + // that observed nothing reconciles to exactly one fragment. + GStatus AudioTranscriptionStream( + ServerContext *, const backend::TranscriptRequest *request, + grpc::ServerWriter *writer) override { + try { + GStatus refusal = GStatus::OK; + const auto model = snapshot_for(request, refusal); + if (model == nullptr) { + return refusal; + } + + audiocpp_backend::RequestShape shape; + shape.has_prompt_text = !request->prompt().empty(); + shape.pinned_task = model->pinned_task(); + + // Before the lane and before the file read, for the reasons spelled + // out in Diarize. + model->check_can_serve(audiocpp_backend::Rpc::AudioTranscriptionStream, + shape); + + // TranscriptRequest.dst is the INPUT audio path, despite the field + // name; nothing is written back. Read at 16 kHz mono for the reasons + // in kSpeechSampleRate, which apply identically here: the streaming + // sessions build their spans in the same feature domain their + // offline halves do. + auto audio = audiocpp_backend::read_audio_file(request->dst(), + kSpeechSampleRate); + const int sample_rate = audio.sample_rate; + const float duration = audio_duration_seconds(audio); + + audiocpp_backend::LaneEntry lane = model->acquire(0); + const auto session = model->session_for( + audiocpp_backend::Rpc::AudioTranscriptionStream, shape, lane); + + // session.task, not the request: for Asr the prompt is decoding + // context, for Alignment it is the transcript to align. + const auto task = build_transcription_request(*request, session.task, + std::move(audio)); + if (!task.audio_input.has_value()) { + // Unreachable through build_transcription_request, which always + // sets it. Checked because the alternative is dereferencing an + // empty optional below. + throw std::runtime_error( + "audio-cpp: transcription request carries no audio"); + } + + audiocpp_backend::TranscriptDeltaTracker tracker; + bool client_gone = false; + const auto write_delta = [&](const std::string &fragment) { + if (fragment.empty() || client_gone) { + return; + } + backend::TranscriptStreamResponse response; + response.set_delta(fragment); + if (!writer->Write(response)) { + client_gone = true; + } + }; + const auto emit = [&](const engine::runtime::StreamEvent &event) { + if (!event.partial_text.has_value()) { + return; + } + write_delta(tracker.observe(event.partial_text->text)); + }; + + engine::runtime::TaskResult result; + if (session.mode == audiocpp_backend::Mode::Streaming) { + // The buffer inside the request is the chunk source, so the + // audio is held once rather than copied for the feed. + result = audiocpp_backend::run_streaming_audio( + session, task, *task.audio_input, emit, lane); + } else { + result = audiocpp_backend::run_offline(session, task, lane); + } + + // The reconciliation, and the only place the offline fallback needs + // to be thought about: with no partials observed this IS the single + // delta carrying the whole transcript. + if (result.text_output.has_value()) { + write_delta(tracker.reconcile(result.text_output->text)); + } + + backend::TranscriptStreamResponse final_response; + // sample_rate is the BUFFER's, which after read_audio_file is + // kSpeechSampleRate and not necessarily the file's; it is the domain + // the result spans come back in. + audiocpp_backend::fill_transcript_result( + result, sample_rate, duration, + final_response.mutable_final_result()); + if (!client_gone) { + writer->Write(final_response); + } + return GStatus::OK; + } catch (const std::exception &err) { + return to_status(err); + } + } + + // The one BIDIRECTIONAL stream this backend serves: live microphone ASR. + // The client sends a TranscriptLiveConfig, then TranscriptLiveAudio frames; + // this answers with a ready ack, then deltas and words as the audio + // arrives, then one message carrying final_result once the read side + // closes. + // + // THERE IS NO OFFLINE FALLBACK, unlike AudioTranscriptionStream. Live + // transcription has to consume audio incrementally, so a family with no + // streaming ASR is refused rather than served a batch run at the end; the + // Streaming-only mode_candidates list for this RPC is where that is + // expressed, and this handler needs no branch for it. + // + // THE READY ACK IS A CONTRACT, not a courtesy. core/backend/transcript_live.go + // sends the config and then BLOCKS on Recv, treating an Unimplemented there + // as "this backend cannot do live transcription, degrade to the file path" + // and a first response carrying data as a broken backend. So routing must be + // refused before the ack, and the ack must go out before the first audio + // frame is read, or the client never sends one and both sides wait. + // + // eou AND eob STAY FALSE. TranscriptLiveResponse carries them for + // cache-aware models that emit end-of-utterance and end-of-backchannel + // tokens; audio.cpp's StreamEvent has no equivalent signal, and inferring + // one from silence would be a guess wearing a protocol flag's clothes. A + // client uses eou to decide the speaker yielded the turn, so a wrong one + // cuts people off mid-sentence. False is not a degradation here, it is the + // truth: this backend does not know. + GStatus AudioTranscriptionLive( + ServerContext *context, + grpc::ServerReaderWriter *stream) override { + try { + // snapshot_unchecked, which every other model-touching handler is + // forbidden to call. TranscriptLiveRequest carries NO ModelIdentity + // field, in either arm of its oneof, so snapshot_for does not even + // instantiate for it: there is nothing to compare. See case 2 at + // snapshot_unchecked, including why the fix is a proto change. + const auto model = snapshot_unchecked(); + if (model == nullptr) { + return GStatus(grpc::StatusCode::FAILED_PRECONDITION, + "audio-cpp: no model is loaded; call LoadModel first"); + } + + backend::TranscriptLiveRequest incoming; + if (!stream->Read(&incoming)) { + // Opened and closed with nothing said. A clean exit, not an + // error: there is nothing to transcribe and nobody to tell. + return GStatus::OK; + } + if (!incoming.has_config()) { + return GStatus(grpc::StatusCode::INVALID_ARGUMENT, + "audio-cpp: the first AudioTranscriptionLive " + "message must carry a Config"); + } + const backend::TranscriptLiveConfig config = incoming.config(); + + audiocpp_backend::RequestShape shape; + // No request signal chooses the task here: TranscriptLiveConfig has + // no prompt field, so has_prompt_text stays false and the only + // candidate is Asr. pinned_task is still copied, because without it + // the model's `task:` option is dead on this RPC. + shape.pinned_task = model->pinned_task(); + + // FIRST, before the lane and BEFORE the rate check below, like every + // other handler here. A family that cannot stream ASR is answering a + // question about itself, and that answer outranks any complaint + // about the request: pkg/grpc/grpcerrors/errors.go degrades to the + // file path on UNIMPLEMENTED and on nothing else, so a live-incapable + // model asked at a wrong rate must not answer INVALID_ARGUMENT and + // cost the caller its fallback. Unreachable from LocalAI today, + // which always sends 16000, and wrong on its own terms regardless. + model->check_can_serve(audiocpp_backend::Rpc::AudioTranscriptionLive, + shape); + + // 16 kHz mono or nothing, and this is a REFUSAL rather than a + // resample for the reason spelled out at kSpeechSampleRate: the + // streaming families build their spans in their own 16 kHz feature + // domain whatever the input rate was, so honouring an 8 kHz session + // would return word timestamps 2x off with a 200 and no diagnostic. + // The file-fed handlers resample their input to make the two + // domains the same one; live audio arrives a frame at a time from a + // client that already picked a rate, and read_audio_file's + // resampler cannot be applied to it incrementally, so the only + // honest answers are this refusal or a wrong timestamp. + // + // ZERO means 16000, which is what backend.proto documents. Every + // OTHER value is refused, including a negative one: -1 is malformed + // rather than absent, and quietly handing it the default would be + // this handler inventing a request the caller did not make. + // + // Costs nothing today: core/backend/transcript_live.go hardcodes + // liveSampleRate = 16000. + if (config.sample_rate() != 0 && + config.sample_rate() != kSpeechSampleRate) { + return GStatus(grpc::StatusCode::INVALID_ARGUMENT, + "audio-cpp: live transcription accepts " + + std::to_string(kSpeechSampleRate) + + " Hz mono PCM only (or 0 to mean it), got " + + std::to_string(config.sample_rate()) + + " Hz; resample on the client side"); + } + const int sample_rate = kSpeechSampleRate; + + // THE LANE IS HELD FOR THE WHOLE STREAM, which is longer than any + // other handler holds it: as long as the user keeps talking. The + // streaming session is stateful and cached, so a concurrent run on + // the same model would interleave its audio with this one's and + // corrupt both transcripts. Named local because LaneEntry is + // immovable and has to span session_for, every chunk and finalize. + audiocpp_backend::LaneEntry lane = model->acquire(0); + const auto session = model->session_for( + audiocpp_backend::Rpc::AudioTranscriptionLive, shape, lane); + + engine::runtime::TaskRequest task; + // An EMPTY buffer carrying the CONTRACT only: the rate and the + // channel count of the frames about to arrive, and no samples, + // because none exist yet. Not decoration. build_preparation_request + // derives the audio contract from this field, and nemotron_asr's + // streaming prepare() throws "Nemotron ASR streaming prepare() + // requires an audio contract" when it is absent, so a live stream + // without it fails on the first model that serves it. + engine::runtime::AudioBuffer contract; + contract.sample_rate = sample_rate; + // Mono, from the proto: TranscriptLiveAudio.pcm is documented as + // "mono PCM in [-1,1] at config.sample_rate" and there is no channel + // field to say otherwise. + contract.channels = 1; + task.audio_input = std::move(contract); + + if (!config.language().empty()) { + task.options["language"] = config.language(); + } + for (const auto ¶m : config.params()) { + task.options[param.first] = param.second; + } + + bool client_gone = false; + backend::TranscriptLiveResponse ack; + ack.set_ready(true); + if (!stream->Write(ack)) { + // The client hung up between opening the stream and the ack. + // Nothing was consumed and nobody is listening. + return GStatus::OK; + } + + audiocpp_backend::TranscriptDeltaTracker tracker; + const auto write_delta = [&](const std::string &fragment) { + if (fragment.empty() || client_gone) { + return; + } + backend::TranscriptLiveResponse response; + response.set_delta(fragment); + if (!stream->Write(response)) { + client_gone = true; + } + }; + + // The SAME reconciliation AudioTranscriptionStream uses, and for the + // same reason: the families disagree about whether partial_text is a + // fragment or the whole hypothesis, one of them delivers the same + // event twice, and a delta that ends inside a UTF-8 sequence makes + // the Go client's Unmarshal fail and costs the client every + // remaining message. None of that is re-derived here. See + // stream_delta.h. + const auto emit = [&](const engine::runtime::StreamEvent &event) { + if (event.partial_text.has_value()) { + write_delta(tracker.observe(event.partial_text->text)); + } + if (event.word_timestamps.empty() || client_gone) { + return; + } + // FORWARD TOLERANT, and dead against the pinned upstream: no + // family puts word timestamps on a StreamEvent, only on the + // TaskResult, so today every word a live client sees arrives in + // final_result below. Kept because the field exists, a family + // that fills it then works with no change here, and the + // conversion is the same one final_result gets. + backend::TranscriptLiveResponse response; + for (const auto &word : event.word_timestamps) { + auto *out = response.add_words(); + // NANOSECONDS. TranscriptWord is the only unit in this proto + // that is not seconds, and core/backend reads it straight + // into a time.Duration. + out->set_start(audiocpp_backend::samples_to_nanoseconds( + word.span.start_sample, sample_rate)); + out->set_end(audiocpp_backend::samples_to_nanoseconds( + word.span.end_sample, sample_rate)); + out->set_text(word.word); + } + if (!stream->Write(response)) { + client_gone = true; + } + }; + + // ARMED ONLY NOW, which is after the lane was taken and before the + // first audio frame is read, and DISARMED when the read side closes. + // Both ends matter: + // + // - not earlier, because acquire() above can legitimately block + // for as long as another live stream is running, and cancelling + // a caller for waiting its turn would be this backend punishing + // its own queue. The read of the CONFIG is likewise unwatched: + // it holds no lane, so a client that stalls there wedges + // nothing. + // - not later, because everything past the read loop is OUR + // compute. A decode that outruns the window is not a peer that + // went quiet, and cancelling there would throw away the + // transcript the client is waiting for. + // + // The status the client actually receives is CANCELLED, not the + // DEADLINE_EXCEEDED returned below: TryCancel is the only way to + // unblock a synchronous Read, and it decides the wire status itself. + // The point of this is not the status, it is that the lane comes + // back. + audiocpp_backend::IdleWatchdog idle( + std::chrono::milliseconds(model->live_idle_timeout_ms()), + [context, &model] { + // Logged, because from the client's side this looks like an + // unexplained cancellation and there is nowhere else to say + // why. + std::cerr << "audio-cpp: live stream idle for " + << model->live_idle_timeout_ms() + << " ms with the send side still open; cancelling " + "it to release the model lane\n"; + context->TryCancel(); + }); + + std::int64_t consumed_frames = 0; + const auto next_frames = [&](std::vector &out) { + while (stream->Read(&incoming)) { + if (incoming.has_config()) { + // backend.proto says a second Config resets the decode + // session. audio.cpp cannot do that truthfully: a reset + // would clear the session's audio while the deltas + // already on the wire cannot be taken back, so the final + // text would then describe only the audio after the + // reset and would CONTRADICT the transcript the client + // assembled. Refused loudly rather than ignored: a + // client that believes it reset the decoder and did not + // is handed a transcript that silently continues the + // audio it thought it discarded. + throw audiocpp_backend::ConfigError( + "audio-cpp: family '" + model->family() + + "' cannot reset a live session mid-stream; close " + "this stream and open a new one"); + } + // Touched only for a frame the decoder can actually consume, + // and touched AFTER the filters rather than before them. + // Neither arm of the oneof set, or an empty pcm field, is + // nothing to feed and not worth failing a live session over + // - but it is also not the peer proving it is still there, + // and resetting the window for it let one client hold the + // model's only lane indefinitely by writing empty frames + // faster than the window. See live_frame_carries_audio. + const auto &pcm = incoming.audio().pcm(); + if (!audiocpp_backend::live_frame_carries_audio( + incoming.has_audio(), pcm.empty())) { + continue; + } + idle.touch(); + out.assign(pcm.begin(), pcm.end()); + // Mono, so floats and frames are the same count. Only used + // for the duration reported in final_result. + consumed_frames += static_cast(pcm.size()); + return true; + } + // The read side is closed, one way or the other. Stop the timer + // BEFORE saying so, since what follows is finalize(). + idle.disarm(); + if (idle.fired()) { + // Read returned false because we cancelled the RPC, not + // because the client closed. Thrown rather than returned so + // the driver does not go on to finalize: there is nobody + // left to send a transcript to, and the whole point was to + // stop occupying the lane. + throw LiveIdleTimeout( + "audio-cpp: no audio frame arrived within " + + std::to_string(model->live_idle_timeout_ms()) + + " ms and the client did not close the stream; cancelled " + "to release the model"); + } + return false; + }; + + // The driver owns the streaming state obligation and the stream + // event sink: nemotron_asr, the only family that serves this RPC + // today, reports EVERY partial through the sink from inside + // finalize(), and its process_audio_chunk returns a bare event. It + // also buffers the wire's frames up to the family's own preferred + // window, which is one second for nemotron_asr and nothing like the + // 512-sample frames a client's audio callback produces. + // + // "LIVE" HERE MEANS INCREMENTAL INPUT, NOT LOW LATENCY, and with the + // families in the pinned checkout it does not yet mean incremental + // OUTPUT either. nemotron_asr's process_audio_chunk only appends to + // its buffer (src/models/nemotron_asr/session.cpp:424-436); the whole + // decode, and therefore every delta, happens inside finalize(). So + // against nemotron_asr every delta arrives AFTER the client closes + // its send side, and the driver's policy-window buffering changes + // nothing measurable. It is not inert for vibevoice_asr or + // higgs_audio_stt, which decode per chunk. Nobody should benchmark + // time-to-first-delta against nemotron_asr and conclude this is + // broken: it is the family, and the fix is a family that streams its + // decode. + const auto result = audiocpp_backend::run_streaming_live( + session, task, next_frames, emit, lane); + + // Makes concatenating every delta equal final_result's text, which + // is what the HTTP layer above assembles. Also what flushes a + // held-back partial UTF-8 sequence. + if (result.text_output.has_value()) { + write_delta(tracker.reconcile(result.text_output->text)); + } + + backend::TranscriptLiveResponse final_response; + audiocpp_backend::fill_transcript_result( + result, sample_rate, + audiocpp_backend::samples_to_seconds(consumed_frames, sample_rate), + final_response.mutable_final_result()); + if (!client_gone) { + stream->Write(final_response); + } + return GStatus::OK; + } catch (const LiveIdleTimeout &err) { + // For the server's own record, and for a client that somehow reads + // it: TryCancel has already decided the wire status is CANCELLED. + // The lane was released by unwinding to here, which is the point. + return GStatus(grpc::StatusCode::DEADLINE_EXCEEDED, err.what()); + } catch (const std::exception &err) { + return to_status(err); + } + } + + // The five surfaces this backend cannot serve at all. + // + // They are overridden rather than left to the generated base class on + // purpose. The base returns UNIMPLEMENTED with an empty message, which + // tells the caller nothing: it cannot distinguish "audio-cpp will never do + // this" from "the model you loaded cannot, try another one" from "this + // build is old". Each override names the loaded family, what that family + // does support, and the upstream limitation, so the answer to "why" is on + // the wire rather than in someone's head. capability_routing.cpp holds the + // reasons; every one of them was checked against the pinned checkout. + // + // None of these is deferred work. Each is an RPC whose LocalAI contract has + // no counterpart in audio.cpp's VoiceTaskKind, and each becomes an ordinary + // Rpc routing entry the day upstream grows one. + + GStatus AudioEncode(ServerContext *, const backend::AudioEncodeRequest *, + backend::AudioEncodeResult *) override { + return refuse_surface(audiocpp_backend::UnsupportedRpc::AudioEncode); + } + + GStatus AudioDecode(ServerContext *, const backend::AudioDecodeRequest *, + backend::AudioDecodeResult *) override { + return refuse_surface(audiocpp_backend::UnsupportedRpc::AudioDecode); + } + + // Refused WITHOUT reading the stream. A bidirectional handler that returns + // a status closes the call, and gRPC delivers that status to the client on + // its next Read or Finish, so the client learns why whether or not it has + // already written frames. Draining first would only delay the same answer + // for as long as the client keeps talking. + GStatus AudioTransformStream( + ServerContext *, + grpc::ServerReaderWriter *) override { + return refuse_surface( + audiocpp_backend::UnsupportedRpc::AudioTransformStream); + } + + // Same, and see the note above about not reading the stream. + GStatus AudioToAudioStream( + ServerContext *, + grpc::ServerReaderWriter *) override { + return refuse_surface( + audiocpp_backend::UnsupportedRpc::AudioToAudioStream); + } + + // The one refusal whose request carries a ModelIdentity, so it is the one + // that has to run the #10952 check before answering. The order is the same + // one snapshot_for enforces everywhere else and for the same reason: a + // request that names a DIFFERENT model must get NOT_FOUND plus the router's + // sentinel, not this UNIMPLEMENTED. Otherwise a stale route pointed at this + // process answers "audio.cpp cannot embed speakers" about a model that is + // not loaded here and may well live on a backend that can, and the router + // retires a working capability on the strength of it. Case 3 at + // snapshot_unchecked explains why snapshot_for cannot be used instead. + GStatus VoiceEmbed(ServerContext *, const backend::VoiceEmbedRequest *request, + backend::VoiceEmbedResponse *) override { + const auto model = snapshot_unchecked(); + if (model != nullptr) { + const GStatus identity = check_model_identity(*model, request); + if (!identity.ok()) { + return identity; + } + } + return refuse_surface(audiocpp_backend::UnsupportedRpc::VoiceEmbed, + model.get()); + } +}; + +void RunServer(const std::string &addr) { + AudioCppBackend service; + grpc::EnableDefaultHealthCheckService(true); + grpc::reflection::InitProtoReflectionServerBuilderPlugin(); + + ServerBuilder builder; + builder.AddListeningPort(addr, grpc::InsecureServerCredentials()); + builder.RegisterService(&service); + // Audio payloads (PCM buffers, encoded frames) are far larger than the + // 4 MiB gRPC default. + builder.SetMaxReceiveMessageSize(256 * 1024 * 1024); + builder.SetMaxSendMessageSize(256 * 1024 * 1024); + + std::unique_ptr server(builder.BuildAndStart()); + if (!server) { + std::cerr << "audio-cpp grpc-server: failed to bind " << addr << "\n"; + std::exit(1); + } + std::cerr << "audio-cpp grpc-server listening on " << addr << "\n"; + + // Wait() runs off the main thread so the main thread stays free to notice + // the shutdown flag and call Shutdown itself, outside any signal context. + std::thread serving([&server] { server->Wait(); }); + + // Polled rather than waited on: notifying a condition variable from a + // signal handler is not async-signal-safe either, so a condition variable + // would move the same defect rather than fix it. Ten wakeups a second on an + // otherwise idle process is not worth a more elaborate scheme. + while (!g_shutdown_requested.load(std::memory_order_relaxed)) { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + + // In-flight RPCs get three seconds to drain before the server stops hard. + server->Shutdown(std::chrono::system_clock::now() + std::chrono::seconds(3)); + serving.join(); +} + +void signal_handler(int) { + // Sets a lock-free atomic and returns. Nothing else belongs here. + // + // Calling grpc::Server::Shutdown directly from a signal handler, which is + // what this used to do, takes an absl::Mutex. That is not async-signal-safe: + // the handler can interrupt a thread that already holds the same mutex, and + // abseil's deadlock detector sees the reentrant acquisition and aborts the + // process. The observed result was exit 134 and a "dying due to potential + // deadlock" stack on every SIGTERM, which is exactly the crash signature a + // real teardown bug would have to compete with. + g_shutdown_requested.store(true, std::memory_order_relaxed); +} + +} // namespace + +int main(int argc, char *argv[]) { + std::string addr = "127.0.0.1:50051"; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + const std::string addr_flag = "--addr="; + if (a.rfind(addr_flag, 0) == 0) { + addr = a.substr(addr_flag.size()); + } else if (a == "--addr" && i + 1 < argc) { + addr = argv[++i]; + } else if (a == "--help" || a == "-h") { + std::cout << "Usage: grpc-server --addr=HOST:PORT\n"; + return 0; + } + } + std::signal(SIGINT, signal_handler); + std::signal(SIGTERM, signal_handler); + RunServer(addr); + return 0; +} diff --git a/backend/cpp/audio-cpp/inference_lane.cpp b/backend/cpp/audio-cpp/inference_lane.cpp new file mode 100644 index 000000000..e619afcaf --- /dev/null +++ b/backend/cpp/audio-cpp/inference_lane.cpp @@ -0,0 +1,114 @@ +#include "inference_lane.h" + +#include +#include + +namespace audiocpp_backend { + +std::int64_t monotonic_millis() { + return std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count(); +} + +int resolve_wait_budget_ms(int policy_ceiling_ms, int request_hint_ms) { + // Both sides normalise to 0 for "not specified", which is also the value + // that means unbounded on the way out, so the unspecified cases fall out of + // the arithmetic instead of needing their own branches. + const int ceiling = policy_ceiling_ms > 0 ? policy_ceiling_ms : 0; + const int hint = request_hint_ms > 0 ? request_hint_ms : 0; + + if (ceiling == 0) { + return hint; + } + if (hint == 0) { + return ceiling; + } + // The hint can tighten the ceiling but never loosen it. + return std::min(ceiling, hint); +} + +bool run_exceeds_budget(bool lane_occupied, std::int64_t run_started_ms, + std::int64_t now_ms, int budget_ms) { + if (!lane_occupied || budget_ms <= 0) { + return false; + } + return now_ms - run_started_ms > static_cast(budget_ms); +} + +namespace { + +std::string busy_prefix(const std::string &model_label) { + return "inference lane for model '" + model_label + "' is busy: "; +} + +// States the measurement and nothing else. A short-budget caller meeting a +// legitimately long run lands here too, so this text must not declare the run +// broken; the numbers let a reader decide that for themselves. +std::string overrun_message(const std::string &model_label, + std::int64_t run_age_ms, int budget_ms) { + return busy_prefix(model_label) + "the in-flight run has been running for " + + std::to_string(run_age_ms) + " ms, longer than this request's " + + std::to_string(budget_ms) + " ms wait budget"; +} + +std::string wait_exhausted_message(const std::string &model_label, + int budget_ms) { + return busy_prefix(model_label) + "timed out after " + + std::to_string(budget_ms) + + " ms waiting for the in-flight run to finish"; +} + +} // namespace + +void InferenceLane::occupy(int budget_ms) { + std::unique_lock lock(state_mutex_); + + // Checked once, on arrival: if the run already in the lane has outlived what + // this caller brought, no amount of waiting can help it, and queueing here + // is exactly how a wedged run swallows every handler thread. + const std::int64_t arrived_ms = monotonic_millis(); + if (run_exceeds_budget(occupied_, run_started_ms_, arrived_ms, budget_ms)) { + throw LaneUnavailable(overrun_message( + model_label_, arrived_ms - run_started_ms_, budget_ms)); + } + + if (budget_ms > 0) { + if (!vacated_.wait_for(lock, std::chrono::milliseconds(budget_ms), + [this] { return !occupied_; })) { + throw LaneUnavailable( + wait_exhausted_message(model_label_, budget_ms)); + } + } else { + vacated_.wait(lock, [this] { return !occupied_; }); + } + + // Only now, holding both the mutex and the lane. Stamping any earlier would + // restart the age of the run for every caller behind this one and make a + // genuinely wedged holder look freshly started forever. + occupied_ = true; + run_started_ms_ = monotonic_millis(); +} + +void InferenceLane::vacate() { + { + std::lock_guard lock(state_mutex_); + occupied_ = false; + } + // notify_all, not notify_one: a waiter that times out concurrently with a + // notification can consume it, and losing the only wakeup would park the + // remaining waiters for the rest of the run's lifetime. Waking all of them + // still admits exactly one, since the rest re-test occupancy under the mutex + // and go back to waiting, and ordering among waiters is not a requirement. + vacated_.notify_all(); +} + +LaneEntry::LaneEntry(InferenceLane &lane, int budget_ms) : lane_(lane) { + // If this throws, the object never existed, so ~LaneEntry does not run and + // cannot hand back a lane this caller never held. + lane_.occupy(budget_ms); +} + +LaneEntry::~LaneEntry() { lane_.vacate(); } + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/inference_lane.h b/backend/cpp/audio-cpp/inference_lane.h new file mode 100644 index 000000000..c9df3aa43 --- /dev/null +++ b/backend/cpp/audio-cpp/inference_lane.h @@ -0,0 +1,138 @@ +#pragma once + +// Serializes inference against the single audio.cpp model a backend process +// owns. Standard library only. +// +// Why a plain mutex is not enough: an audio.cpp session is not reentrant, so +// concurrent gRPC handlers have to take turns. But once a GPU call stops making +// progress there is nothing a host thread can do to take it back, and an +// unbounded queue behind such a run would absorb the handler threads one by one +// until nothing is left to answer with. A caller therefore needs to be able to +// walk away, and needs to be able to tell "the lane is busy with normal work and +// I ran out of patience" apart from "the run in the lane has already outlived +// the patience I brought". + +#include +#include +#include +#include +#include + +namespace audiocpp_backend { + +// Reading off a monotonic clock, in milliseconds. Monotonic on purpose: a wall +// clock adjustment must never make an in-flight run look younger or older than +// it is, because that reading decides whether callers give up. +std::int64_t monotonic_millis(); + +// Collapses the per-model configured ceiling and the optional per-request hint +// into the wait budget a caller actually gets. Returns 0 for "wait +// indefinitely". +// +// Either input may be non-positive, which means "not specified": +// - an unspecified hint yields the ceiling, +// - an unspecified ceiling means no policy limit, so the hint stands, +// - unspecified on both sides is unbounded. +// A specified hint may only tighten the ceiling. A client asking for a longer +// wait than the model's policy allows does not get it, because that would let a +// request weaken an operator's choice. +int resolve_wait_budget_ms(int policy_ceiling_ms, int request_hint_ms); + +// True when a caller carrying budget_ms should give up on arrival rather than +// queue up. Deliberately a pure function of the lane's observable state so the +// decision can be tested without threads or sleeping. +// +// Only occupancy makes a start timestamp meaningful: a lane nobody holds is +// never overrunning, whatever timestamp the last holder left behind. An +// unbounded caller (non-positive budget) has no budget to exceed. And the +// comparison is strict, so a run whose age exactly equals the budget still has +// its last millisecond. +bool run_exceeds_budget(bool lane_occupied, std::int64_t run_started_ms, + std::int64_t now_ms, int budget_ms); + +// Thrown when a caller cannot take the lane, in either of the two situations +// resolve_wait_budget_ms allows for. The message distinguishes them; callers +// that need to report a status code can treat them alike. +class LaneUnavailable : public std::runtime_error { + public: + explicit LaneUnavailable(const std::string &reason) + : std::runtime_error(reason) {} +}; + +class LaneEntry; + +// One lane per loaded model. Shared by every handler thread; not copyable. +class InferenceLane { + public: + explicit InferenceLane(std::string model_label) + : model_label_(std::move(model_label)) {} + + InferenceLane(const InferenceLane &) = delete; + InferenceLane &operator=(const InferenceLane &) = delete; + + const std::string &model_label() const { return model_label_; } + + private: + // Occupancy is only reachable through LaneEntry, so there is no way to take + // the lane without also having something that gives it back. + friend class LaneEntry; + + void occupy(int budget_ms); + void vacate(); + + const std::string model_label_; + + std::mutex state_mutex_; + std::condition_variable vacated_; + bool occupied_ = false; + // Only meaningful while occupied_ is true. + std::int64_t run_started_ms_ = 0; +}; + +// Scoped occupancy of a lane. Construct it where the inference happens and it +// is given back on every exit from that scope, including an exception and +// including a caller that returns from the middle of a long stream. Throws +// LaneUnavailable if the lane could not be taken, in which case there is no +// object and nothing to release. +// +// Not reentrant, and it does not detect reentrancy: a second entry constructed +// while the calling thread already holds the same lane waits for a lane only +// that thread can release. With a positive budget that surfaces as +// LaneUnavailable, but in unbounded mode the thread parks with no diagnostic at +// all. Keep entries one per call: a handler that holds one across a stream must +// not let a helper it calls construct another. +class LaneEntry { + public: + // budget_ms <= 0 waits indefinitely. Pass the output of + // resolve_wait_budget_ms. + LaneEntry(InferenceLane &lane, int budget_ms); + ~LaneEntry(); + + LaneEntry(const LaneEntry &) = delete; + LaneEntry &operator=(const LaneEntry &) = delete; + + // Deliberately immovable rather than carefully movable. A moved-from entry + // would have to stop releasing the lane while the lane still records it as + // occupied, and that hazard is not worth the convenience: the lane can only + // be recovered by whoever took it. + // + // To hold a lane for longer than one scope, construct the entry in place + // instead of moving one in. Two shapes work: + // std::optional held; // member or local + // held.emplace(lane, budget_ms); // takes the lane, held.reset() gives it back + // auto held = std::make_unique(lane, budget_ms); // also returnable + // Both outlive the acquiring scope and still release exactly once, when they + // are reset or destroyed. Prefer the optional for a member whose lifetime is + // the handler's; use the unique_ptr when the entry has to be returned, since + // an optional of an immovable type is itself immovable and cannot be. A + // factory may instead write `return LaneEntry(lane, budget_ms);`, which C++17 + // guarantees to elide, whereas `LaneEntry entry(...); return entry;` does not + // compile, because that form is a move. + LaneEntry(LaneEntry &&) = delete; + LaneEntry &operator=(LaneEntry &&) = delete; + + private: + InferenceLane &lane_; +}; + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/inference_lane_test.cpp b/backend/cpp/audio-cpp/inference_lane_test.cpp new file mode 100644 index 000000000..bee36b04b --- /dev/null +++ b/backend/cpp/audio-cpp/inference_lane_test.cpp @@ -0,0 +1,618 @@ +// Unit tests for inference_lane. Standard library only. The harness compiles +// this file as a single translation unit, so the implementation is included +// directly rather than linked. +// +// Two kinds of test live here: +// +// * The pure ones (budget negotiation, the overrun predicate) run with no +// threads and no sleeping. They carry the arithmetic, so they are the tests +// that must be exhaustive. +// * The threaded ones exercise the lane itself. Every one of them is bounded: +// contenders use a generous wait budget instead of the unbounded mode +// wherever the point of the test does not require unbounded, and a watchdog +// in main() puts a ceiling on the whole file. A broken implementation must +// go red, not hang, because a hung job costs CI far more than a red one. +// +// Wall-clock margins are called out individually. The rule applied throughout: +// a margin is only allowed if a slow or loaded machine pushes the measurement +// deeper into the passing region. + +#include "inference_lane.cpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static int failures = 0; + +static void check(bool ok, const std::string &name) { + if (!ok) { + failures++; + fprintf(stderr, "FAIL: %s\n", name.c_str()); + } else { + fprintf(stderr, "ok: %s\n", name.c_str()); + } +} + +using audiocpp_backend::InferenceLane; +using audiocpp_backend::LaneEntry; +using audiocpp_backend::LaneUnavailable; +using audiocpp_backend::resolve_wait_budget_ms; +using audiocpp_backend::run_exceeds_budget; + +// A one-shot, level-triggered signal with a bounded wait. Preferred over sleeps +// for "the other thread got there" so the tests do not encode a guess about +// scheduling. +class Signal { + public: + void raise() { + { + std::lock_guard lock(mutex_); + raised_ = true; + } + cv_.notify_all(); + } + + bool await(int timeout_ms) { + std::unique_lock lock(mutex_); + return cv_.wait_for(lock, std::chrono::milliseconds(timeout_ms), + [this] { return raised_; }); + } + + private: + std::mutex mutex_; + std::condition_variable cv_; + bool raised_ = false; +}; + +static std::int64_t elapsed_ms_since( + const std::chrono::steady_clock::time_point &start) { + return std::chrono::duration_cast( + std::chrono::steady_clock::now() - start) + .count(); +} + +static void nap(int ms) { + std::this_thread::sleep_for(std::chrono::milliseconds(ms)); +} + +static bool mentions(const std::string &haystack, const std::string &needle) { + return haystack.find(needle) != std::string::npos; +} + +// Wording that separates the two failure modes. Kept here so a message reword +// that erases the distinction breaks these tests loudly. +static const char *const kTimedOutPhrase = "timed out after"; +static const char *const kStillRunningPhrase = "has been running for"; + +// Attempts an entry and reports what happened, so the threaded tests can assert +// on the message rather than only on the exception type. +struct EntryOutcome { + bool acquired = false; + std::string message; + std::int64_t took_ms = 0; +}; + +static EntryOutcome try_entry(InferenceLane &lane, int budget_ms) { + EntryOutcome out; + const auto start = std::chrono::steady_clock::now(); + try { + LaneEntry entry(lane, budget_ms); + out.acquired = true; + } catch (const LaneUnavailable &refused) { + out.message = refused.what(); + } + out.took_ms = elapsed_ms_since(start); + return out; +} + +// --------------------------------------------------------------------------- +// B9: budget negotiation. Pure, no threads. +// --------------------------------------------------------------------------- +static void test_budget_negotiation() { + check(resolve_wait_budget_ms(0, 0) == 0, + "B9 no ceiling and no request hint is unbounded"); + check(resolve_wait_budget_ms(-1, -1) == 0, + "B9 negative ceiling and negative hint is unbounded"); + + check(resolve_wait_budget_ms(5000, 0) == 5000, + "B9 absent hint yields the ceiling"); + check(resolve_wait_budget_ms(5000, -250) == 5000, + "B9 negative hint yields the ceiling"); + + check(resolve_wait_budget_ms(5000, 1200) == 1200, + "B9 a shorter request hint is granted"); + check(resolve_wait_budget_ms(5000, 1) == 1, + "B9 a much shorter request hint is granted"); + + check(resolve_wait_budget_ms(5000, 9000) == 5000, + "B9 a longer request hint cannot weaken the ceiling"); + check(resolve_wait_budget_ms(5000, 5001) == 5000, + "B9 a hint one ms over the ceiling is clamped"); + check(resolve_wait_budget_ms(5000, 5000) == 5000, + "B9 a hint equal to the ceiling is the ceiling"); + + check(resolve_wait_budget_ms(0, 1200) == 1200, + "B9 without a policy limit the request hint applies"); + check(resolve_wait_budget_ms(-5, 1200) == 1200, + "B9 a negative ceiling is no policy limit"); + + check(resolve_wait_budget_ms(1, 0) == 1, + "B9 a one ms ceiling survives negotiation"); +} + +// --------------------------------------------------------------------------- +// B5 and B6: the overrun predicate. Pure, no threads. +// --------------------------------------------------------------------------- +static void test_overrun_predicate() { + // B5: strictly longer. + check(run_exceeds_budget(true, 1000, 1100, 100) == false, + "B5 elapsed exactly equal to the budget is not an overrun"); + check(run_exceeds_budget(true, 1000, 1101, 100) == true, + "B5 one ms past the budget is an overrun"); + check(run_exceeds_budget(true, 1000, 1099, 100) == false, + "B5 one ms short of the budget is not an overrun"); + check(run_exceeds_budget(true, 0, 1, 1) == false, + "B5 a one ms budget at one ms elapsed is not an overrun"); + check(run_exceeds_budget(true, 0, 2, 1) == true, + "B5 a one ms budget at two ms elapsed is an overrun"); + + // B6: an unoccupied lane is never stuck, whatever the leftover timestamp + // says. This is the pure half of B6; the wiring half is threaded below. + check(run_exceeds_budget(false, 0, 10000000, 1) == false, + "B6 an idle lane with an ancient start timestamp is not an overrun"); + check(run_exceeds_budget(false, 5, 5, 5) == false, + "B6 an idle lane is not an overrun at any elapsed value"); + + // An unbounded caller has no budget to exceed, so it never fails fast. + check(run_exceeds_budget(true, 0, 10000000, 0) == false, + "B2 an unbounded caller never sees an overrun"); + check(run_exceeds_budget(true, 0, 10000000, -1) == false, + "B2 a negative budget never sees an overrun"); +} + +// --------------------------------------------------------------------------- +// B1: mutual exclusion under real contention, and a release admits a waiter. +// --------------------------------------------------------------------------- +static void test_mutual_exclusion() { + InferenceLane lane("exclusion-model"); + + constexpr int kContenders = 4; // "at least three simultaneous contenders" + constexpr int kHoldMs = 15; + // Generous on purpose: the point of this test is exclusion, not timeouts. + // A larger budget only makes a healthy run more likely to pass, while still + // bounding a broken one at roughly five seconds instead of forever. + constexpr int kBudgetMs = 5000; + + std::atomic in_flight{0}; + std::atomic peak_in_flight{0}; + std::atomic completed{0}; + std::atomic refused{0}; + + Signal go; + std::vector contenders; + for (int i = 0; i < kContenders; i++) { + contenders.emplace_back([&] { + go.await(5000); + try { + LaneEntry entry(lane, kBudgetMs); + const int now_inside = in_flight.fetch_add(1) + 1; + int seen = peak_in_flight.load(); + while (now_inside > seen && + !peak_in_flight.compare_exchange_weak(seen, now_inside)) { + // retry with the refreshed value + } + nap(kHoldMs); + in_flight.fetch_sub(1); + completed.fetch_add(1); + } catch (const LaneUnavailable &) { + refused.fetch_add(1); + } + }); + } + go.raise(); + for (auto &t : contenders) { + t.join(); + } + + check(refused.load() == 0, "B1 no contender was refused within its budget"); + check(completed.load() == kContenders, + "B1 every contender eventually got the lane"); + check(peak_in_flight.load() == 1, + "B1 never more than one holder inside the lane at once"); +} + +// --------------------------------------------------------------------------- +// B2: unbounded mode waits out a run longer than any bound would allow. +// --------------------------------------------------------------------------- +static void test_unbounded_waits_out_the_holder() { + InferenceLane lane("patient-model"); + constexpr int kHoldMs = 300; + + Signal held; + Signal patient_done; + std::int64_t patient_wait_ms = -1; + bool patient_acquired = false; + + std::thread holder([&] { + LaneEntry entry(lane, 0); + held.raise(); + nap(kHoldMs); + }); + check(held.await(5000), "B2 holder took the lane"); + + std::thread patient([&] { + const auto start = std::chrono::steady_clock::now(); + try { + LaneEntry entry(lane, 0); + patient_acquired = true; + } catch (const LaneUnavailable &) { + patient_acquired = false; + } + patient_wait_ms = elapsed_ms_since(start); + patient_done.raise(); + }); + + // Bound on the unbounded mode: a lane that never wakes its waiters goes red + // here instead of hanging in join(). The named failure is printed before the + // join, so even a hard hang leaves a diagnosis behind for the watchdog. + check(patient_done.await(10000), "B2 unbounded caller returned at all"); + patient.join(); + holder.join(); + + check(patient_acquired, "B2 unbounded caller acquired instead of failing"); + // Margin: the holder holds for 300 ms, so the waiter must block for about + // that long. Asserting only half of it means a loaded machine, which makes + // the wait longer, drifts further into passing. + check(patient_wait_ms >= kHoldMs / 2, + "B2 unbounded caller actually waited for the in-flight run"); +} + +// --------------------------------------------------------------------------- +// B3: a bounded caller that cannot get in gives up with the timeout wording. +// --------------------------------------------------------------------------- +static void test_bounded_wait_times_out() { + InferenceLane lane("impatient-model"); + constexpr int kBudgetMs = 120; + + Signal held; + Signal release; + std::thread holder([&] { + LaneEntry entry(lane, 0); + held.raise(); + release.await(10000); + }); + check(held.await(5000), "B3 holder took the lane"); + + // The waiter arrives immediately, so the holder's elapsed time is far below + // the budget and the fail-fast path must not trigger here. + // + // Load-sensitive margin, and the tightest one in this file: what makes this + // the timeout path rather than the fail-fast path is the holder's age + // staying under 120 ms at the arrival check. All that sits between the + // holder's stamp and this call is one signal handover, microseconds against + // a 120 ms allowance, but unlike the other margins here load pushes this one + // toward failing rather than away from it. If it ever does flip, the symptom + // is the wording assertions below going red, not a hang, and the fix is a + // larger budget rather than a weaker assertion. + const EntryOutcome outcome = try_entry(lane, kBudgetMs); + release.raise(); + holder.join(); + + check(!outcome.acquired, "B3 bounded caller did not acquire a held lane"); + check(mentions(outcome.message, kTimedOutPhrase), + "B3 failure names the exhausted wait, not an overrunning run"); + check(!mentions(outcome.message, kStillRunningPhrase), + "B3 failure is not worded as an overrun"); + check(mentions(outcome.message, "impatient-model"), + "B3 failure names the model"); + check(mentions(outcome.message, "120"), + "B3 failure reports the budget it waited out"); + // Margin: wait_for cannot return before its deadline, so the true value is + // at least 120 ms and load only raises it. Asserting 100 leaves room for + // clock granularity while still catching an implementation that returns + // early without waiting. + check(outcome.took_ms >= 100, "B3 bounded caller waited out its budget"); +} + +// --------------------------------------------------------------------------- +// B4: a caller whose budget is already exceeded fails at once. +// --------------------------------------------------------------------------- +static void test_fail_fast_against_a_long_run() { + InferenceLane lane("wedged-model"); + constexpr int kBudgetMs = 200; + constexpr int kRunAgeMs = 400; + + Signal held; + Signal release; + std::thread holder([&] { + LaneEntry entry(lane, 0); + held.raise(); + release.await(10000); + }); + check(held.await(5000), "B4 holder took the lane"); + + // Margin: the arriving caller needs the holder's elapsed time to exceed + // 200 ms. Sleeping 400 ms means a loaded machine oversleeps and pushes the + // elapsed time further past the budget, never below it. + nap(kRunAgeMs); + const EntryOutcome outcome = try_entry(lane, kBudgetMs); + release.raise(); + holder.join(); + + check(!outcome.acquired, "B4 caller did not acquire a long-running lane"); + check(mentions(outcome.message, kStillRunningPhrase), + "B4 failure states the measured age of the in-flight run"); + check(!mentions(outcome.message, kTimedOutPhrase), + "B4 failure is not worded as an exhausted wait"); + check(mentions(outcome.message, "wedged-model"), + "B4 failure names the model"); + // Margin: a fail-fast return takes microseconds. 150 ms of headroom under a + // 200 ms budget separates "returned at once" from "waited out the budget" + // by a wide enough gap that scheduler noise cannot close it. The message + // assertions above are the load-independent proof; this one pins the timing. + check(outcome.took_ms < 150, "B4 caller failed without waiting out its budget"); +} + +// --------------------------------------------------------------------------- +// B6 wiring: an idle lane never looks stuck, however old the last run is. +// --------------------------------------------------------------------------- +static void test_idle_lane_is_never_stuck() { + InferenceLane lane("idle-model"); + + { + LaneEntry entry(lane, 0); + } + // Ages the leftover start timestamp well past the tiny budget used below. + // A longer sleep only makes a stale-timestamp bug more visible, so load + // helps this test rather than hurting it. + nap(80); + + const EntryOutcome first = try_entry(lane, 20); + check(first.acquired, "B6 tiny budget still acquires an idle lane"); + + nap(80); + const EntryOutcome second = try_entry(lane, 1); + check(second.acquired, "B6 a one ms budget still acquires an idle lane"); +} + +// --------------------------------------------------------------------------- +// B7: every ownership exit clears the busy state, including an exception +// thrown from inside the guarded region. +// --------------------------------------------------------------------------- +static void test_release_on_exception() { + InferenceLane lane("throwing-model"); + + struct GuardedRegionFailure {}; + bool propagated = false; + try { + LaneEntry entry(lane, 0); + throw GuardedRegionFailure{}; + } catch (const GuardedRegionFailure &) { + propagated = true; + } + check(propagated, "B7 the guarded region's own exception propagated"); + + // If the throw had leaked the busy state, this tiny budget would fail. + const EntryOutcome after_throw = try_entry(lane, 20); + check(after_throw.acquired, "B7 lane is free after an exception unwound it"); + + // Same check for a holder that unwinds on another thread, which is the shape + // a gRPC handler failing mid-inference actually has. + Signal thrown; + std::thread unlucky([&] { + try { + LaneEntry entry(lane, 0); + throw GuardedRegionFailure{}; + } catch (const GuardedRegionFailure &) { + thrown.raise(); + } + }); + check(thrown.await(5000), "B7 worker thread unwound its guarded region"); + unlucky.join(); + + const EntryOutcome after_worker = try_entry(lane, 20); + check(after_worker.acquired, + "B7 lane is free after a worker thread unwound it"); +} + +// --------------------------------------------------------------------------- +// B8: a waiter must not publish itself as the holder. If it did, its arrival +// would restart the elapsed-time measurement and hide the real holder. +// +// Timeline, with the holder taking the lane at t0 and never letting go: +// +// t0 holder acquires, elapsed measurement starts here and only here +// t0+100 waiter arrives with a 400 ms budget, blocks, and times out +// t0+500 late caller arrives with a 450 ms budget +// +// A correct lane measures 500 ms of holding at the late arrival, which is more +// than 450, so the late caller fails fast. An implementation that let the +// waiter stamp itself as holder measures only the 400 ms since the waiter +// arrived, which is under 450, so the late caller would queue behind a stuck +// run instead. The two paths are told apart by their wording. +// --------------------------------------------------------------------------- +static void test_waiter_does_not_become_the_holder() { + InferenceLane lane("stamp-model"); + constexpr int kWaiterArrivesAfterMs = 100; + constexpr int kWaiterBudgetMs = 400; + constexpr int kLateBudgetMs = 450; + + Signal held; + Signal release; + std::thread holder([&] { + LaneEntry entry(lane, 0); + held.raise(); + release.await(10000); + }); + check(held.await(5000), "B8 holder took the lane"); + + // Margin: the waiter must not fail fast on arrival, which needs the + // holder's elapsed time to stay under 400 ms. Arriving at 100 ms leaves + // 300 ms of slack, so oversleeping under load does not flip the path. + nap(kWaiterArrivesAfterMs); + + EntryOutcome waiter_outcome; + std::thread waiter([&] { waiter_outcome = try_entry(lane, kWaiterBudgetMs); }); + waiter.join(); + check(!waiter_outcome.acquired, "B8 mid-queue waiter did not acquire"); + check(mentions(waiter_outcome.message, kTimedOutPhrase), + "B8 mid-queue waiter waited out its budget and timed out"); + + // Margin: the holder has now been in the lane for at least 500 ms against a + // 450 ms budget. Load lengthens both sleeps, so the measured age only grows + // and the fail-fast path only becomes more certain. + const EntryOutcome late = try_entry(lane, kLateBudgetMs); + release.raise(); + holder.join(); + + check(!late.acquired, "B8 late caller did not acquire"); + check(mentions(late.message, kStillRunningPhrase), + "B8 elapsed time is still measured from the real holder's acquisition"); + check(late.took_ms < 200, + "B8 late caller failed fast rather than queueing behind the holder"); +} + +// --------------------------------------------------------------------------- +// B8, other direction: the age of a run is measured from the moment its holder +// acquired, not from the moment that holder arrived. A caller that queued for a +// while and then got in is starting a fresh run, and its time in the queue must +// not be billed to it: if it were, every handover would hand the new holder a +// head start towards looking overrun, and short-budget callers would be turned +// away from a run that has barely begun. +// +// t0 first holder acquires and holds for 300 ms +// t0 second caller arrives and queues +// t0+300 second caller acquires, so its own run age is ~0 here +// t0+300 a third caller arrives with a 300 ms budget +// +// Correct: the third caller sees a run that just started, so it queues and then +// times out. Billing the queue time to the second caller would show a 300 ms old +// run instead, and the third caller would be turned away as an overrun. +// --------------------------------------------------------------------------- +static void test_run_age_starts_at_acquisition() { + InferenceLane lane("handover-model"); + constexpr int kFirstHoldMs = 300; + constexpr int kThirdBudgetMs = 200; + + Signal first_held; + Signal handed_over; + Signal release_second; + + std::thread first([&] { + LaneEntry entry(lane, 0); + first_held.raise(); + nap(kFirstHoldMs); + }); + // Ordering matters only for the diagnosis, not for the assertion: waiting + // for the first holder guarantees the second caller really does queue, which + // is what gives it queue time to be wrongly billed for. + check(first_held.await(5000), "B8 first holder took the lane"); + + std::thread second([&] { + LaneEntry entry(lane, 0); + handed_over.raise(); + release_second.await(10000); + }); + check(handed_over.await(10000), "B8 queued caller was handed the lane"); + + // Margin: the new holder's run is a few ms old against a 200 ms budget, so + // this caller must queue. Load can only add a few ms of handover latency, + // well inside that slack, while it lengthens the queue time that the buggy + // version would bill, making the bug more visible rather than less. + const EntryOutcome third = try_entry(lane, kThirdBudgetMs); + release_second.raise(); + second.join(); + first.join(); + + check(!third.acquired, "B8 lane was still held by the queued caller"); + check(mentions(third.message, kTimedOutPhrase), + "B8 a fresh holder's run age excludes the time it spent queueing"); +} + +// --------------------------------------------------------------------------- +// B10: both failure modes carry a usable message, and the fail-fast one reports +// a measurement rather than diagnosing a cause. +// --------------------------------------------------------------------------- +static void test_failure_messages_are_diagnosable() { + InferenceLane lane("diagnosable-model"); + + Signal held; + Signal release; + std::thread holder([&] { + LaneEntry entry(lane, 0); + held.raise(); + release.await(10000); + }); + check(held.await(5000), "B10 holder took the lane"); + + // The two budgets have to straddle the run's age, or both callers take the + // same path and the comparisons below are between two fail-fast messages + // that differ only in the budget they print. + // + // Margin: 30 ms is far under the ~120 ms age, and load only ages the run + // further, so `fast` fails fast. 400 ms is far over it, with the same + // slack and the same safe direction as the B3 test, so `slow` queues and + // then times out. + nap(120); + const EntryOutcome fast = try_entry(lane, 30); + const EntryOutcome slow = try_entry(lane, 400); + release.raise(); + holder.join(); + + check(mentions(fast.message, kStillRunningPhrase) && + mentions(slow.message, kTimedOutPhrase), + "B10 one caller took the fail-fast path and the other timed out"); + check(fast.message != slow.message, + "B10 the two failure modes do not share one message"); + check(!fast.message.empty() && !slow.message.empty(), + "B10 both failures carry text"); + check(mentions(fast.message, "diagnosable-model") && + mentions(slow.message, "diagnosable-model"), + "B10 both failures name the model"); + // The fail-fast wording must not accuse the run of being stuck: a short + // budget meeting a legitimately long run reaches this path too. + for (const char *verdict : {"stuck", "wedged", "hung", "deadlock"}) { + check(!mentions(fast.message, verdict), + std::string("B10 fail-fast message avoids diagnosing '") + + verdict + "'"); + } +} + +int main() { + // Last resort only. Every threaded test above is individually bounded, so + // this should never fire; it exists so that an implementation which parks a + // thread forever still ends the job instead of occupying a CI runner. + std::thread watchdog([] { + std::this_thread::sleep_for(std::chrono::seconds(30)); + fprintf(stderr, "FAIL: watchdog fired, an inference_lane test hung\n"); + fflush(stderr); + std::_Exit(1); + }); + watchdog.detach(); + + test_budget_negotiation(); + test_overrun_predicate(); + test_mutual_exclusion(); + test_unbounded_waits_out_the_holder(); + test_bounded_wait_times_out(); + test_fail_fast_against_a_long_run(); + test_idle_lane_is_never_stuck(); + test_release_on_exception(); + test_waiter_does_not_become_the_holder(); + test_run_age_starts_at_acquisition(); + test_failure_messages_are_diagnosable(); + + if (failures == 0) { + fprintf(stderr, "\nAll inference_lane tests passed.\n"); + return 0; + } + fprintf(stderr, "\n%d inference_lane test(s) failed.\n", failures); + return 1; +} diff --git a/backend/cpp/audio-cpp/live_watchdog.cpp b/backend/cpp/audio-cpp/live_watchdog.cpp new file mode 100644 index 000000000..94263c782 --- /dev/null +++ b/backend/cpp/audio-cpp/live_watchdog.cpp @@ -0,0 +1,75 @@ +#include "live_watchdog.h" + +#include + +namespace audiocpp_backend { + +IdleWatchdog::IdleWatchdog(std::chrono::milliseconds window, + std::function on_idle) + : window_(window), on_idle_(std::move(on_idle)), + last_(std::chrono::steady_clock::now()) { + if (window_.count() <= 0) { + // Disabled: no thread at all, rather than a thread with an infinite + // deadline. A thread that exists is a thread that has to be joined on + // every exit path, and there is nothing for this one to do. + return; + } + thread_ = std::thread([this] { run(); }); +} + +IdleWatchdog::~IdleWatchdog() { disarm(); } + +void IdleWatchdog::touch() { + std::lock_guard lock(mu_); + last_ = std::chrono::steady_clock::now(); + // Deliberately does NOT notify. The waiter recomputes its deadline from + // last_ every time it wakes, so a touch that lands mid-window is picked up + // when the old deadline expires, and a touch is the hot path: it runs once + // per frame on the wire. +} + +void IdleWatchdog::disarm() { + { + std::lock_guard lock(mu_); + stop_ = true; + } + cv_.notify_all(); + if (thread_.joinable()) { + thread_.join(); + } +} + +bool IdleWatchdog::fired() const { + std::lock_guard lock(mu_); + return fired_; +} + +void IdleWatchdog::run() { + std::unique_lock lock(mu_); + while (!stop_) { + const auto deadline = last_ + window_; + if (cv_.wait_until(lock, deadline, [this] { return stop_; })) { + return; // disarmed + } + // The deadline passed, but last_ may have moved while this thread was + // waiting, and a condition variable may also wake spuriously. Re-read + // it: without this check a touch that landed mid-window would still be + // followed by a cancellation, i.e. a live client cut off mid-sentence. + if (std::chrono::steady_clock::now() < last_ + window_) { + continue; + } + fired_ = true; + auto callback = on_idle_; + lock.unlock(); + if (callback) { + callback(); + } + return; // one shot + } +} + +bool live_frame_carries_audio(bool has_audio, bool pcm_empty) { + return has_audio && !pcm_empty; +} + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/live_watchdog.h b/backend/cpp/audio-cpp/live_watchdog.h new file mode 100644 index 000000000..7f82ec124 --- /dev/null +++ b/backend/cpp/audio-cpp/live_watchdog.h @@ -0,0 +1,99 @@ +#pragma once + +// A one-shot idle timer for a bidirectional stream. Standard library only, so +// it is tested without an audio.cpp checkout or a gRPC server. +// +// WHY IT EXISTS. AudioTranscriptionLive holds the model's inference lane for the +// whole stream, because the streaming session is stateful and a concurrent run +// would interleave two callers' audio. Every other RPC in this backend holds the +// lane across COMPUTE, or across a write to a slow reader, and both of those +// terminate on their own. A live stream instead blocks in a client-driven read, +// and a peer that goes silent WITHOUT closing the stream never terminates +// anything: the lane stays taken and every other RPC against that model queues +// behind a client that has stopped speaking. A websocket death does cancel the +// RPC and free it, but "the peer's TCP connection eventually dies" is not a +// bound anyone can state, so this supplies one. +// +// HOW IT ENDS THE STREAM, and the part that is not obvious: gRPC's synchronous +// ServerReaderWriter::Read has no timeout and cannot be given one. The only way +// to unblock it from another thread is ServerContext::TryCancel, which is what +// the callback is for. That means the client sees CANCELLED rather than whatever +// status the handler goes on to return: the returned status is for the server's +// own record. Releasing the lane is the point. +// +// ONE SHOT on purpose. Once the callback has run the stream is being torn down, +// so there is nothing left to watch, and a repeating timer would call TryCancel +// on a context the handler may already have returned from. + +#include +#include +#include +#include +#include + +namespace audiocpp_backend { + +class IdleWatchdog { +public: + // A window that is not positive DISABLES the watchdog entirely: no thread is + // started and fired() never becomes true. That is the operator's escape + // hatch for a client that legitimately holds a stream open through long + // pauses, and it is why the option carrying it documents 0 as "no limit" + // rather than as "expire immediately". + // + // `on_idle` runs on the watchdog's own thread with no lock held. It must be + // safe to call while the watched thread is blocked in a read, which is the + // only reason this class exists; ServerContext::TryCancel is documented as + // exactly that. + IdleWatchdog(std::chrono::milliseconds window, std::function on_idle); + + // Joins the thread, so the callback can safely capture anything that + // outlives this object's scope and nothing else has to be reasoned about. + ~IdleWatchdog(); + + IdleWatchdog(const IdleWatchdog &) = delete; + IdleWatchdog &operator=(const IdleWatchdog &) = delete; + + // Restarts the window. Call it whenever the peer proves it is still there. + void touch(); + + // Stops watching and joins. Idempotent, and REQUIRED before any long + // non-read work the window must not cover: the caller's own decode is not + // the peer going quiet, and cancelling in the middle of it would throw away + // a transcript the client is waiting for. + void disarm(); + + // True once the window elapsed and the callback ran. Stays true after + // disarm, so the caller can tell "the peer closed" from "we cancelled it". + bool fired() const; + +private: + void run(); + + const std::chrono::milliseconds window_; + std::function on_idle_; + + mutable std::mutex mu_; + std::condition_variable cv_; + std::chrono::steady_clock::time_point last_; + bool stop_ = false; + bool fired_ = false; + std::thread thread_; +}; + +// Whether one message read off a live stream is a frame the decoder can +// actually consume, which is the ONLY thing that counts as the peer proving it +// is still there. +// +// Split out of the read loop so the distinction is testable, and because +// getting it wrong is silent. The loop used to touch the watchdog on ANY +// message, before it filtered on has_audio and on an empty pcm field, so a peer +// writing unset-oneof or zero-length frames faster than the window held the +// lane forever: no audio was ever fed, no work was ever done, and the timer +// that exists to break exactly that grip was reset by the frames doing it. +// There is one lane per model and one model per process, so that is a single +// client denying the whole backend. The thrown message already said "no audio +// frame arrived"; this is the code agreeing with it. +bool live_frame_carries_audio(bool has_audio, bool pcm_empty); + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/live_watchdog_test.cpp b/backend/cpp/audio-cpp/live_watchdog_test.cpp new file mode 100644 index 000000000..184655b2b --- /dev/null +++ b/backend/cpp/audio-cpp/live_watchdog_test.cpp @@ -0,0 +1,225 @@ +// Unit tests for the live stream idle watchdog. Standard library only; the +// harness compiles this as a single translation unit, so the implementation is +// included directly. +// +// These are TIMING tests, which is unavoidable: what is under test is a +// deadline. Every window here is short and every assertion waits several +// multiples of it, so a loaded machine slows the test down rather than +// flipping its answer. The one thing never asserted is how SOON something +// happens, only that it eventually does or never does. + +#include "live_watchdog.cpp" + +#include +#include +#include +#include + +using namespace std::chrono_literals; + +static int failures = 0; + +static void check(bool ok, const std::string &name) { + if (!ok) { + failures++; + fprintf(stderr, "FAIL: %s\n", name.c_str()); + } else { + fprintf(stderr, "ok: %s\n", name.c_str()); + } +} + +// A peer that goes quiet without closing. This is the whole point: the lane it +// holds has to come back. +static void test_it_fires_when_nothing_touches_it() { + std::atomic calls{0}; + audiocpp_backend::IdleWatchdog watchdog(100ms, [&calls] { ++calls; }); + std::this_thread::sleep_for(600ms); + check(watchdog.fired(), "a window that elapses untouched fires"); + check(calls.load() == 1, "the callback runs exactly once, not once per window"); +} + +// A peer that is still speaking must never be cut off. Touches land at a third +// of the window, for six windows' worth of wall clock. +static void test_touching_defers_it_indefinitely() { + std::atomic calls{0}; + audiocpp_backend::IdleWatchdog watchdog(300ms, [&calls] { ++calls; }); + for (int i = 0; i < 20; ++i) { + std::this_thread::sleep_for(100ms); + watchdog.touch(); + } + check(!watchdog.fired(), + "a stream touched inside every window is never cancelled"); + check(calls.load() == 0, "no callback runs while the peer is still there"); +} + +// Disarm is what the handler calls when the read side closes, before a decode +// that can take longer than the window. Firing after that would throw away the +// transcript the client is waiting for. +static void test_disarm_stops_it_before_the_window() { + std::atomic calls{0}; + audiocpp_backend::IdleWatchdog watchdog(200ms, [&calls] { ++calls; }); + std::this_thread::sleep_for(20ms); + watchdog.disarm(); + std::this_thread::sleep_for(600ms); + check(!watchdog.fired(), "a disarmed watchdog does not fire"); + check(calls.load() == 0, "a disarmed watchdog runs no callback"); +} + +static void test_disarm_is_idempotent() { + audiocpp_backend::IdleWatchdog watchdog(50ms, [] {}); + watchdog.disarm(); + watchdog.disarm(); + watchdog.disarm(); + check(true, "disarming three times joins once and does not abort"); +} + +// The operator's escape hatch, for a client that legitimately holds a stream +// open through long pauses. Not "expire immediately", which is what a naive +// reading of a zero timeout would give. +static void test_a_non_positive_window_disables_it() { + std::atomic calls{0}; + { + audiocpp_backend::IdleWatchdog watchdog(0ms, [&calls] { ++calls; }); + std::this_thread::sleep_for(300ms); + check(!watchdog.fired(), "a zero window never fires"); + } + { + audiocpp_backend::IdleWatchdog watchdog(-5ms, [&calls] { ++calls; }); + std::this_thread::sleep_for(300ms); + check(!watchdog.fired(), "a negative window never fires"); + } + check(calls.load() == 0, "a disabled watchdog runs no callback"); +} + +// fired() has to survive the disarm, because the handler reads it AFTER the +// read loop ends to tell "the peer closed" from "we cancelled the peer", and +// those two get different statuses. +static void test_fired_survives_a_later_disarm() { + audiocpp_backend::IdleWatchdog watchdog(80ms, [] {}); + std::this_thread::sleep_for(500ms); + watchdog.disarm(); + check(watchdog.fired(), "a watchdog that fired still says so after disarm"); +} + +// The destructor joins, so a callback capturing the handler's frame cannot run +// after that frame is gone. Without the join this is a use after free that only +// shows up under load. +// +// The window is LONGER than the scope on purpose. An earlier version of this +// test slept past the window inside the scope, so the callback had already run +// by the time the object was destroyed and a destructor that DETACHED the thread +// instead of joining it passed unnoticed. Mutation testing is what found that; +// the shape below kills it, because a detached thread wakes after the object is +// gone and calls a callback that must never run. +static void test_the_destructor_joins() { + std::atomic calls{0}; + std::atomic alive{true}; + { + audiocpp_backend::IdleWatchdog watchdog(200ms, [&calls, &alive] { + check(alive.load(), + "the callback never runs after the watched scope ended"); + ++calls; + }); + std::this_thread::sleep_for(20ms); + } + alive.store(false); + std::this_thread::sleep_for(600ms); + check(calls.load() == 0, + "destruction stops the timer rather than leaving it running against a " + "dead frame"); +} + +// The other half of that pair: a callback that DOES fire inside the scope runs +// exactly once, so the test above is not passing merely because nothing ever +// fires. +static void test_a_firing_watchdog_still_joins_cleanly() { + std::atomic calls{0}; + { + audiocpp_backend::IdleWatchdog watchdog(50ms, [&calls] { ++calls; }); + std::this_thread::sleep_for(400ms); + } + check(calls.load() == 1, "the callback ran once, inside the scope"); +} + +// The predicate the live read loop filters on. +static void test_only_a_frame_with_audio_counts() { + using audiocpp_backend::live_frame_carries_audio; + check(live_frame_carries_audio(true, false), + "a frame with a non-empty pcm field carries audio"); + check(!live_frame_carries_audio(true, true), + "an empty pcm field does not"); + check(!live_frame_carries_audio(false, false), + "an unset audio oneof does not, whatever the pcm field looks like"); + check(!live_frame_carries_audio(false, true), + "and neither does an unset oneof with an empty pcm field"); +} + +// The defect this closes, expressed as behaviour rather than as a call order: +// a peer writing frames the decoder cannot consume, faster than the window, +// used to hold the model's only inference lane forever, because the read loop +// touched the watchdog before it filtered them out. One lane per model and one +// model per process, so that is a single client denying the whole backend, +// which is exactly what the watchdog exists to prevent. +static void test_empty_frames_do_not_hold_the_lane() { + std::atomic cancels{0}; + std::atomic stop{false}; + audiocpp_backend::IdleWatchdog watchdog(80ms, [&cancels] { ++cancels; }); + + // The read loop with the real filter in it: frames arrive continuously, + // none of them carries audio, and only a frame that does may touch. + std::thread peer([&] { + while (!stop.load()) { + if (audiocpp_backend::live_frame_carries_audio(false, true)) { + watchdog.touch(); + } + std::this_thread::sleep_for(5ms); + } + }); + + std::this_thread::sleep_for(600ms); + stop.store(true); + peer.join(); + watchdog.disarm(); + + check(cancels.load() == 1, + "a flood of frames with no audio in them still releases the lane"); + + // The mirror image, so this cannot pass merely because the watchdog always + // fires: a peer that keeps sending audio is left alone, exactly as before. + std::atomic live_cancels{0}; + std::atomic live_stop{false}; + audiocpp_backend::IdleWatchdog live(80ms, [&live_cancels] { ++live_cancels; }); + std::thread speaker([&] { + while (!live_stop.load()) { + if (audiocpp_backend::live_frame_carries_audio(true, false)) { + live.touch(); + } + std::this_thread::sleep_for(5ms); + } + }); + std::this_thread::sleep_for(600ms); + live_stop.store(true); + speaker.join(); + live.disarm(); + check(live_cancels.load() == 0, + "a peer that keeps sending audio is never cancelled"); +} + +int main() { + test_it_fires_when_nothing_touches_it(); + test_touching_defers_it_indefinitely(); + test_disarm_stops_it_before_the_window(); + test_disarm_is_idempotent(); + test_a_non_positive_window_disables_it(); + test_fired_survives_a_later_disarm(); + test_the_destructor_joins(); + test_a_firing_watchdog_still_joins_cleanly(); + test_only_a_frame_with_audio_counts(); + test_empty_frames_do_not_hold_the_lane(); + if (failures) { + fprintf(stderr, "%d check(s) failed\n", failures); + return 1; + } + fprintf(stderr, "all live_watchdog checks passed\n"); + return 0; +} diff --git a/backend/cpp/audio-cpp/loaded_model.cpp b/backend/cpp/audio-cpp/loaded_model.cpp new file mode 100644 index 000000000..779a19ab6 --- /dev/null +++ b/backend/cpp/audio-cpp/loaded_model.cpp @@ -0,0 +1,802 @@ +#include "loaded_model.h" + +#include "family_gate.h" + +#include "engine/framework/assets/tensor_source.h" + +#include +#include +#include +#include +#include + +#if defined(__APPLE__) +#include +#include +#include +#elif defined(__linux__) +#include +#endif + +namespace audiocpp_backend { + +// -------------------------------------------------------------------------- +// Enum coupling +// +// audiocpp_backend::Task mirrors engine::runtime::VoiceTaskKind positionally so +// capability_routing can stay stdlib-only and testable without an audio.cpp +// checkout. Nothing about that mirroring is enforced by the type system, and a +// drift is silent in the worst possible way: every unit still compiles, every +// test still passes, and the backend runs a different task than the one the +// caller asked for. +// +// Two mechanisms pin it, and both are needed because they catch different edits: +// +// 1. The assertions below pin every enumerator's value on both sides. An +// insertion or a reorder anywhere before the last member shifts the values +// after it and fails the build here. +// 2. An enumerator APPENDED after the last one shifts nothing, so no value +// assertion can see it. What sees it is the switch in from_engine_task, +// which covers the engine enum with no `default:` label. CMakeLists.txt +// compiles this file with -Werror=switch so that omission is an error +// rather than a warning nobody reads. +// +// Neither mechanism catches a pure RENAME of an upstream enumerator, but that +// does not need catching: the switch stops naming an enumerator that exists and +// the build fails on its own. +// -------------------------------------------------------------------------- + +namespace { + +constexpr int kEngine(engine::runtime::VoiceTaskKind kind) { + return static_cast(kind); +} +constexpr int kMirror(Task task) { return static_cast(task); } + +} // namespace + +static_assert(kEngine(engine::runtime::VoiceTaskKind::Vad) == 0, "VoiceTaskKind drifted"); +static_assert(kEngine(engine::runtime::VoiceTaskKind::Asr) == 1, "VoiceTaskKind drifted"); +static_assert(kEngine(engine::runtime::VoiceTaskKind::Diarization) == 2, "VoiceTaskKind drifted"); +static_assert(kEngine(engine::runtime::VoiceTaskKind::SourceSeparation) == 3, "VoiceTaskKind drifted"); +static_assert(kEngine(engine::runtime::VoiceTaskKind::AudioGeneration) == 4, "VoiceTaskKind drifted"); +static_assert(kEngine(engine::runtime::VoiceTaskKind::Tts) == 5, "VoiceTaskKind drifted"); +static_assert(kEngine(engine::runtime::VoiceTaskKind::VoiceCloning) == 6, "VoiceTaskKind drifted"); +static_assert(kEngine(engine::runtime::VoiceTaskKind::VoiceConversion) == 7, "VoiceTaskKind drifted"); +static_assert(kEngine(engine::runtime::VoiceTaskKind::SpeechToSpeech) == 8, "VoiceTaskKind drifted"); +static_assert(kEngine(engine::runtime::VoiceTaskKind::Alignment) == 9, "VoiceTaskKind drifted"); +static_assert(kEngine(engine::runtime::VoiceTaskKind::VoiceDesign) == 10, "VoiceTaskKind drifted"); +static_assert(kEngine(engine::runtime::VoiceTaskKind::SpeakerRecognition) == 11, "VoiceTaskKind drifted"); +// The last member. Pinning it pins the member count too, as long as the +// enumerators stay contiguous and unassigned, which upstream's declaration is. +static_assert(kEngine(engine::runtime::VoiceTaskKind::Svc) == 12, + "engine::runtime::VoiceTaskKind gained, lost or reordered a member. " + "audiocpp_backend::Task mirrors it positionally: update capability_routing.h, " + "to_engine_task and from_engine_task together, then move this pin."); + +static_assert(kMirror(Task::Vad) == 0, "Task drifted from VoiceTaskKind"); +static_assert(kMirror(Task::Asr) == 1, "Task drifted from VoiceTaskKind"); +static_assert(kMirror(Task::Diarization) == 2, "Task drifted from VoiceTaskKind"); +static_assert(kMirror(Task::SourceSeparation) == 3, "Task drifted from VoiceTaskKind"); +static_assert(kMirror(Task::AudioGeneration) == 4, "Task drifted from VoiceTaskKind"); +static_assert(kMirror(Task::Tts) == 5, "Task drifted from VoiceTaskKind"); +static_assert(kMirror(Task::VoiceCloning) == 6, "Task drifted from VoiceTaskKind"); +static_assert(kMirror(Task::VoiceConversion) == 7, "Task drifted from VoiceTaskKind"); +static_assert(kMirror(Task::SpeechToSpeech) == 8, "Task drifted from VoiceTaskKind"); +static_assert(kMirror(Task::Alignment) == 9, "Task drifted from VoiceTaskKind"); +static_assert(kMirror(Task::VoiceDesign) == 10, "Task drifted from VoiceTaskKind"); +static_assert(kMirror(Task::SpeakerRecognition) == 11, "Task drifted from VoiceTaskKind"); +static_assert(kMirror(Task::Svc) == 12, "Task drifted from VoiceTaskKind"); + +static_assert(static_cast(engine::runtime::RunMode::Offline) == 0, "RunMode drifted"); +static_assert(static_cast(engine::runtime::RunMode::Streaming) == 1, + "engine::runtime::RunMode gained, lost or reordered a member. " + "audiocpp_backend::Mode mirrors it positionally."); +static_assert(static_cast(Mode::Offline) == 0, "Mode drifted from RunMode"); +static_assert(static_cast(Mode::Streaming) == 1, "Mode drifted from RunMode"); + +namespace { + +engine::core::BackendType parse_backend_type(const std::string &value) { + if (value == "cuda") { + return engine::core::BackendType::Cuda; + } + if (value == "vulkan") { + return engine::core::BackendType::Vulkan; + } + if (value == "metal") { + return engine::core::BackendType::Metal; + } + if (value == "best") { + return engine::core::BackendType::BestAvailable; + } + if (value == "cpu" || value.empty()) { + return engine::core::BackendType::Cpu; + } + throw ConfigError("audio-cpp: unknown backend option '" + value + + "'. Known backends: cpu, cuda, vulkan, metal, best"); +} + +std::filesystem::path executable_directory() { +#if defined(__APPLE__) + std::uint32_t size = 0; + _NSGetExecutablePath(nullptr, &size); + std::vector buffer(size + 1, '\0'); + if (_NSGetExecutablePath(buffer.data(), &size) != 0) { + return std::filesystem::current_path(); + } + return std::filesystem::path(buffer.data()).parent_path(); +#elif defined(__linux__) + std::error_code ec; + const auto self = std::filesystem::read_symlink("/proc/self/exe", ec); + if (ec) { + return std::filesystem::current_path(); + } + return self.parent_path(); +#else + return std::filesystem::current_path(); +#endif +} + +// Runs the load gate. Throws ConfigError rather than returning a decision, +// because the only caller is a delegating constructor whose member initializer +// list has nowhere to put a failure. +std::string require_family(const std::string &resolved_path, + const ModelOptions &options) { + const std::filesystem::path path(resolved_path); + std::error_code ec; + if (!std::filesystem::exists(path, ec)) { + throw ConfigError("audio-cpp: model path does not exist: " + resolved_path); + } + + const bool is_gguf = path_looks_like_gguf(resolved_path); + // Only a GGUF is asked for embedded metadata. A directory has no single + // file to read it from, and probing one would make the gate's refusal + // depend on which file happened to be inside. + const std::string embedded = + is_gguf ? read_gguf_family(resolved_path) : std::string(); + + const FamilyDecision decision = + decide_family(is_gguf, embedded, options.family); + if (!decision.ok) { + throw ConfigError(decision.error); + } + return decision.family; +} + +// Refuses a GGUF whose weights are stored in a dtype the family cannot survive. +// +// The POLICY lives in family_gate's weight_dtype_is_supported, which is +// stdlib-only and therefore testable; this is only the part that needs a file +// and an engine to read one. See the table there for why an entry exists and +// what has to be run before deleting it. +// +// Only GGUF paths are inspected. A directory of safetensors carries its dtypes +// per file and has not been tested against this failure, so it is passed +// through rather than guessed at. +void require_supported_weight_dtypes(const std::string &family, + const std::string &resolved_path) { + // Asked as "is there an entry", not as "is the description non-empty": an + // entry with an empty allow list describes a family that can run nothing, + // and reading the description would skip the check on exactly that entry + // while weight_dtype_is_supported refused every dtype. No such entry exists + // today; the two questions are different ones and only one of them is this + // guard's. + if (!family_has_weight_dtype_allow_list(family) || + !path_looks_like_gguf(resolved_path)) { + return; + } + + std::string offending_dtype; + std::string offending_tensor; + try { + const auto source = + engine::assets::open_tensor_source(std::filesystem::path(resolved_path)); + if (source == nullptr) { + return; + } + for (const auto &tensor : source->tensors()) { + if (!weight_dtype_is_supported(family, tensor.dtype)) { + offending_dtype = tensor.dtype; + offending_tensor = tensor.name; + break; + } + } + } catch (const std::exception &) { + // Unreadable as a tensor source. Not this guard's problem to report: + // the registry load below produces a message naming the real fault, and + // refusing here would turn every unusual packaging into this error. + return; + } + + if (offending_dtype.empty()) { + return; + } + throw ConfigError( + "audio-cpp: family '" + family + "' cannot run weights stored as '" + + offending_dtype + "' (tensor '" + offending_tensor + "' in " + + resolved_path + + "); it aborts the backend process on the first request rather than " + "failing the request. Use the 'orig' GGUF package, whose weights are " + + supported_weight_dtypes(family) + "."); +} + +} // namespace + +engine::runtime::VoiceTaskKind to_engine_task(Task task) { + using K = engine::runtime::VoiceTaskKind; + // An explicit switch, never a cast: a cast would keep compiling through + // exactly the drift the assertions above exist to catch. + switch (task) { + case Task::Vad: return K::Vad; + case Task::Asr: return K::Asr; + case Task::Diarization: return K::Diarization; + case Task::SourceSeparation: return K::SourceSeparation; + case Task::AudioGeneration: return K::AudioGeneration; + case Task::Tts: return K::Tts; + case Task::VoiceCloning: return K::VoiceCloning; + case Task::VoiceConversion: return K::VoiceConversion; + case Task::SpeechToSpeech: return K::SpeechToSpeech; + case Task::Alignment: return K::Alignment; + case Task::VoiceDesign: return K::VoiceDesign; + case Task::SpeakerRecognition: return K::SpeakerRecognition; + case Task::Svc: return K::Svc; + } + // Unreachable for any valid enumerator. No `default:` label, so -Wswitch + // still reports a member this switch stops covering. + return K::Vad; +} + +Task from_engine_task(engine::runtime::VoiceTaskKind kind) { + using K = engine::runtime::VoiceTaskKind; + switch (kind) { + case K::Vad: return Task::Vad; + case K::Asr: return Task::Asr; + case K::Diarization: return Task::Diarization; + case K::SourceSeparation: return Task::SourceSeparation; + case K::AudioGeneration: return Task::AudioGeneration; + case K::Tts: return Task::Tts; + case K::VoiceCloning: return Task::VoiceCloning; + case K::VoiceConversion: return Task::VoiceConversion; + case K::SpeechToSpeech: return Task::SpeechToSpeech; + case K::Alignment: return Task::Alignment; + case K::VoiceDesign: return Task::VoiceDesign; + case K::SpeakerRecognition: return Task::SpeakerRecognition; + case K::Svc: return Task::Svc; + } + return Task::Vad; +} + +engine::runtime::RunMode to_engine_mode(Mode mode) { + using M = engine::runtime::RunMode; + switch (mode) { + case Mode::Offline: return M::Offline; + case Mode::Streaming: return M::Streaming; + } + return M::Offline; +} + +Mode from_engine_mode(engine::runtime::RunMode mode) { + using M = engine::runtime::RunMode; + switch (mode) { + case M::Offline: return Mode::Offline; + case M::Streaming: return Mode::Streaming; + } + return Mode::Offline; +} + +Capabilities to_capabilities(const std::string &family, + const engine::runtime::CapabilitySet &set) { + Capabilities caps; + caps.family = family; + caps.tasks.reserve(set.supported_tasks.size()); + for (const auto &supported : set.supported_tasks) { + TaskCapability capability; + capability.task = from_engine_task(supported.task); + capability.modes.reserve(supported.modes.size()); + for (const auto mode : supported.modes) { + capability.modes.push_back(from_engine_mode(mode)); + } + caps.tasks.push_back(std::move(capability)); + } + return caps; +} + +std::string read_gguf_family(const std::string &path) { + try { + const auto spec = + engine::assets::read_gguf_embedded_model_spec(std::filesystem::path(path)); + if (spec.has_value()) { + return spec->family; + } + } catch (...) { + // A file that is not a readable GGUF simply has no family. The load + // gate turns that into a clear refusal; a throw here would surface as + // an opaque internal error during backend probing. + } + return {}; +} + +std::string resolve_model_path(const std::string &model_path_dir, + const std::string &model_file, + const std::string &model_name) { + const std::string candidate = !model_file.empty() ? model_file : model_name; + + // The bundled: form is looked for in BOTH fields, and in ModelOptions.Model + // FIRST, because that is the only field it survives in. LocalAI fills + // ModelFile by joining ModelPath onto the configured model string + // (pkg/model/loader.go, LoadModelWithFile), so a model YAML saying + // `model: bundled:silero_vad` arrives here as ModelFile + // "/models/bundled:silero_vad" and Model "bundled:silero_vad". Testing + // `candidate` alone therefore made the zero-download VAD path reachable only + // from a hand-written LoadModel call that left ModelFile empty, and every + // model YAML using it failed with "model path does not exist". + const std::string bundled_prefix = "bundled:"; + for (const std::string *field : {&model_name, &model_file}) { + if (field->rfind(bundled_prefix, 0) == 0) { + const std::string name = field->substr(bundled_prefix.size()); + return (executable_directory() / "assets" / name).string(); + } + } + + std::filesystem::path path(candidate); + if (path.is_absolute() || model_path_dir.empty()) { + return path.string(); + } + return (std::filesystem::path(model_path_dir) / path).string(); +} + +LoadedModel::LoadedModel(const std::string &resolved_path, + const ModelOptions &options, + std::string model_identity) + : LoadedModel(resolved_path, options, require_family(resolved_path, options), + std::move(model_identity)) {} + +LoadedModel::LoadedModel(const std::string &resolved_path, + const ModelOptions &options, std::string family, + std::string model_identity) + : lane_(family), registry_(engine::runtime::make_default_registry()), + identity_(std::move(model_identity)) { + if (!registry_.supports_family(family)) { + throw ConfigError("audio-cpp: unknown audio.cpp family '" + family + "'"); + } + + // Before the load, for the same reason parse_backend_type runs before it: + // a refusal a metadata read can produce should not cost a full model load. + // More importantly it must precede the FIRST REQUEST, since that is where + // an unsupported dtype aborts the process rather than failing. + require_supported_weight_dtypes(family, resolved_path); + + // Session options are built BEFORE the load, because parse_backend_type + // rejects an unknown backend name. Validating after the load would make + // `backend:cudaa` cost a full model load, on a fault a string comparison + // could have caught. + session_options_.backend.type = parse_backend_type(options.backend); + session_options_.backend.device = options.device; + if (options.threads > 0) { + session_options_.backend.threads = options.threads; + } + for (const auto &entry : options.session_options) { + session_options_.options[entry.first] = entry.second; + } + + pinned_task_ = options.task; + wait_budget_ceiling_ms_ = options.busy_timeout_ms; + live_idle_timeout_ms_ = options.live_idle_timeout_ms; + + engine::runtime::ModelLoadRequest request; + request.model_path = std::filesystem::path(resolved_path); + request.family_hint = family; + if (!options.model_spec_override.empty()) { + request.model_spec_override = + std::filesystem::path(options.model_spec_override); + } + for (const auto &entry : options.load_options) { + request.options[entry.first] = entry.second; + } + + try { + model_ = registry_.load(request); + } catch (const std::exception &err) { + throw ConfigError("audio-cpp: failed to load family '" + family + + "' from " + resolved_path + ": " + err.what()); + } + if (model_ == nullptr) { + throw ConfigError("audio-cpp: the registry returned no model for " + + resolved_path); + } + + const auto &metadata = model_->metadata(); + const auto &engine_caps = model_->capabilities(); + variant_ = metadata.variant; + description_ = metadata.description; + languages_ = engine_caps.languages; + supports_timestamps_ = engine_caps.supports_timestamps; + capabilities_ = to_capabilities(family, engine_caps); +} + +Route LoadedModel::check_can_serve(Rpc rpc, const RequestShape &shape) const { + const Route route = resolve_route(rpc, shape, capabilities_); + if (!route.ok) { + throw CapabilityError(route.error); + } + // Returned so a handler can act on the task before running it. It is the + // same route session_for will resolve, since both read the immutable + // capabilities_ from the same shape. + return route; +} + +LoadedModel::Session LoadedModel::session_for(Rpc rpc, const RequestShape &shape, + LaneEntry &lane) { + // Proof of holding only. Nothing here reads it, and nothing should: its + // whole job is to make a caller that has not taken the lane fail to + // compile. Non-const so it cannot bind to an inline acquire(), whose + // temporary would be released at the end of this call. + (void)lane; + const Route route = resolve_route(rpc, shape, capabilities_); + if (!route.ok) { + throw CapabilityError(route.error); + } + + const SessionKey key{static_cast(route.task), static_cast(route.mode)}; + auto found = sessions_.find(key); + const bool cache_hit = found != sessions_.end(); + if (!cache_hit) { + engine::runtime::TaskSpec spec; + spec.task = to_engine_task(route.task); + spec.mode = to_engine_mode(route.mode); + std::unique_ptr created; + try { + created = model_->create_task_session(spec, session_options_); + } catch (const std::exception &err) { + // NOT a CapabilityError. The family said it supports this pair, and + // a throw from here is overwhelmingly an environment fault: a ggml + // backend .so that package.sh did not ship, an out of memory, a CUDA + // device that is not there. UNIMPLEMENTED would tell LocalAI and + // every client "this model cannot do this, never retry", and send an + // operator hunting a capability bug instead of a packaging one. A + // plain runtime_error maps to INTERNAL, which is what a fixable + // deployment fault should look like. + throw std::runtime_error( + std::string("audio-cpp: family '") + capabilities_.family + + "' advertises " + task_name(route.task) + "/" + + mode_name(route.mode) + " but refused to create the session: " + + err.what()); + } + if (created == nullptr) { + // A null return with no throw is the family declining, which is a + // genuine capability answer and stays UNIMPLEMENTED. + throw CapabilityError(std::string("audio-cpp: family '") + + capabilities_.family + + "' returned no session for " + + task_name(route.task) + "/" + + mode_name(route.mode)); + } + found = sessions_.emplace(key, std::move(created)).first; + } + + Session session; + session.task = route.task; + session.mode = route.mode; + engine::runtime::IVoiceTaskSession *raw = found->second.get(); + if (route.mode == Mode::Streaming) { + session.streaming = + dynamic_cast(raw); + if (session.streaming == nullptr) { + throw CapabilityError(std::string("audio-cpp: family '") + + capabilities_.family + + "' advertises " + task_name(route.task) + + "/streaming but its session is not streaming"); + } + // Deliberately NOT reset here, though a cached streaming session does + // carry state across chunks. reset() is not callable at this point: + // silero_vad's implementation throws "session prepare() must be called + // before Silero VAD reset()", so resetting on a cache hit would turn an + // ordinary second fetch into a hard error, which is worse than the leak + // it would prevent. + // + // The state is instead cleared by the sequence every streaming caller + // owes anyway. IStreamingVoiceTaskSession::start_stream's base + // implementation IS a call to reset(), so a caller that runs + // prepare(...) then start_stream(...) at the top of each stream gets a + // clean session for free. See the STATE CONTRACT in loaded_model.h. + } else { + session.offline = + dynamic_cast(raw); + if (session.offline == nullptr) { + throw CapabilityError(std::string("audio-cpp: family '") + + capabilities_.family + + "' advertises " + task_name(route.task) + + "/offline but its session is not offline"); + } + } + return session; +} + +LaneEntry LoadedModel::acquire(int requested_timeout_ms) { + // Constructed straight into the return value. C++17 requires that, which is + // what lets an immovable type be returned at all; a named local here would + // not compile. + return LaneEntry(lane_, + resolve_wait_budget_ms(wait_budget_ceiling_ms_, + requested_timeout_ms)); +} + +std::unique_ptr LoadedModel::acquire_owned(int requested_timeout_ms) { + return std::make_unique( + lane_, + resolve_wait_budget_ms(wait_budget_ceiling_ms_, requested_timeout_ms)); +} + +engine::runtime::TaskResult run_offline(const LoadedModel::Session &session, + const engine::runtime::TaskRequest &request, + LaneEntry &lane) { + // Proof of holding only, as in session_for. + (void)lane; + if (session.offline == nullptr) { + throw CapabilityError("audio-cpp: no offline session for this request"); + } + session.offline->prepare(engine::runtime::build_preparation_request(request)); + return session.offline->run(request); +} + +namespace { + +engine::runtime::IStreamingVoiceTaskSession & +require_streaming(const LoadedModel::Session &session) { + if (session.streaming == nullptr) { + throw CapabilityError("audio-cpp: no streaming session for this request"); + } + return *session.streaming; +} + +// Clears the stream event sink on every exit from the driver, including the +// exception path. The session is CACHED and outlives the call that installed +// the sink, so a std::function left behind holding references into that call's +// frame is called with dangling captures by whoever streams next. +class ScopedStreamSink { +public: + ScopedStreamSink(engine::runtime::IStreamingVoiceTaskSession &session, + engine::runtime::StreamEventCallback sink) + : session_(session) { + session_.set_stream_event_sink(std::move(sink)); + } + ~ScopedStreamSink() { session_.set_stream_event_sink(nullptr); } + + ScopedStreamSink(const ScopedStreamSink &) = delete; + ScopedStreamSink &operator=(const ScopedStreamSink &) = delete; + +private: + engine::runtime::IStreamingVoiceTaskSession &session_; +}; + +// Frames per chunk to feed a streaming session, from its own policy. +// +// FRAMES, not floats. preferred_audio_chunk_samples is a per-channel count +// everywhere upstream sets it (nemotron_asr uses its frontend sample rate, +// i.e. one second), and vibevoice_asr refuses a chunk whose float count is not +// divisible by its channel count, so slicing on floats would both mis-size the +// window and hand a family a half frame. +std::int64_t chunk_frames_for(const engine::runtime::StreamingPolicy &policy, + int sample_rate) { + if (policy.preferred_audio_chunk_samples > 0) { + return policy.preferred_audio_chunk_samples; + } + // higgs_audio_stt states its window in seconds (4.0) and leaves the sample + // count at zero, so this branch is real rather than defensive. + if (policy.preferred_audio_chunk_seconds > 0.0 && sample_rate > 0) { + const auto frames = static_cast( + policy.preferred_audio_chunk_seconds * static_cast(sample_rate)); + if (frames > 0) { + return frames; + } + } + // The interface's own default, from IStreamingVoiceTaskSession::streaming_policy. + return 512; +} + +} // namespace + +void begin_stream(const LoadedModel::Session &session, + const engine::runtime::TaskRequest &request, LaneEntry &lane) { + // Proof of holding only, as in session_for. + (void)lane; + auto &streaming = require_streaming(session); + // Order is load-bearing: start_stream's reset() is illegal before prepare(). + streaming.prepare(engine::runtime::build_preparation_request(request)); + streaming.start_stream(request); +} + +engine::runtime::TaskResult run_streaming_pull( + const LoadedModel::Session &session, + const engine::runtime::TaskRequest &request, + const std::function &on_event, + LaneEntry &lane) { + auto &streaming = require_streaming(session); + begin_stream(session, request, lane); + while (const auto event = streaming.next_stream_event()) { + if (on_event) { + on_event(*event); + } + // No pinned family sets is_final on a pulled event, so this is not what + // ends the loop today; the nullopt above is. Honoured anyway, because a + // family that does set it is saying the stream is over and pulling once + // more would be asking a finished session for another chunk. + if (event->is_final) { + break; + } + } + return streaming.finish_stream(); +} + +engine::runtime::TaskResult run_streaming_audio( + const LoadedModel::Session &session, + const engine::runtime::TaskRequest &request, + const engine::runtime::AudioBuffer &audio, + const std::function &on_event, + LaneEntry &lane) { + auto &streaming = require_streaming(session); + + const int channels = audio.channels > 0 ? audio.channels : 1; + // REFUSED, not rounded away, and checked before anything is touched so a + // refusal leaves no half-started stream on a cached session. + // + // An interleaved buffer whose float count is not a whole number of frames + // is a truncated input, and the integer division below would silently drop + // the tail floats: they would never be fed, never reach the transcript, and + // nothing would say so. Upstream refuses the same condition rather than + // tolerating it, in two places: vibevoice_asr's audio_frame_count throws + // "VibeVoice-ASR audio samples must be divisible by channel count" + // (session.cpp:70-76), and its process_audio_chunk throws the same with + // "streamed" in the text about the chunks this driver hands it + // (session.cpp:742-747). + // + // ConfigError, i.e. INVALID_ARGUMENT, because the buffer came from the + // caller's file. read_audio_file's positive-rate path always answers mono + // and so cannot reach this, but its native-rate path passes the reader's + // sample count through unchanged, and a driver does not get to assume which + // path its caller took. + if (audio.samples.size() % static_cast(channels) != 0) { + throw ConfigError( + "audio-cpp: streaming input is not a whole number of frames: " + + std::to_string(audio.samples.size()) + " samples across " + + std::to_string(channels) + " channels"); + } + + // Installed BEFORE the stream begins, so a family that reports during + // start_stream is not silently dropped, and destroyed after finish_stream, + // because nemotron_asr emits every one of its partials from inside + // finalize(). + ScopedStreamSink sink(streaming, + [&on_event](const engine::runtime::StreamEvent &event) { + if (on_event) { + on_event(event); + } + }); + + begin_stream(session, request, lane); + + const auto total_frames = + static_cast(audio.samples.size() / static_cast(channels)); + const std::int64_t chunk_frames = + chunk_frames_for(streaming.streaming_policy(), audio.sample_rate); + + for (std::int64_t offset = 0; offset < total_frames; offset += chunk_frames) { + const std::int64_t end = std::min(offset + chunk_frames, total_frames); + engine::runtime::AudioChunk chunk; + chunk.sample_rate = audio.sample_rate; + chunk.channels = channels; + // A FRAME index, which is what every span in a returned event is + // expressed in. vibevoice_asr adds the chunk's own frame count to it to + // offset the spans it reports, so a float index here would place every + // span of a stereo stream at twice its real time. + chunk.start_sample = offset; + chunk.samples.assign( + audio.samples.begin() + static_cast(offset * channels), + audio.samples.begin() + static_cast(end * channels)); + const auto event = streaming.process_audio_chunk(chunk); + if (on_event) { + on_event(event); + } + } + return streaming.finish_stream(); +} + +engine::runtime::TaskResult run_streaming_live( + const LoadedModel::Session &session, + const engine::runtime::TaskRequest &request, + const std::function &)> &next_frames, + const std::function &on_event, + LaneEntry &lane) { + auto &streaming = require_streaming(session); + + // The contract is the only thing that says what rate and layout the frames + // about to arrive are in, and prepare() needs it: see the header. + if (!request.audio_input.has_value()) { + throw ConfigError( + "audio-cpp: a live streaming request carries no audio contract"); + } + const int sample_rate = request.audio_input->sample_rate; + const int channels = + request.audio_input->channels > 0 ? request.audio_input->channels : 1; + + // Installed BEFORE the stream begins and cleared on every exit, including + // the exception path, for the reasons spelled out in run_streaming_audio. + ScopedStreamSink sink(streaming, + [&on_event](const engine::runtime::StreamEvent &event) { + if (on_event) { + on_event(event); + } + }); + + begin_stream(session, request, lane); + + const std::int64_t chunk_frames = + chunk_frames_for(streaming.streaming_policy(), sample_rate); + // chunk_frames_for never returns a non-positive count, so this is never + // zero and the accumulation loop below always terminates. + const std::size_t chunk_floats = static_cast(chunk_frames) * + static_cast(channels); + + std::int64_t fed_frames = 0; + const auto feed = [&](std::vector samples) { + engine::runtime::AudioChunk chunk; + chunk.sample_rate = sample_rate; + chunk.channels = channels; + // A FRAME index, counted across the whole stream: vibevoice_asr offsets + // every span it reports by it, so restarting it per chunk would put + // every word at the top of the recording. + chunk.start_sample = fed_frames; + chunk.samples = std::move(samples); + fed_frames += + static_cast(chunk.samples.size()) / channels; + const auto event = streaming.process_audio_chunk(chunk); + if (on_event) { + on_event(event); + } + }; + + std::vector pending; + std::vector incoming; + while (true) { + incoming.clear(); + if (!next_frames(incoming)) { + break; + } + pending.insert(pending.end(), incoming.begin(), incoming.end()); + while (pending.size() >= chunk_floats) { + std::vector window(pending.begin(), + pending.begin() + + static_cast(chunk_floats)); + pending.erase(pending.begin(), + pending.begin() + + static_cast(chunk_floats)); + feed(std::move(window)); + } + } + + if (!pending.empty()) { + // The tail is whatever did not fill a window. Refused rather than + // truncated when it is not a whole number of frames, exactly as in + // run_streaming_audio: the division above would drop the stray floats + // from the transcript with no diagnostic. Unreachable for a mono live + // stream, which is every live stream today. + if (pending.size() % static_cast(channels) != 0) { + throw ConfigError( + "audio-cpp: live stream ended mid-frame: " + + std::to_string(pending.size()) + " trailing samples across " + + std::to_string(channels) + " channels"); + } + feed(std::move(pending)); + } + + if (fed_frames == 0) { + // Nothing was spoken. See the header: finalizing an empty stream is not + // legal for every family, and an empty transcript is the truthful + // answer rather than an engine-internal INTERNAL. + return engine::runtime::TaskResult{}; + } + return streaming.finish_stream(); +} + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/loaded_model.h b/backend/cpp/audio-cpp/loaded_model.h new file mode 100644 index 000000000..c132fc72f --- /dev/null +++ b/backend/cpp/audio-cpp/loaded_model.h @@ -0,0 +1,378 @@ +#pragma once + +// Owns one audio.cpp model for the life of the process, plus a lazily created +// session per (task, mode) so the same model answers both TTS and TTSStream. +// This is the only unit that converts between the stdlib-only mirror types and +// engine::runtime types. + +#include "capability_routing.h" +#include "inference_lane.h" +#include "model_options.h" + +#include "engine/framework/runtime/model.h" +#include "engine/framework/runtime/registry.h" +#include "engine/framework/runtime/session.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace audiocpp_backend { + +// User-fixable configuration problem. grpc-server maps this to INVALID_ARGUMENT. +class ConfigError : public std::runtime_error { +public: + using std::runtime_error::runtime_error; +}; + +// The family cannot serve the requested RPC. Maps to UNIMPLEMENTED. +class CapabilityError : public std::runtime_error { +public: + using std::runtime_error::runtime_error; +}; + +engine::runtime::VoiceTaskKind to_engine_task(Task task); +engine::runtime::RunMode to_engine_mode(Mode mode); +Task from_engine_task(engine::runtime::VoiceTaskKind kind); +Mode from_engine_mode(engine::runtime::RunMode mode); +Capabilities to_capabilities(const std::string &family, + const engine::runtime::CapabilitySet &set); + +// Reads audiocpp.model_spec.family from a GGUF. Returns an empty string when +// the file is not a GGUF, carries no audio.cpp spec, or cannot be read. Never +// throws: an unreadable file is the load gate's problem, not a crash. +std::string read_gguf_family(const std::string &path); + +// Builds the absolute model path from LocalAI's (ModelPath, ModelFile, Model) +// triple. Either the Model or the ModelFile field may carry the form +// "bundled:", which resolves to /assets/, where +// package.sh puts upstream's bundled silero_vad and marblenet_vad assets. BOTH +// are checked because LocalAI fills ModelFile by joining ModelPath onto the +// configured model string, so a model YAML using the form has it intact only in +// Model. +std::string resolve_model_path(const std::string &model_path_dir, + const std::string &model_file, + const std::string &model_name); + +class LoadedModel { +public: + struct Session { + Task task = Task::Tts; + Mode mode = Mode::Offline; + // Exactly one of these is non-null, matching the resolved mode. + engine::runtime::IOfflineVoiceTaskSession *offline = nullptr; + engine::runtime::IStreamingVoiceTaskSession *streaming = nullptr; + }; + + // Throws ConfigError when the path does not exist, the family cannot be + // determined, or the registry rejects the family. + // + // `model_identity` is ModelOptions.Model verbatim: the UNTRANSLATED + // controller-side name. It is a constructor argument rather than a setter + // so identity and model are inseparable. llama-cpp keeps its equivalent in + // a separate global from the model, which leaves a window where a handler + // can read one without the other; here a handler that holds the model + // through snapshot() necessarily holds the identity it was loaded with. + LoadedModel(const std::string &resolved_path, const ModelOptions &options, + std::string model_identity); + + LoadedModel(const LoadedModel &) = delete; + LoadedModel &operator=(const LoadedModel &) = delete; + + const std::string &family() const noexcept { return capabilities_.family; } + // Empty when the controller predates ModelOptions.ModelIdentity, which the + // identity check reads as "skip". See check_model_identity in grpc-server. + const std::string &identity() const noexcept { return identity_; } + const std::string &variant() const noexcept { return variant_; } + const std::string &description() const noexcept { return description_; } + const std::vector &languages() const noexcept { return languages_; } + const Capabilities &capabilities() const noexcept { return capabilities_; } + bool supports_timestamps() const noexcept { return supports_timestamps_; } + const engine::runtime::SessionOptions &session_options() const noexcept { + return session_options_; + } + + // The model's `task:` option, empty when unset. Every handler must copy it + // into RequestShape::pinned_task before calling session_for: routing is + // otherwise derived from the RPC alone, and this is the option's only route + // from the load to the request that honours it. + const std::string &pinned_task() const noexcept { return pinned_task_; } + + // The `live_idle_timeout_ms` option: how long AudioTranscriptionLive waits + // for the next audio frame before cancelling the stream to give this + // model's lane back. 0 means no limit. See the option in model_options.h + // for why it exists and how the default was chosen. + int live_idle_timeout_ms() const noexcept { return live_idle_timeout_ms_; } + + // Throws the same CapabilityError session_for would throw when this family + // cannot serve the RPC, and RETURNS THE RESOLVED ROUTE otherwise. + // + // The route is returned rather than computed and dropped because a handler + // often has to know which task it is about to run BEFORE running it. + // AudioTransform refuses params[stem] on any route but source separation, + // and reading that off the route costs microseconds where reading it off + // the result costs a whole inference first. A caller with no such need + // ignores the value, which is what the three transcription-shaped handlers + // do. + // + // It exists so a refusal does not have to buy a place in the queue first. + // resolve_route is a pure function of capabilities_, which is fixed at + // construction and never written again, so unlike the session cache it + // needs no lane and no lock: a model that cannot transcribe can say so + // while another request is halfway through a thirty second run. Without + // this the refusal waits for that run to finish only to be told no. + // + // It does NOT replace the routing inside session_for, and must not be made + // to: session_for still needs the route to key the session cache. The two + // calls agree because both read the same immutable capabilities. What this + // one adds is only the ordering, so call it before acquire(). + // + // Const and lane-free on purpose. If a future edit makes routing depend on + // mutable state, this must grow the lane parameter its siblings carry. + Route check_can_serve(Rpc rpc, const RequestShape &shape) const; + + // Routes the RPC and returns the cached session, creating it on first use. + // Throws CapabilityError when this family cannot serve the RPC, and a plain + // runtime_error when it can but the session could not be built, which is an + // environment fault rather than a capability answer. + // + // The `lane` parameter is a PROOF OF HOLDING and is otherwise unused: it + // exists so the rule below is a compile error rather than prose. The + // session cache is an unsynchronised std::map and the sessions themselves + // are not reentrant, so this must only be called with the lane held; the + // lane admits one caller at a time, which is exactly the constraint the + // sessions impose. Pass the LaneEntry from acquire(). + // + // NON-CONST reference on purpose, and do not "tidy" it to const. A const + // reference binds to a temporary, which makes this compile: + // + // auto session = model->session_for(rpc, shape, model->acquire(0)); + // auto result = run_offline(session, task, model->acquire(0)); + // + // and each temporary dies at the end of its own full-expression, so the + // lane is released between the two calls. That is precisely the split this + // parameter exists to prevent, and it is the form a future caller is most + // likely to reach for because it reads as tidy. Requiring an lvalue forces + // a named entry whose scope spans both calls. + // + // What it proves is bounded, so do not over-trust it: it proves A lane was + // taken, not THIS model's lane. A caller determined to defeat it can + // construct an entry on an unrelated InferenceLane and pass that. It + // therefore catches the two mistakes that actually happen, forgetting the + // lane entirely and taking it after routing, and does not catch lane + // identity. + // + // STATE CONTRACT, and it is the CALLER'S to honour. Sessions are cached per + // (task, mode), so a streaming session is normally the same warm object the + // previous stream used, carrying that stream's state. session_for hands it + // back as it is. + // + // Every streaming caller must therefore begin a stream through + // begin_stream() below, which is prepare() then start_stream() in that + // order and is the ONLY implementation of that sequence. start_stream's + // base implementation is a call to reset(), which is what clears the + // previous stream, and reset() is only legal after prepare(): silero_vad + // throws "session prepare() must be called before Silero VAD reset()" + // otherwise. That ordering constraint is also why session_for cannot do + // this for you. Skipping it does not raise an error, it silently continues + // the previous stream. + // + // Offline sessions need no such care: their interface has no reset and + // run() takes a whole request. + Session session_for(Rpc rpc, const RequestShape &shape, LaneEntry &lane); + + // Takes the inference lane, or throws LaneUnavailable. Serializes runs + // against this model. `requested_timeout_ms` is a per-request wait hint + // where a value <= 0 means "use the model's configured ceiling"; a hint may + // only tighten that ceiling, never loosen it. + // + // LaneEntry is deliberately immovable, so bind the result to a named local + // in the scope the inference happens in: + // + // LaneEntry entry = model.acquire(request_hint_ms); + // + // which C++17 initializes in place. A handler that has to keep the lane + // beyond one scope, for instance in a member that outlives the call that + // took it, wants acquire_owned instead. + LaneEntry acquire(int requested_timeout_ms); + + // Same lane, heap-allocated so it can be stored or handed on. Prefer + // acquire: this one adds a null state that the scoped form does not have. + std::unique_ptr acquire_owned(int requested_timeout_ms); + +private: + // Keyed by the enum values so the map needs no custom comparator. + using SessionKey = std::pair; + + // The public constructor runs the load gate, then delegates here. The + // detour exists because lane_ has to be built from the family in the member + // initializer list, and the family is only known after the gate has run. + // Four parameters rather than three so it cannot be confused with the + // public constructor, whose third argument is also a std::string. + LoadedModel(const std::string &resolved_path, const ModelOptions &options, + std::string family, std::string model_identity); + + // MEMBER ORDER IS LOAD-BEARING BELOW THIS LINE. Members are destroyed in + // reverse declaration order. + // + // lane_ is first so it is destroyed last: nothing that runs during teardown + // can then find a lane that has already gone. + InferenceLane lane_; + // registry_ before model_: the registry owns the loader that produced the + // model, and the model may hold loader-owned state. + engine::runtime::ModelRegistry registry_; + // model_ before sessions_, so sessions_ is destroyed FIRST and the model + // second. A session is created from the model and must not outlive it. Do + // not reorder these two. + std::unique_ptr model_; + std::map> sessions_; + + engine::runtime::SessionOptions session_options_; + Capabilities capabilities_; + std::string variant_; + std::string description_; + std::vector languages_; + std::string pinned_task_; + std::string identity_; + bool supports_timestamps_ = false; + int wait_budget_ceiling_ms_ = 0; + int live_idle_timeout_ms_ = 0; +}; + +// Prepares and runs an offline session. prepare() is called for every run +// rather than once per session, because SessionPreparationRequest is derived +// from the request itself (audio contract, text, voice condition) and not from +// the model: a second request with a different sample rate or length would +// otherwise run against the first request's contract. +// +// `lane` is a PROOF OF HOLDING, unused at runtime, for the same reason +// session_for takes one: the session is not reentrant and prepare() mutates it, +// so running without the lane is a data race. Making it a parameter turns that +// into a compile error instead of a comment. Non-const for the same reason as +// session_for's: a const reference would bind to `model.acquire(0)` written +// inline, and that temporary dies at the end of this call, releasing the lane +// before the caller's next one. +// +// Throws CapabilityError when the session is not an offline one. That should be +// unreachable through session_for, which already refuses a non-offline session +// for an offline route, and is checked anyway because the alternative is a null +// dereference. +engine::runtime::TaskResult run_offline(const LoadedModel::Session &session, + const engine::runtime::TaskRequest &request, + LaneEntry &lane); + +// THE ONE IMPLEMENTATION of the streaming state obligation described in +// session_for's STATE CONTRACT: prepare(), then start_stream(), in that order. +// +// It is a function rather than a comment because the obligation is invisible +// when it is broken. Streaming sessions are CACHED per (task, mode), so the +// object a second stream gets is the warm one the first stream left behind, +// still holding its audio, its tokens and its started flag. What clears it is +// start_stream, whose base implementation IS a reset() and whose seven family +// overrides (nemotron_asr, vibevoice_asr, higgs_audio_stt, voxtral_realtime, +// supertonic, omnivoice, voxcpm2) every one call reset() as their first +// statement, verified in the pinned checkout. Nothing in the type system pins +// that. A future override that dropped the reset would break every call site +// at once with no compile error and no exception, only a second transcript +// that begins with the first one's audio, so the fewer call sites there are to +// break, the better: this is the only one. +// +// prepare() must come first and cannot be folded into session_for, because +// reset() is illegal before prepare() (silero_vad throws "session prepare() +// must be called before Silero VAD reset()"), and because the preparation +// request is derived from the REQUEST, not the model: build_preparation_request +// reads the audio contract, the text and the voice condition off it, so a +// second stream with a different sample rate or length would otherwise run +// against the first stream's contract. +// +// `lane` is a PROOF OF HOLDING, unused at runtime, exactly as in run_offline. +// +// Throws CapabilityError when the session is not a streaming one. +void begin_stream(const LoadedModel::Session &session, + const engine::runtime::TaskRequest &request, LaneEntry &lane); + +// Drives a streaming session that takes NO incremental input, which is the TTS +// shape (StreamingInputKind::None, StreamingOutputKind::PullEvents): begin the +// stream, pull events until the session says there are no more, then finish. +// +// NO STREAM EVENT SINK IS INSTALLED HERE, and that is deliberate rather than an +// omission. voxcpm2's start_stream runs the whole synthesis and pushes every +// chunk to the sink, then its next_stream_event replays those same chunks out +// of the stored result, so a sink on this path would put every chunk of audio +// on the wire twice. supertonic and omnivoice ignore set_stream_event_sink +// outright. The pull loop is therefore the single delivery channel. +// +// The returned TaskResult is the session's own merged whole for all three +// families, NOT a tail the pull loop missed. A caller that already emitted the +// pulled events must not also emit its audio; see the TTSStream handler. +engine::runtime::TaskResult run_streaming_pull( + const LoadedModel::Session &session, + const engine::runtime::TaskRequest &request, + const std::function &on_event, + LaneEntry &lane); + +// Drives a streaming session that CONSUMES audio chunks, which is the ASR shape +// (StreamingInputKind::AudioChunks): begin the stream, feed the buffer in +// policy-sized chunks, then finalize. +// +// A STREAM EVENT SINK IS INSTALLED HERE, and it is not optional: nemotron_asr +// reports its partial text ONLY through the sink, and only from inside +// finalize(), because its decode does not start until the audio is complete. +// Without the sink that family streams a transcript with no partials at all. +// The sink is cleared again before returning, including on the exception path: +// the session is cached and outlives this call, so a sink left holding a +// reference to the caller's frame is a use after free waiting for the next +// stream. +// +// Both delivery channels are consumed, the sink and the value process_audio_chunk +// returns, because the families do not agree on which they use, and +// voxtral_realtime uses BOTH for the same event. The duplicate that produces is +// absorbed by TranscriptDeltaTracker in stream_delta.h rather than here. +engine::runtime::TaskResult run_streaming_audio( + const LoadedModel::Session &session, + const engine::runtime::TaskRequest &request, + const engine::runtime::AudioBuffer &audio, + const std::function &on_event, + LaneEntry &lane); + +// Drives the same ASR shape as run_streaming_audio when the audio DOES NOT +// EXIST YET, which is the live-microphone case: instead of slicing a buffer it +// pulls frames from the caller until the input side closes. +// +// `next_frames` fills `out` with interleaved float PCM and returns true, or +// returns false when there is no more input. It is expected to BLOCK, since the +// only real implementation is a gRPC stream Read, and it may throw: a request +// the handler has to refuse mid-stream unwinds through here, and the sink is +// cleared on that path like every other. +// +// The audio contract comes from `request.audio_input`, which for a live stream +// is an EMPTY buffer carrying only the sample rate and channel count. It is not +// optional: nemotron_asr's streaming prepare() throws "Nemotron ASR streaming +// prepare() requires an audio contract" without one, and there is no buffer to +// derive it from here. +// +// Frames are BUFFERED to the family's own preferred window rather than fed in +// whatever sizes the wire delivered them in, because that window is a family's +// statement about what it can decode (nemotron_asr asks for one second, higgs +// for four), and a 512-sample gRPC frame is a property of the client's audio +// callback rather than of the model. The tail shorter than a window is fed at +// the end. +// +// A stream that carried NO AUDIO returns an empty TaskResult and never calls +// finish_stream. Finalizing an empty stream is not universally legal: +// nemotron_asr throws "Nemotron ASR finalize requires streamed audio", so a +// client that opens a session and closes it without speaking would receive an +// INTERNAL naming an engine internal instead of an empty transcript, which is +// the truthful answer to "transcribe nothing". +engine::runtime::TaskResult run_streaming_live( + const LoadedModel::Session &session, + const engine::runtime::TaskRequest &request, + const std::function &)> &next_frames, + const std::function &on_event, + LaneEntry &lane); + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/model_options.cpp b/backend/cpp/audio-cpp/model_options.cpp new file mode 100644 index 000000000..7a4b8c277 --- /dev/null +++ b/backend/cpp/audio-cpp/model_options.cpp @@ -0,0 +1,153 @@ +#include "model_options.h" + +#include +#include +#include +#include + +namespace audiocpp_backend { +namespace { + +std::string trim(const std::string &value) { + size_t begin = 0; + while (begin < value.size() && + std::isspace(static_cast(value[begin])) != 0) { + ++begin; + } + size_t end = value.size(); + while (end > begin && + std::isspace(static_cast(value[end - 1])) != 0) { + --end; + } + return value.substr(begin, end - begin); +} + +// Parses a non-negative integer. Returns false on anything else, including +// empty strings, signs, trailing garbage, and values too large for int. +// +// strtol rather than atoi: atoi is undefined behaviour once the digits exceed +// long, and in practice it hands back a wrapped value. That would let +// "device:2147483648" through as -2147483648 and send a negative index to the +// ggml backend selector, from a function whose error text promises the caller a +// non-negative integer. +bool parse_non_negative_int(const std::string &value, int &out) { + if (value.empty()) { + return false; + } + for (const char ch : value) { + if (std::isdigit(static_cast(ch)) == 0) { + return false; + } + } + + errno = 0; + char *end = nullptr; + const long parsed = std::strtol(value.c_str(), &end, 10); + if (errno == ERANGE || end == nullptr || *end != '\0') { + return false; + } + if (parsed < 0 || parsed > INT_MAX) { + return false; + } + + out = static_cast(parsed); + return true; +} + +bool starts_with(const std::string &value, const std::string &prefix) { + return value.size() >= prefix.size() && + value.compare(0, prefix.size(), prefix) == 0; +} + +} // namespace + +ParsedOptions parse_model_options(const std::vector &entries) { + ParsedOptions parsed; + + for (const auto &raw : entries) { + const std::string entry = trim(raw); + if (entry.empty()) { + continue; + } + + // Split on the FIRST colon: values are often paths that contain more. + const size_t sep = entry.find(':'); + if (sep == std::string::npos) { + parsed.error = "audio-cpp: option '" + entry + + "' is not in key:value form"; + return parsed; + } + + const std::string key = trim(entry.substr(0, sep)); + const std::string value = trim(entry.substr(sep + 1)); + + if (starts_with(key, "load.")) { + const std::string inner = key.substr(5); + if (inner.empty()) { + parsed.error = "audio-cpp: option '" + entry + + "' has an empty load option name"; + return parsed; + } + parsed.options.load_options[inner] = value; + continue; + } + if (starts_with(key, "session.")) { + const std::string inner = key.substr(8); + if (inner.empty()) { + parsed.error = "audio-cpp: option '" + entry + + "' has an empty session option name"; + return parsed; + } + parsed.options.session_options[inner] = value; + continue; + } + + if (key == "family") { + parsed.options.family = value; + } else if (key == "task") { + parsed.options.task = value; + } else if (key == "backend") { + parsed.options.backend = value; + } else if (key == "model_spec_override") { + parsed.options.model_spec_override = value; + } else if (key == "device") { + if (!parse_non_negative_int(value, parsed.options.device)) { + parsed.error = "audio-cpp: option 'device' needs a non-negative " + "integer, got '" + value + "'"; + return parsed; + } + parsed.options.device_set = true; + } else if (key == "threads") { + if (!parse_non_negative_int(value, parsed.options.threads)) { + parsed.error = "audio-cpp: option 'threads' needs a non-negative " + "integer, got '" + value + "'"; + return parsed; + } + } else if (key == "busy_timeout_ms") { + if (!parse_non_negative_int(value, parsed.options.busy_timeout_ms)) { + parsed.error = "audio-cpp: option 'busy_timeout_ms' needs a " + "non-negative integer, got '" + value + "'"; + return parsed; + } + } else if (key == "live_idle_timeout_ms") { + if (!parse_non_negative_int(value, + parsed.options.live_idle_timeout_ms)) { + parsed.error = "audio-cpp: option 'live_idle_timeout_ms' needs a " + "non-negative integer, got '" + value + "'"; + return parsed; + } + } else { + // Quotes the whole entry, not just the key: an entry like ":value" + // has an empty key and would otherwise leave nothing to grep for. + parsed.error = "audio-cpp: unknown option key '" + entry + + "'. Known keys: family, task, backend, device, " + "threads, model_spec_override, busy_timeout_ms, " + "live_idle_timeout_ms, load., session."; + return parsed; + } + } + + return parsed; +} + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/model_options.h b/backend/cpp/audio-cpp/model_options.h new file mode 100644 index 000000000..9cb4464f8 --- /dev/null +++ b/backend/cpp/audio-cpp/model_options.h @@ -0,0 +1,66 @@ +#pragma once + +// Parses the model YAML's `options:` list (ModelOptions.Options in +// backend.proto) into a struct. Standard library only: this unit is compiled +// and tested by backend/cpp/run-unit-tests.sh without an audio.cpp checkout. + +#include +#include +#include + +namespace audiocpp_backend { + +struct ModelOptions { + // audio.cpp model family. Empty means "derive from the GGUF's embedded + // audiocpp.model_spec.family key"; a non-GGUF path with an empty family is + // rejected at load time, not here. + std::string family; + // Pins the audio.cpp task, overriding RPC-based routing. Empty means route. + std::string task; + // ggml backend: cpu, cuda, vulkan, metal, best. + std::string backend = "cpu"; + int device = 0; + // True once a `device:` entry has been seen. 0 is both the default and a + // legitimate device index, so the value alone cannot tell an explicit + // `device:0` from an unset option, and a caller merging in its own fallback + // would silently override the explicit choice. + bool device_set = false; + // 0 means "let the runtime decide". + int threads = 0; + std::string model_spec_override; + // 0 disables the run guard's fail-fast, restoring an unbounded wait. + int busy_timeout_ms = 0; + // How long AudioTranscriptionLive waits for the next audio frame before it + // cancels the stream and gives the model's lane back. 0 means NO LIMIT. + // + // It exists because that RPC holds the lane for the whole stream, so a peer + // that stops sending WITHOUT closing blocks every other request against this + // model for as long as its socket stays up. No other RPC can do that: they + // hold the lane across compute, which ends on its own. + // + // 30 seconds, and the number is picked from what the only in-tree client + // does. core/http/endpoints/openai/realtime.go drives a 300 ms ticker and + // feeds every tick that produced new audio while a turn is open, so 30 s of + // silence is a hundred ticks that delivered nothing: the peer is gone, or + // its socket is wedged. It is also comfortably longer than any pause a + // speaker takes mid-utterance, which is the case that must never be cut off, + // and backend.proto allows one stream to span many utterances, so a client + // that pauses for longer than this between them should raise it rather than + // discover it. Lowering it below a few seconds risks cancelling a live + // speaker; 0 turns the limit off for a client that legitimately idles. + int live_idle_timeout_ms = 30000; + // `load.:` entries, prefix stripped. + std::map load_options; + // `session.:` entries, prefix stripped. + std::map session_options; +}; + +struct ParsedOptions { + ModelOptions options; + // Non-empty means the caller must fail the load with INVALID_ARGUMENT. + std::string error; +}; + +ParsedOptions parse_model_options(const std::vector &entries); + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/model_options_test.cpp b/backend/cpp/audio-cpp/model_options_test.cpp new file mode 100644 index 000000000..c3cb2f75a --- /dev/null +++ b/backend/cpp/audio-cpp/model_options_test.cpp @@ -0,0 +1,165 @@ +// Unit tests for model_options. Standard library only, so +// backend/cpp/run-unit-tests.sh picks this up with no engine checkout. +// +// The harness compiles this file as a single translation unit with no other +// sources, so the implementation is included directly rather than linked. +// +// Build and run standalone: +// g++ -std=c++17 -I. model_options_test.cpp -o t && ./t + +#include "model_options.cpp" + +#include +#include +#include +#include + +static int failures = 0; + +static void check(bool ok, const std::string &name) { + if (!ok) { + failures++; + fprintf(stderr, "FAIL: %s\n", name.c_str()); + } else { + fprintf(stderr, "ok: %s\n", name.c_str()); + } +} + +// Returns the mapped value, or an empty string when the key is absent. map::at +// would throw on a miss and abort the whole binary, so one prefix off-by-one +// would hide every check that follows it instead of failing a single one. +static std::string lookup(const std::map &values, + const std::string &key) { + const auto found = values.find(key); + return found == values.end() ? std::string() : found->second; +} + +using audiocpp_backend::parse_model_options; + +static void test_defaults() { + auto r = parse_model_options({}); + check(r.error.empty(), "empty option list is not an error"); + check(r.options.family.empty(), "family defaults to empty"); + check(r.options.task.empty(), "task defaults to empty"); + check(r.options.backend == "cpu", "backend defaults to cpu"); + check(r.options.device == 0, "device defaults to 0"); + check(r.options.threads == 0, "threads defaults to 0"); + check(r.options.busy_timeout_ms == 0, "busy_timeout_ms defaults to 0"); + // NOT zero, unlike every other numeric option here. A live stream holds the + // model's lane while it waits on the client, so the default has to bound + // that wait; 0 is the explicit "no limit" the operator opts into. + check(r.options.live_idle_timeout_ms == 30000, + "live_idle_timeout_ms defaults to 30000"); + check(r.options.load_options.empty(), "load_options defaults empty"); + check(r.options.session_options.empty(), "session_options defaults empty"); +} + +static void test_scalar_options() { + auto r = parse_model_options({ + "family:qwen3_tts", + "task:tts", + "backend:cuda", + "device:1", + "threads:8", + "busy_timeout_ms:30000", + "live_idle_timeout_ms:5000", + }); + check(r.error.empty(), "scalar options parse without error"); + check(r.options.family == "qwen3_tts", "family parsed"); + check(r.options.task == "tts", "task parsed"); + check(r.options.backend == "cuda", "backend parsed"); + check(r.options.device == 1, "device parsed"); + check(r.options.threads == 8, "threads parsed"); + check(r.options.busy_timeout_ms == 30000, "busy_timeout_ms parsed"); + check(r.options.live_idle_timeout_ms == 5000, "live_idle_timeout_ms parsed"); + check(parse_model_options({"live_idle_timeout_ms:0"}).options.live_idle_timeout_ms == 0, + "an explicit 0 turns the live idle limit off rather than reverting to " + "the default"); +} + +// Values containing colons must survive: split on the FIRST colon only. +static void test_value_containing_colon() { + auto r = parse_model_options({"model_spec_override:/models/a:b/spec.json"}); + check(r.error.empty(), "colon-bearing value is not an error"); + check(r.options.model_spec_override == "/models/a:b/spec.json", + "value keeps every colon after the first separator"); +} + +static void test_namespaced_options() { + auto r = parse_model_options({ + "load.weight_type:q8_0", + "session.miocodec.weight_type:f16", + "session.graph_capacity:tiered", + }); + check(r.error.empty(), "namespaced options parse without error"); + check(r.options.load_options.size() == 1, "one load option"); + check(lookup(r.options.load_options, "weight_type") == "q8_0", "load prefix stripped"); + check(r.options.session_options.size() == 2, "two session options"); + check(lookup(r.options.session_options, "miocodec.weight_type") == "f16", + "session prefix stripped, inner dots kept"); + check(lookup(r.options.session_options, "graph_capacity") == "tiered", + "second session option parsed"); +} + +static void test_errors() { + check(!parse_model_options({"family"}).error.empty(), + "entry without a colon is rejected"); + check(!parse_model_options({"nonsense:1"}).error.empty(), + "unknown key is rejected"); + check(!parse_model_options({"device:abc"}).error.empty(), + "non-numeric device is rejected"); + check(!parse_model_options({"threads:-1"}).error.empty(), + "negative threads is rejected"); + check(!parse_model_options({"load.:x"}).error.empty(), + "empty load key is rejected"); + check(!parse_model_options({"session.:x"}).error.empty(), + "empty session key is rejected"); + check(!parse_model_options({"busy_timeout_ms:abc"}).error.empty(), + "non-numeric busy_timeout_ms is rejected"); + check(!parse_model_options({"live_idle_timeout_ms:abc"}).error.empty(), + "non-numeric live_idle_timeout_ms is rejected"); + check(!parse_model_options({"live_idle_timeout_ms:-1"}).error.empty(), + "negative live_idle_timeout_ms is rejected"); + check(!parse_model_options({"device:-1"}).error.empty(), + "negative device is rejected"); + check(!parse_model_options({"threads:x"}).error.empty(), + "non-numeric threads is rejected"); + + // Values too large for int must be rejected, not silently wrapped into a + // negative device index that then reaches the ggml backend selector. + check(!parse_model_options({"device:2147483648"}).error.empty(), + "device above INT_MAX is rejected"); + check(!parse_model_options({"threads:99999999999999"}).error.empty(), + "threads above INT_MAX is rejected"); + + // The error text must name the offending entry so a user can fix their YAML. + const auto r = parse_model_options({"nonsense:1"}); + check(r.error.find("nonsense") != std::string::npos, + "error names the offending key"); + + // An empty key still has to give the user something to grep for. + const auto empty_key = parse_model_options({":value"}); + check(empty_key.error.find(":value") != std::string::npos, + "unknown-key error names the entry even when the key is empty"); +} + +static void test_blank_entries_ignored() { + auto r = parse_model_options({"", " ", "family:supertonic"}); + check(r.error.empty(), "blank entries are skipped, not rejected"); + check(r.options.family == "supertonic", "real entry still parsed"); +} + +int main() { + test_defaults(); + test_scalar_options(); + test_value_containing_colon(); + test_namespaced_options(); + test_errors(); + test_blank_entries_ignored(); + if (failures) { + fprintf(stderr, "%d check(s) failed\n", failures); + return 1; + } + fprintf(stderr, "all model_options checks passed\n"); + return 0; +} diff --git a/backend/cpp/audio-cpp/package.sh b/backend/cpp/audio-cpp/package.sh new file mode 100755 index 000000000..ba91d3746 --- /dev/null +++ b/backend/cpp/audio-cpp/package.sh @@ -0,0 +1,228 @@ +#!/bin/bash +# Assemble backend/cpp/audio-cpp/package, which becomes the whole content of the +# FROM scratch backend image. Nothing outside this directory exists at run time. +set -euo pipefail + +CURDIR=$(dirname "$(realpath "$0")") +REPO_ROOT="${CURDIR}/../../.." +PACKAGE_DIR="$CURDIR/package" +BUILD_DIR="$CURDIR/build" + +rm -rf "$PACKAGE_DIR" +mkdir -p "$PACKAGE_DIR/lib" "$PACKAGE_DIR/assets" + +cp -avf "$CURDIR/grpc-server" "$PACKAGE_DIR/" +cp -fv "$CURDIR/run.sh" "$PACKAGE_DIR/" + +# ENGINE_ENABLE_CPU_ALL_VARIANTS builds the ggml backends as shared objects that +# are dlopened at run time, so ldd cannot see them and the dependency walk below +# would leave the image with no CPU backend at all. They also cannot go in lib/: +# ggml DISCOVERS them by listing dirname(/proc/self/exe) and the current +# directory, so being on a library path is not enough, they have to be in a +# directory ggml scans. run.sh execs the bundled loader from the package root +# exactly so that directory is this one. cmake writes them to build/bin, not +# next to build/grpc-server, which is why this reads from bin/. +# +# -a keeps the libggml.so -> libggml.so.0 -> libggml.so.0.12.0 symlink chain, +# so the SONAME the binary asks for still names a file here. +for pattern in '*.so*' '*.dylib*'; do + if compgen -G "$BUILD_DIR/bin/$pattern" > /dev/null; then + # shellcheck disable=SC2086 + cp -avf "$BUILD_DIR/bin/"$pattern "$PACKAGE_DIR/" + fi +done + +# Upstream ships silero_vad and marblenet_vad as small runtime assets. +# resolve_model_path() expands "bundled:" to +# dirname(/proc/self/exe)/assets/, so copying them here is what makes VAD +# work with nothing downloaded. +for asset in silero_vad marblenet_vad; do + src="$CURDIR/audio.cpp/assets/framework/models/$asset" + if [ -d "$src" ]; then + cp -rfv "$src" "$PACKAGE_DIR/assets/" + else + echo "package.sh: bundled asset missing: $src" >&2 + echo "package.sh: run 'make audio.cpp' before packaging" >&2 + exit 1 + fi +done + +# Everything below this point is Linux-only: a bundled ELF loader, an ldd walk +# and an ld.so --list validation. The macOS equivalent is the otool -L closure in +# scripts/build/audio-cpp-darwin.sh, which picks up from the exit 0 below. +# +# WARNING FOR ANYONE REWORKING THAT SCRIPT. The obvious move is to copy +# scripts/build/privacy-filter-darwin.sh, and that script assembles its own +# package under build/darwin and never calls package.sh at all. Adapted as-is it +# will silently omit assets/, and the bundled: model path form then resolves to +# nothing, which takes the only zero-download verification path in this backend +# with it. That is why audio-cpp-darwin.sh copies THIS directory instead of +# rebuilding one. The Darwin package needs the same root-level layout as the +# Linux one: grpc-server, run.sh, the ggml dylibs and assets/ in ONE directory, +# with lib/ for the rest. run.sh's Darwin branch execs grpc-server directly, so +# _NSGetExecutablePath already names the package root; nothing else is needed +# beyond putting the files there. +UNAME_S=$(uname -s) +if [ "$UNAME_S" = "Darwin" ]; then + echo "package.sh: Darwin dylib bundling is deferred to scripts/build/audio-cpp-darwin.sh" + ls -lah "$PACKAGE_DIR/" "$PACKAGE_DIR/assets/" + exit 0 +fi + +# The loader goes in the package ROOT, not in lib/. run.sh explains why at +# length; the short version is that exec'ing it makes dirname(/proc/self/exe) +# the directory it sits in, and both the ggml backend scan and the bundled: +# asset lookup need that to be the package root. +if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then + cp -arfLv /lib64/ld-linux-x86-64.so.2 "$PACKAGE_DIR/ld.so" +elif [ -f "/lib/ld-linux-aarch64.so.1" ]; then + cp -arfLv /lib/ld-linux-aarch64.so.1 "$PACKAGE_DIR/ld.so" +else + echo "package.sh: unknown architecture" >&2 + exit 1 +fi + +# THE LAYOUT ASSERTION. Everything else in this script checks that the package +# can LINK. This checks that it can RESOLVE, which is a different property and +# the one with no other guard on it. +# +# The loader, assets/ and the dlopened ggml objects only agree while they share +# one directory, because run.sh execs the loader and all three are reached +# through dirname(/proc/self/exe). A tidy-up that moves the loader into lib/, +# following the llama-cpp layout, produces a package that builds, ships, and +# then fails at run time with "model path does not exist: /lib/assets/..." +# or "Failed to initialize CPU backend". Fail the build instead. +# +# This sits immediately after the loader copy rather than at the end of the +# script on purpose: everything below dereferences $PACKAGE_DIR/ld.so, so a +# misplaced loader would otherwise surface as "No such file or directory" from +# the validation gate and never reach an assertion that could explain it. +if [ ! -f "$PACKAGE_DIR/ld.so" ]; then + echo "package.sh: the bundled loader must be at the package root, not in lib/." >&2 + echo "package.sh: run.sh execs it, so its directory is dirname(/proc/self/exe)," >&2 + echo "package.sh: which is where resolve_model_path looks for assets/ and where" >&2 + echo "package.sh: ggml looks for the CPU variants." >&2 + exit 1 +fi +if [ ! -d "$PACKAGE_DIR/assets" ]; then + echo "package.sh: assets/ must sit beside the loader at the package root." >&2 + exit 1 +fi +# Only assert the ggml half when this build produced CPU variants at all: a +# cublas or vulkan build links ggml statically and ships none. +if compgen -G "$BUILD_DIR/bin/libggml-cpu-*.so" > /dev/null && \ + ! compgen -G "$PACKAGE_DIR/libggml-cpu-*.so" > /dev/null; then + echo "package.sh: the build produced libggml-cpu-*.so but none reached the" >&2 + echo "package.sh: package root, so ggml's scan of dirname(/proc/self/exe)" >&2 + echo "package.sh: will find no CPU backend." >&2 + exit 1 +fi + +# Libraries the host GPU driver stack owns. package_gpu_libs deliberately ships +# the CUDA/Vulkan runtime but not the driver, because the driver has to match +# the kernel module on whatever host runs the image. Copying the build host's +# copy in would pin it to the build host instead. +# +# One regex, used by both the copy loop and the validation gate below. They have +# to agree: exempting a library from the copy but not from the gate makes the +# gate reject the very absence the copy loop just created. +DRIVER_LIB_RE='^(libcuda\.so|libnvidia-)' +# awk applies string-escape processing to a -v assignment before compiling the +# regex, so a lone backslash is eaten and awk warns about it. Double them here +# rather than keeping a second hand-written copy of the pattern, which is the +# drift this single-source-of-truth exists to prevent. +DRIVER_LIB_RE_AWK=${DRIVER_LIB_RE//\\/\\\\} +is_driver_lib() { + [[ "$(basename "$1")" =~ $DRIVER_LIB_RE ]] +} + +# Bundle the full dependency closure. grpc-server links the distro gRPC, +# protobuf and absl stack; copying only the C/C++ runtime leaves the scratch +# image unable to start. The walk runs over the PACKAGED binary, not the one in +# $CURDIR, because its RUNPATH is $ORIGIN: only from inside the package does +# libggml.so.0 resolve to the copy shipped above rather than to nothing. +# The dlopened ggml objects are walked too, since a dependency of theirs that +# grpc-server does not itself link would otherwise be missed. +{ + ldd "$PACKAGE_DIR/grpc-server" + for so in "$PACKAGE_DIR"/*.so*; do + [ -f "$so" ] || continue + ldd "$so" + done +} | awk '$2 == "=>" && $3 ~ /^\// { print $3 }' | sort -u | \ +while read -r so; do + # Skip what is already inside the package: the ggml objects resolve through + # $ORIGIN and re-copying them into lib/ would ship two copies of each. + case "$so" in "$PACKAGE_DIR"/*) continue ;; esac + if is_driver_lib "$so"; then + echo "package.sh: leaving driver-owned library to the host: $so" + continue + fi + cp -arfLv "$so" "$PACKAGE_DIR/lib/" +done + +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}..." + # shellcheck source=/dev/null + source "$GPU_LIB_SCRIPT" "$PACKAGE_DIR/lib" + package_gpu_libs +fi + +# Resolve every dependency through the same loader and library path the +# from-scratch image uses. Two distinct failures are rejected, because the +# loader can still fall back to the host's default directories: a dependency it +# could not resolve at all, and one it resolved to a file OUTSIDE the package, +# which would validate here and be absent in the image. +# +# The driver libraries are exempt from BOTH rejections, and that exemption is +# load-bearing on GPU builds rather than tidiness. With BUILD_TYPE=cublas ggml +# is static (no CPU_ALL_VARIANTS), and ggml/CMakeLists.txt defaults +# GGML_CUDA_NO_VMM=OFF, so ggml-cuda links CUDA::cuda_driver and grpc-server +# itself carries DT_NEEDED libcuda.so.1. The copy loop above deliberately leaves +# that to the host, so inside the CUDA builder it resolves either to a host path +# or to nothing. Without this exemption every cublas build would fail here and +# CI would produce no image at all. +# +# LD_TRACE_LOADED_OBJECTS + LD_LIBRARY_PATH, NOT `ld.so --library-path --list`, +# and the difference is not cosmetic. Measured on a stub object built to +# DT_NEEDED an absent libcuda.so.1: `--list` refuses to trace at all, printing +# "libdrivertest.so: error while loading shared libraries: libcuda.so.1: cannot +# open shared object file" and exiting 127, so no per-library line is ever +# produced and no exemption below could apply. The env form prints +# "libcuda.so.1 => not found" and exits 0, which is what makes both the +# unresolved rule and its driver exemption reachable. It is also closer to what +# run.sh actually does, since run.sh exports LD_LIBRARY_PATH rather than passing +# --library-path. +validation_failed=0 +validate_object() { + local object="$1" + LD_TRACE_LOADED_OBJECTS=1 LD_LIBRARY_PATH="$PACKAGE_DIR/lib:$PACKAGE_DIR" \ + "$PACKAGE_DIR/ld.so" "$object" | awk -v pkg="$PACKAGE_DIR/" -v obj="$object" \ + -v driver_re="$DRIVER_LIB_RE_AWK" ' + function base(p, n, parts) { n = split(p, parts, "/"); return parts[n] } + $2 == "=>" && $3 == "not" { + if ($1 ~ driver_re) next + print "package.sh: unresolved dependency of " obj ": " $1 > "/dev/stderr" + bad = 1 + } + $2 == "=>" && $3 ~ /^\// && index($3, pkg) != 1 { + if (base($3) ~ driver_re) next + print "package.sh: dependency of " obj " resolved outside the package: " $0 > "/dev/stderr" + bad = 1 + } + END { exit bad } + ' +} + +validate_object "$PACKAGE_DIR/grpc-server" || validation_failed=1 +for so in "$PACKAGE_DIR"/*.so*; do + [ -f "$so" ] || continue + validate_object "$so" || validation_failed=1 +done +if [ "$validation_failed" -ne 0 ]; then + exit 1 +fi + +echo "audio-cpp package contents:" +ls -lah "$PACKAGE_DIR/" "$PACKAGE_DIR/lib/" "$PACKAGE_DIR/assets/" diff --git a/backend/cpp/audio-cpp/result_map.cpp b/backend/cpp/audio-cpp/result_map.cpp new file mode 100644 index 000000000..c9498d1ae --- /dev/null +++ b/backend/cpp/audio-cpp/result_map.cpp @@ -0,0 +1,85 @@ +#include "result_map.h" + +#include "transcript_assembly.h" + +#include +#include + +namespace audiocpp_backend { + +void fill_transcript_result(const engine::runtime::TaskResult &result, + int sample_rate, float duration_seconds, + backend::TranscriptResult *out) { + // No null guard on `out`, deliberately. gRPC always hands a handler a + // response message, so a null here would be a programming error in a + // caller, and a guard that returned quietly would answer the client with an + // untouched, empty transcript and an OK status. That is the same + // indistinguishable-from-silence failure the rest of this unit exists to + // prevent; crashing on the developer's machine is the cheaper outcome. + std::vector speech_segments; + speech_segments.reserve(result.speech_segments.size()); + for (const auto &segment : result.speech_segments) { + speech_segments.push_back( + Span{segment.span.start_sample, segment.span.end_sample}); + } + + std::vector speaker_turns; + speaker_turns.reserve(result.speaker_turns.size()); + for (const auto &turn : result.speaker_turns) { + speaker_turns.push_back( + SpeakerSpan{Span{turn.span.start_sample, turn.span.end_sample}, + turn.speaker_id}); + } + + std::vector words; + words.reserve(result.word_timestamps.size()); + for (const auto &word : result.word_timestamps) { + words.push_back( + WordSpan{Span{word.span.start_sample, word.span.end_sample}, + word.word}); + } + + // The ONLY read of transcript text in this function, and the only one there + // may ever be. See THE RULE in the header. + const std::string text = + result.text_output.has_value() ? result.text_output->text : std::string(); + + const AssembledTranscript assembled = assemble_transcript( + text, speech_segments, speaker_turns, words, sample_rate); + + out->set_text(assembled.text); + // language has no source inside transcript_assembly, which is span-shaped + // only, so it is read straight off the engine result here. Left untouched + // when the family reported no text output at all: an empty string would be + // indistinguishable from a family that genuinely detected no language, and + // the field is documented as optional. + if (result.text_output.has_value()) { + out->set_language(result.text_output->language); + } + out->set_duration(duration_seconds); + + // Cleared rather than appended to. A caller that fills the same message + // twice (a stream's final_result being rebuilt, say) would otherwise emit + // every segment twice, and the second call's ids would restart at 0 and + // collide with the first call's. + out->clear_segments(); + for (const auto &segment : assembled.segments) { + auto *out_segment = out->add_segments(); + out_segment->set_id(segment.id); + // NANOSECONDS. TranscriptSegment and TranscriptWord are the only + // messages in backend.proto that use them; VADSegment and DiarizeSegment + // are float seconds. assemble_transcript has already converted. + out_segment->set_start(segment.start_ns); + out_segment->set_end(segment.end_ns); + out_segment->set_text(segment.text); + out_segment->set_speaker(segment.speaker); + for (const auto &word : segment.words) { + auto *out_word = out_segment->add_words(); + out_word->set_start(word.start_ns); + out_word->set_end(word.end_ns); + out_word->set_text(word.text); + } + } +} + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/result_map.h b/backend/cpp/audio-cpp/result_map.h new file mode 100644 index 000000000..97965877c --- /dev/null +++ b/backend/cpp/audio-cpp/result_map.h @@ -0,0 +1,37 @@ +#pragma once + +// Converts engine::runtime results into LocalAI proto messages. All of the +// non-trivial shaping lives in transcript_assembly, which is stdlib-only and +// unit tested; this unit is the thin engine-typed boundary around it. + +#include "backend.pb.h" + +#include "engine/framework/runtime/session.h" + +namespace audiocpp_backend { + +// Fills text, language, duration, segments and per-segment words. +// +// THE RULE: the top-level text is TaskResult.text_output verbatim. It is never +// derived from segments or words. audio.cpp carries transcript text in +// text_output and nowhere else: speech_segments, speaker_turns and +// word_timestamps carry spans and labels and no text at all. Deriving the +// transcript from them therefore returns an EMPTY text for every producer that +// reports segments without word timing, which real VibeVoice diarized ASR does. +// An earlier attempt at this backend shipped exactly that bug. assemble_transcript +// enforces the rule and is heavily tested; this unit's job is not to re-derive +// it but to not undo it at the proto boundary. +// +// `sample_rate` is the rate the result's spans are expressed in, which is the +// rate of the AudioBuffer that was handed to the session, NOT the rate of the +// file the caller uploaded. Those differ whenever read_audio_file resampled, +// which is why the handler passes the buffer's rate rather than the file's. +// +// Segments are replaced, not appended to, so a message filled twice does not +// accumulate. `out` must be non-null and is not checked; see the note at the +// top of the implementation for why that is not an oversight. +void fill_transcript_result(const engine::runtime::TaskResult &result, + int sample_rate, float duration_seconds, + backend::TranscriptResult *out); + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/result_map_ctest.cpp b/backend/cpp/audio-cpp/result_map_ctest.cpp new file mode 100644 index 000000000..fc12016da --- /dev/null +++ b/backend/cpp/audio-cpp/result_map_ctest.cpp @@ -0,0 +1,247 @@ +// Tests for result_map, the engine-to-proto boundary. +// +// NAMED _ctest AND NOT _test ON PURPOSE. backend/cpp/run-unit-tests.sh globs +// every *_test.cpp under backend/cpp/ and compiles it as a single standalone +// translation unit with no include path beyond its own directory. This file +// needs backend.pb.h and the audio.cpp framework headers, so it is built and +// run by ctest instead: +// +// make -C backend/cpp/audio-cpp test-engine +// +// Renaming it to *_test.cpp would break the standalone suite for every backend. +// +// What is worth testing here is exactly one thing, and it is not the field +// copying: THE RULE. TaskResult carries transcript text in text_output and +// nowhere else, so the proto's text must be that string verbatim. An earlier +// attempt at this backend derived it from the segments, which returns an empty +// transcript for every producer that reports segments without word timing. +// transcript_assembly already enforces the rule and is tested on its own; these +// checks are here so that a future edit cannot undo it at the boundary. + +#include "result_map.h" + +#include +#include + +static int failures = 0; + +static void check(bool ok, const std::string &name) { + if (!ok) { + failures++; + fprintf(stderr, "FAIL: %s\n", name.c_str()); + } else { + fprintf(stderr, "ok: %s\n", name.c_str()); + } +} + +using namespace audiocpp_backend; +namespace rt = engine::runtime; + +static const int kRate = 16000; + +static rt::SpeechSegment speech(std::int64_t start, std::int64_t end) { + rt::SpeechSegment segment; + segment.span.start_sample = start; + segment.span.end_sample = end; + return segment; +} + +static rt::SpeakerTurn turn(std::int64_t start, std::int64_t end, + const std::string &speaker) { + rt::SpeakerTurn out; + out.span.start_sample = start; + out.span.end_sample = end; + out.speaker_id = speaker; + return out; +} + +static rt::WordTimestamp word(std::int64_t start, std::int64_t end, + const std::string &text) { + rt::WordTimestamp out; + out.span.start_sample = start; + out.span.end_sample = end; + out.word = text; + return out; +} + +// THE REGRESSION. A diarized ASR result: real text, real speaker turns, and no +// word timing at all. This is the vibevoice_asr shape, and it is the one that +// came back empty before. +static void test_text_survives_segments_without_words() { + rt::TaskResult result; + rt::Transcript transcript; + transcript.text = "hello there general kenobi"; + transcript.language = "en"; + result.text_output = transcript; + result.speaker_turns.push_back(turn(0, 16000, "speaker_0")); + result.speaker_turns.push_back(turn(16000, 32000, "speaker_1")); + + backend::TranscriptResult out; + fill_transcript_result(result, kRate, 2.0f, &out); + + check(out.text() == "hello there general kenobi", + "diarized result keeps text_output verbatim"); + check(out.language() == "en", "language comes from text_output"); + check(out.segments_size() == 2, "both speaker turns become segments"); + if (out.segments_size() == 2) { + check(out.segments(0).speaker() == "speaker_0", + "first segment keeps its own speaker label"); + check(out.segments(1).speaker() == "speaker_1", + "second segment keeps its own speaker label"); + check(out.segments(1).start() == 1000000000LL, + "segment start is nanoseconds, not samples"); + check(out.segments(1).end() == 2000000000LL, + "segment end is nanoseconds, not samples"); + } +} + +// The same rule seen from the other side: text present, spans present, and the +// per-segment text empty because there is nothing truthful to split. A boundary +// that derived the top-level text from these segments would produce "". +static void test_speech_segments_do_not_supply_the_text() { + rt::TaskResult result; + rt::Transcript transcript; + transcript.text = "one two three"; + result.text_output = transcript; + result.speech_segments.push_back(speech(0, 8000)); + result.speech_segments.push_back(speech(8000, 16000)); + + backend::TranscriptResult out; + fill_transcript_result(result, kRate, 1.0f, &out); + + check(out.text() == "one two three", + "speech segments without words do not empty the transcript"); + check(out.segments_size() == 2, "both speech segments are emitted"); + if (out.segments_size() == 2) { + check(out.segments(0).text().empty() && out.segments(1).text().empty(), + "per-segment text stays empty when there is no word timing"); + } +} + +static void test_words_reach_the_proto_in_nanoseconds() { + rt::TaskResult result; + rt::Transcript transcript; + transcript.text = "hi there"; + result.text_output = transcript; + result.word_timestamps.push_back(word(0, 8000, "hi")); + result.word_timestamps.push_back(word(8000, 16000, "there")); + + backend::TranscriptResult out; + fill_transcript_result(result, kRate, 1.0f, &out); + + check(out.text() == "hi there", "word-timed result keeps text_output"); + check(out.segments_size() == 1, "words with no spans yield one covering segment"); + if (out.segments_size() == 1) { + const auto &segment = out.segments(0); + check(segment.words_size() == 2, "both words are emitted"); + if (segment.words_size() == 2) { + check(segment.words(0).text() == "hi", "first word text"); + check(segment.words(0).start() == 0, "first word start"); + check(segment.words(0).end() == 500000000LL, + "first word end is 0.5 s in nanoseconds"); + check(segment.words(1).start() == 500000000LL, "second word start"); + check(segment.words(1).end() == 1000000000LL, "second word end"); + } + } +} + +// The buffer's rate, not the file's, is what the spans mean. Passing 8000 for +// the same spans has to halve every timestamp, which is what makes resampling +// the input at read time load-bearing rather than cosmetic. +static void test_sample_rate_scales_the_timestamps() { + rt::TaskResult result; + rt::Transcript transcript; + transcript.text = "x"; + result.text_output = transcript; + result.speech_segments.push_back(speech(0, 8000)); + + backend::TranscriptResult out; + fill_transcript_result(result, 8000, 1.0f, &out); + + check(out.segments_size() == 1, "one segment at 8 kHz"); + if (out.segments_size() == 1) { + check(out.segments(0).end() == 1000000000LL, + "8000 samples at 8 kHz is one second"); + } +} + +static void test_duration_is_carried_through() { + rt::TaskResult result; + rt::Transcript transcript; + transcript.text = "x"; + result.text_output = transcript; + + backend::TranscriptResult out; + fill_transcript_result(result, kRate, 14.07f, &out); + + check(out.duration() > 14.06f && out.duration() < 14.08f, + "duration is set from the argument"); +} + +// No text output at all. A VAD-shaped result reaching this boundary must not +// invent a transcript, and must not overwrite a language the caller had already +// decided on. +static void test_missing_text_output_leaves_language_alone() { + rt::TaskResult result; + result.speech_segments.push_back(speech(0, 16000)); + + backend::TranscriptResult out; + out.set_language("it"); + fill_transcript_result(result, kRate, 1.0f, &out); + + check(out.text().empty(), "no text_output means no text"); + check(out.language() == "it", + "a result with no text_output does not clear the language"); + check(out.segments_size() == 1, "spans are still emitted"); +} + +// Filling the same message twice must replace, not accumulate: the second +// call's ids restart at 0 and would collide with the first call's. +static void test_refilling_replaces_the_segments() { + rt::TaskResult first; + rt::Transcript transcript; + transcript.text = "first"; + first.text_output = transcript; + first.speech_segments.push_back(speech(0, 16000)); + first.speech_segments.push_back(speech(16000, 32000)); + + backend::TranscriptResult out; + fill_transcript_result(first, kRate, 2.0f, &out); + + rt::TaskResult second; + rt::Transcript replacement; + replacement.text = "second"; + second.text_output = replacement; + second.speech_segments.push_back(speech(0, 16000)); + fill_transcript_result(second, kRate, 1.0f, &out); + + check(out.text() == "second", "the second fill replaces the text"); + check(out.segments_size() == 1, + "the second fill replaces the segments instead of appending"); +} + +static void test_empty_result_is_empty() { + rt::TaskResult result; + backend::TranscriptResult out; + fill_transcript_result(result, kRate, 0.0f, &out); + + check(out.text().empty(), "empty result has no text"); + check(out.segments_size() == 0, "empty result has no segments"); +} + +int main() { + test_text_survives_segments_without_words(); + test_speech_segments_do_not_supply_the_text(); + test_words_reach_the_proto_in_nanoseconds(); + test_sample_rate_scales_the_timestamps(); + test_duration_is_carried_through(); + test_missing_text_output_leaves_language_alone(); + test_refilling_replaces_the_segments(); + test_empty_result_is_empty(); + if (failures) { + fprintf(stderr, "%d check(s) failed\n", failures); + return 1; + } + fprintf(stderr, "all result_map checks passed\n"); + return 0; +} diff --git a/backend/cpp/audio-cpp/run.sh b/backend/cpp/audio-cpp/run.sh new file mode 100755 index 000000000..2ecd56761 --- /dev/null +++ b/backend/cpp/audio-cpp/run.sh @@ -0,0 +1,58 @@ +#!/bin/bash +# Entry point for the audio-cpp backend image and for BACKEND_BINARY mode. +# +# The image's final stage is FROM scratch, so the package root is / and there is +# no system loader, no system libc and no fallback library path. Everything the +# process opens has to be inside the package, and it has to be findable by the +# two mechanisms that actually do the finding: the dynamic linker, and +# audio.cpp's own directory scans. +set -e + +CURDIR=$(dirname "$(realpath "$0")") + +if [ "$(uname -s)" = "Darwin" ]; then + export DYLD_LIBRARY_PATH="$CURDIR/lib:$CURDIR:$DYLD_LIBRARY_PATH" + exec "$CURDIR/grpc-server" "$@" +fi + +# $CURDIR is on the path as well as $CURDIR/lib: the ggml shared objects the +# CPU-all-variants build produces sit in the package root, next to the binary, +# not in lib/. See the comment below for why they cannot live in lib/. +export LD_LIBRARY_PATH="$CURDIR/lib:$CURDIR:$LD_LIBRARY_PATH" + +# THE BUNDLED LOADER IS AT THE PACKAGE ROOT, NOT AT lib/ld.so. DO NOT MOVE IT. +# +# Exec'ing the loader is what pins the bundled glibc to the matching ld.so, and +# every other C++ backend here does it. The cost is that /proc/self/exe then +# names the LOADER rather than grpc-server, and this backend has two consumers +# of /proc/self/exe that both have to land on the package root: +# +# - ggml's backend registry DISCOVERS the per-microarch libggml-cpu-*.so by +# scanning dirname(/proc/self/exe) and the current directory. Those files +# are dlopened, never linked, so no library path and no RUNPATH reaches +# them: they have to be in a directory ggml scans. +# - resolve_model_path() turns "bundled:" into +# dirname(/proc/self/exe)/assets/, which is how the bundled +# silero_vad and marblenet_vad models resolve with nothing downloaded. +# +# backend/cpp/llama-cpp/package.sh answers the first of these by keeping +# lib/ld.so and moving the ggml objects INTO lib/. That does not generalise +# here, because it would also drag assets/ into lib/ to keep the second +# consumer working. Putting the loader in the package root instead makes +# dirname(/proc/self/exe) the package root, so the binary, the ggml objects and +# assets/ all sit in the one directory that all three mechanisms agree on. +# +# The ggml half has a second chance that the bundled: half does not: LocalAI +# sets the backend process cwd to the directory holding run.sh +# (pkg/model/process.go), so ggml's fs::current_path() fallback would find the +# objects in normal operation whatever the loader's placement. That fallback is +# worth little here. It holds only for the launcher that sets that cwd, it is +# gone the moment anyone runs the binary by hand or through a wrapper that +# chdirs, and resolve_model_path has no equivalent, which would leave the only +# zero-download path in this backend resting on it. Rooting the loader is the +# one layout where all three mechanisms agree without depending on the cwd. +if [ -f "$CURDIR/ld.so" ]; then + exec "$CURDIR/ld.so" "$CURDIR/grpc-server" "$@" +fi + +exec "$CURDIR/grpc-server" "$@" diff --git a/backend/cpp/audio-cpp/stem_selection.cpp b/backend/cpp/audio-cpp/stem_selection.cpp new file mode 100644 index 000000000..da944e6c2 --- /dev/null +++ b/backend/cpp/audio-cpp/stem_selection.cpp @@ -0,0 +1,129 @@ +#include "stem_selection.h" + +#include +#include +#include + +namespace audiocpp_backend { +namespace { + +// The stem every separation family this backend can reach names its lead vocal +// track, and the one a caller who names no stem almost always wants: it is what +// the OpenAI-shaped "isolate the voice" request means. htdemucs (drums, bass, +// other, vocals) and mel_band_roformer (vocals, instrumental) both have it, and +// in htdemucs's case it is NOT the first output, which is the whole reason this +// preference is written down rather than left as "take index 0". +const char *const kPreferredStem = "vocals"; + +std::string join_names(const std::vector &names) { + std::string out; + for (const auto &name : names) { + if (!out.empty()) { + out += ", "; + } + out += name; + } + return out; +} + +// Whether a model-supplied stem name can be used as one component of a file +// name. Deliberately a whitelist of refusals rather than a sanitiser: silently +// rewriting "vo/cals" to "cals" would make the file the caller receives +// disagree with the name they would have to ask for. +bool name_is_writable(const std::string &name) { + if (name.empty() || name == "." || name == "..") { + return false; + } + // Control bytes, NUL above all. GGUF strings are length prefixed and demucs + // reads its source names out of JSON, which can encode one, so a + // std::string holding an embedded NUL survives all the way here. Two such + // names differing only AFTER the NUL are distinct std::strings, so the + // duplicate check below waves them through, and then path::c_str() + // truncates both at the NUL and they open the same file: precisely the + // silent overwrite the duplicate check exists to prevent, with the ".wav" + // stripped off as well. The rest of the range goes with it, since a newline + // or an escape sequence in a file name is a terminal and log injection + // nuisance with no legitimate use. + for (const char byte : name) { + const auto value = static_cast(byte); + if (value < 0x20 || value == 0x7f) { + return false; + } + } + // Both separators, not just the host's. A GGUF is a downloaded file and its + // strings are not this host's to trust, so a name written on Windows must + // not become a directory traversal wherever the check happens to run. + return name.find('/') == std::string::npos && + name.find('\\') == std::string::npos; +} + +} // namespace + +StemChoice select_named_output(const std::vector &names, + const std::string &requested) { + StemChoice choice; + if (names.empty()) { + // Not an error here. The caller distinguishes "this family produces one + // unnamed output" from "this family produced nothing", and only it can + // tell them apart. + return choice; + } + + // Every name is checked, not merely the selected one, because every stem is + // written. A bad name in the fourth output would otherwise be discovered + // only after three files had already been created. + for (std::size_t i = 0; i < names.size(); ++i) { + if (!name_is_writable(names[i])) { + choice.error = "audio-cpp: this model names an output stem '" + + names[i] + + "' that cannot be used as a file name; stems: " + + join_names(names); + return choice; + } + for (std::size_t seen = 0; seen < i; ++seen) { + if (names[seen] == names[i]) { + choice.error = + "audio-cpp: this model produces two output stems both named '" + + names[i] + "'; one would silently overwrite the other"; + return choice; + } + } + } + + if (!requested.empty()) { + for (std::size_t i = 0; i < names.size(); ++i) { + if (names[i] == requested) { + choice.index = static_cast(i); + return choice; + } + } + choice.error = "audio-cpp: no stem named '" + requested + + "' in this model's output; available stems: " + + join_names(names); + return choice; + } + + for (std::size_t i = 0; i < names.size(); ++i) { + if (names[i] == kPreferredStem) { + choice.index = static_cast(i); + return choice; + } + } + choice.index = 0; + return choice; +} + +std::string sibling_stem_path(const std::string &dst, const std::string &name) { + const std::filesystem::path path(dst); + // The dst extension is reused rather than forced to ".wav" so the siblings + // look like the file the caller named. write_audio_file writes WAV bytes + // whatever the extension says, for dst as much as for the siblings, so this + // keeps the set consistent instead of making the siblings honest about a + // format dst is already lying about. + const std::string extension = + path.has_extension() ? path.extension().string() : std::string(".wav"); + return (path.parent_path() / (path.stem().string() + "." + name + extension)) + .string(); +} + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/stem_selection.h b/backend/cpp/audio-cpp/stem_selection.h new file mode 100644 index 000000000..f889639c5 --- /dev/null +++ b/backend/cpp/audio-cpp/stem_selection.h @@ -0,0 +1,63 @@ +#pragma once + +// Decides which of a separation model's named stems the AudioTransform response +// carries in `dst`, and names the sibling files every other stem is written to. +// +// It exists because AudioTransformResult carries ONE dst while htdemucs and +// mel_band_roformer produce several named outputs from a single run. Running +// once per stem would cost four full inferences for a four stem model, so the +// handler runs once, writes every stem beside dst, and puts the selected one in +// dst itself. +// +// Standard library only, so backend/cpp/run-unit-tests.sh compiles and runs its +// test without an audio.cpp checkout. grpc-server.cpp flattens the engine's +// NamedAudioBuffer list into the plain name vector taken here; nothing in this +// unit knows about engine::runtime. + +#include +#include + +namespace audiocpp_backend { + +struct StemChoice { + // Index into the `names` vector. -1 means nothing was chosen, which happens + // for an empty list (the family produced a single unnamed output) and for + // every refusal. + // + // THE CONTRACT THE CALLER INDEXES ON: when `error` is empty and `names` was + // not, this is always a valid index into `names`. It is never -1 in that + // case, so a caller that checks `error` first can index without a further + // guard, and a caller that does not check `error` first would index with a + // negative value. Check the error. + int index = -1; + // Non-empty when the request must be refused, and suitable verbatim as an + // INVALID_ARGUMENT message. Two things land here: the caller named a stem + // this model does not produce, and the model named stems that cannot both + // be written (an unusable file name, or two stems sharing one). + std::string error; +}; + +// Picks the stem that goes to dst. Preference order: an explicit `requested`, +// then "vocals", then the first output. +// +// An explicit but unknown `requested` is an ERROR rather than a fallback. A +// caller who asks for "drums" and silently receives "vocals" gets a 200 and a +// wrong file, which is the failure mode nobody can see; the message therefore +// lists the stem names this model really has. +// +// The names are also validated, because they come from the MODEL (htdemucs +// reads them from the GGUF's config.sources) and each one becomes a component +// of a file path this backend writes. A name carrying a path separator would +// write outside the caller's output directory, and two stems sharing a name +// would silently overwrite each other. Both are refused before anything is +// written, which is also why selection has to happen before the first write +// rather than after the loop: a refused request must leave no files behind. +StemChoice select_named_output(const std::vector &names, + const std::string &requested); + +// "/generated/transform-1.wav" + "drums" -> "/generated/transform-1.drums.wav". +// A dst with no extension gets ".wav", since that is what write_audio_file +// produces whatever the caller called the file. +std::string sibling_stem_path(const std::string &dst, const std::string &name); + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/stem_selection_test.cpp b/backend/cpp/audio-cpp/stem_selection_test.cpp new file mode 100644 index 000000000..1316e2aee --- /dev/null +++ b/backend/cpp/audio-cpp/stem_selection_test.cpp @@ -0,0 +1,250 @@ +// Unit tests for stem_selection. Standard library only. The harness +// (backend/cpp/run-unit-tests.sh) compiles this as a single translation unit, +// so the implementation is included directly. +// +// What is actually at stake here: AudioTransformResult carries one dst, a +// separation model produces several stems, and the caller cannot see which one +// they got. Every check below is about a wrong file arriving with a 200. + +#include "stem_selection.cpp" + +#include +#include +#include + +static int failures = 0; + +static void check(bool ok, const std::string &name) { + if (!ok) { + failures++; + fprintf(stderr, "FAIL: %s\n", name.c_str()); + } else { + fprintf(stderr, "ok: %s\n", name.c_str()); + } +} + +static void check_equal(const std::string &got, const std::string &want, + const std::string &name) { + check(got == want, name + " (got \"" + got + "\", want \"" + want + "\")"); +} + +using namespace audiocpp_backend; + +// htdemucs's real source order, taken from its GGUF config.sources. vocals is +// LAST, which is why "first output" is not the default. +static const std::vector kDemucs = {"drums", "bass", "other", + "vocals"}; +// mel_band_roformer's, where vocals is first. +static const std::vector kRoformer = {"vocals", "instrumental"}; + +static void test_default_selection() { + const auto demucs = select_named_output(kDemucs, ""); + check(demucs.error.empty(), "an unrequested selection is not an error"); + check(demucs.index == 3, "no stem asked for picks vocals, not the first output"); + + const auto roformer = select_named_output(kRoformer, ""); + check(roformer.index == 0, "vocals is picked when it is already first"); + + // No vocals anywhere: the first output is the documented fallback. + const auto novocals = select_named_output({"accompaniment", "drums"}, ""); + check(novocals.index == 0 && novocals.error.empty(), + "a family with no vocals stem falls back to the first output"); + + // Substring matches must not count: "vocals_2" is a different stem. + const auto near = select_named_output({"vocals_2", "backing"}, ""); + check(near.index == 0 && near.error.empty(), + "'vocals_2' is not 'vocals', so the fallback and not the preference applies"); + const auto near_second = select_named_output({"backing", "vocals_2"}, ""); + check(near_second.index == 0, + "a near miss on the preferred name does not pull it to the front"); +} + +static void test_explicit_selection() { + for (int i = 0; i < 4; ++i) { + const auto choice = select_named_output(kDemucs, kDemucs[static_cast(i)]); + check(choice.index == i && choice.error.empty(), + "explicit '" + kDemucs[static_cast(i)] + "' selects its own index"); + } + // Including the one the default would have chosen anyway: asking for it + // must not be treated as "no request". + const auto vocals = select_named_output(kDemucs, "vocals"); + check(vocals.index == 3, "explicitly asking for vocals still selects vocals"); +} + +static void test_unknown_stem_is_refused() { + const auto choice = select_named_output(kDemucs, "kazoo"); + check(choice.index == -1, "an unknown stem selects nothing"); + check(!choice.error.empty(), "an unknown stem is refused rather than substituted"); + check(choice.error.find("kazoo") != std::string::npos, + "the refusal names the stem that was asked for"); + // The real names, so the caller can fix the request without guessing. + for (const auto &name : kDemucs) { + check(choice.error.find(name) != std::string::npos, + "the refusal lists the real stem '" + name + "'"); + } + check(choice.error.find("drums, bass, other, vocals") != std::string::npos, + "the refusal lists the stems in the model's own order"); + + // Case matters: the engine's ids are exact, so a wrong case is a wrong name + // rather than a near miss to be forgiven. + const auto wrong_case = select_named_output(kDemucs, "Vocals"); + check(wrong_case.index == -1 && !wrong_case.error.empty(), + "stem names are matched case sensitively"); +} + +static void test_no_named_outputs() { + const auto choice = select_named_output({}, ""); + check(choice.index == -1, "an empty output list selects nothing"); + check(choice.error.empty(), + "an empty output list is not an error here: the caller decides"); + const auto requested = select_named_output({}, "vocals"); + check(requested.index == -1 && requested.error.empty(), + "an empty output list stays the caller's decision even when a stem was asked for"); +} + +static void test_unwritable_names_are_refused() { + // Model-supplied names become file path components. A separator would write + // outside the caller's output directory. + const std::vector traversal = {"vocals", "../../etc/passwd"}; + const auto escaped = select_named_output(traversal, "vocals"); + check(escaped.index == -1 && !escaped.error.empty(), + "a stem name containing a path separator is refused"); + check(escaped.error.find("../../etc/passwd") != std::string::npos, + "the refusal names the offending stem"); + + check(!select_named_output({"vo\\cals", "drums"}, "").error.empty(), + "a backslash separator is refused too"); + check(!select_named_output({"drums", ""}, "").error.empty(), + "an empty stem name is refused"); + check(!select_named_output({"drums", "."}, "").error.empty(), + "a stem named '.' is refused"); + check(!select_named_output({"drums", ".."}, "").error.empty(), + "a stem named '..' is refused"); + + // The check covers EVERY name, not only the selected one: all of them are + // written, so a bad fourth name must not be found after three files exist. + const auto late = select_named_output({"vocals", "drums", "bass", "a/b"}, "vocals"); + check(late.index == -1 && !late.error.empty(), + "an unwritable name after the selected one still refuses the whole request"); + + // Control bytes, and the NUL case is why the whole range is refused. These + // two names are DIFFERENT std::strings, so the duplicate check does not + // fire, yet both truncate to "vocals" at path::c_str() and would open one + // file: the silent overwrite the duplicate check exists to prevent, with + // the ".wav" stripped off into the bargain. + const std::string nul_a("vocals\0drums", 12); + const std::string nul_b("vocals\0bass", 11); + check(nul_a != nul_b, "the two NUL names really are distinct std::strings"); + check(std::string(nul_a.c_str()) == "vocals" && + std::string(nul_b.c_str()) == "vocals", + "and both truncate to the same C string, which is the hazard"); + const auto nul_pair = select_named_output({nul_a, nul_b}, ""); + check(nul_pair.index == -1 && !nul_pair.error.empty(), + "two stem names differing only after an embedded NUL are refused"); + check(!select_named_output({"drums", std::string("vo\0cals", 7)}, "").error.empty(), + "a single embedded NUL is refused on its own"); + check(!select_named_output({"drums", "voc\nals"}, "").error.empty(), + "a newline in a stem name is refused"); + check(!select_named_output({"drums", "voc\tals"}, "").error.empty(), + "a tab in a stem name is refused"); + check(!select_named_output({"drums", "voc\033[31mals"}, "").error.empty(), + "an escape sequence in a stem name is refused"); + check(!select_named_output({"drums", "voc\177als"}, "").error.empty(), + "DEL in a stem name is refused"); + + // The boundary below the refused range is the space, which is an ordinary + // file name character and must stay usable, or this check would be + // refusing real stem names. + const auto spaced = select_named_output({"lead vocals", "drums"}, "lead vocals"); + check(spaced.index == 0 && spaced.error.empty(), + "a space is not a control character and stays usable"); + // And every byte above DEL: a UTF-8 stem name is ordinary, and signed char + // would make those bytes compare as negative. + const auto utf8 = select_named_output({"vocals", "b\xc3\xa4sse"}, "b\xc3\xa4sse"); + check(utf8.index == 1 && utf8.error.empty(), + "a UTF-8 stem name is not mistaken for a control character"); + + // A leading dot is not a traversal and must stay usable. + const auto dotted = select_named_output({".vocals", "drums"}, ".vocals"); + check(dotted.index == 0 && dotted.error.empty(), + "a leading dot in a stem name is allowed"); +} + +static void test_duplicate_names_are_refused() { + const auto choice = select_named_output({"vocals", "drums", "vocals"}, "vocals"); + check(choice.index == -1 && !choice.error.empty(), + "two stems sharing a name are refused: one file would overwrite the other"); + check(choice.error.find("vocals") != std::string::npos, + "the duplicate refusal names the repeated stem"); +} + +static void test_sibling_paths() { + check_equal(sibling_stem_path("/generated/transform-1.wav", "drums"), + "/generated/transform-1.drums.wav", "sibling beside an absolute dst"); + check_equal(sibling_stem_path("sep.wav", "vocals"), "sep.vocals.wav", + "sibling of a bare file name has no directory"); + check_equal(sibling_stem_path("/out/sep", "vocals"), "/out/sep.vocals.wav", + "an extensionless dst gets .wav"); + check_equal(sibling_stem_path("/out/take.2.wav", "bass"), "/out/take.2.bass.wav", + "only the final extension is treated as the extension"); + check_equal(sibling_stem_path("/out/sep.WAV", "bass"), "/out/sep.bass.WAV", + "the caller's extension spelling is preserved"); + check_equal(sibling_stem_path("/a b/c d.wav", "other"), "/a b/c d.other.wav", + "spaces in the destination survive"); + + // The property that matters: no stem can ever be written over dst itself, + // or the "dst holds the selected stem" contract would depend on write order. + const std::string dst = "/out/sep.wav"; + for (const auto &name : kDemucs) { + check(sibling_stem_path(dst, name) != dst, + "the sibling for '" + name + "' is not dst itself"); + } + // Distinct stems must land in distinct files. + check(sibling_stem_path(dst, "drums") != sibling_stem_path(dst, "bass"), + "two stems get two different sibling paths"); +} + +// The contract grpc-server.cpp indexes on: an accepted choice over a non-empty +// name list is always in range, so the handler needs no bounds guard of its own. +// A -1 reaching the subscript would become a colossal size_t. +static void test_accepted_index_is_always_in_range() { + const std::vector> lists = { + kDemucs, kRoformer, {"solo"}, {"accompaniment", "drums"}, {"a", "b", "c"}}; + const std::vector requests = {"", "vocals", "drums", "solo", "c", + "kazoo", "..", "a/b"}; + for (const auto &names : lists) { + for (const auto &requested : requests) { + const auto choice = select_named_output(names, requested); + if (!choice.error.empty()) { + check(choice.index == -1, + "a refusal never carries an index (request '" + requested + "')"); + continue; + } + check(choice.index >= 0 && + choice.index < static_cast(names.size()), + "an accepted choice is in range (request '" + requested + "')"); + // And the selected name is the one that was asked for, when one was. + if (!requested.empty()) { + check(names[static_cast(choice.index)] == requested, + "an accepted explicit request selects that exact name"); + } + } + } +} + +int main() { + test_accepted_index_is_always_in_range(); + test_default_selection(); + test_explicit_selection(); + test_unknown_stem_is_refused(); + test_no_named_outputs(); + test_unwritable_names_are_refused(); + test_duplicate_names_are_refused(); + test_sibling_paths(); + if (failures != 0) { + fprintf(stderr, "%d check(s) failed\n", failures); + return 1; + } + fprintf(stderr, "all stem_selection checks passed\n"); + return 0; +} diff --git a/backend/cpp/audio-cpp/stream_delta.cpp b/backend/cpp/audio-cpp/stream_delta.cpp new file mode 100644 index 000000000..97e1b3750 --- /dev/null +++ b/backend/cpp/audio-cpp/stream_delta.cpp @@ -0,0 +1,204 @@ +#include "stream_delta.h" + +namespace audiocpp_backend { +namespace { + +// True when `text` begins with `prefix`. An empty prefix matches everything, +// which is what makes the first fragment take the cumulative branch and the +// incremental branch alike: they agree there. +bool starts_with(const std::string &text, const std::string &prefix) { + return text.size() >= prefix.size() && + text.compare(0, prefix.size(), prefix) == 0; +} + +// Number of leading bytes of `text` that cannot BEGIN a character, i.e. orphan +// continuation bytes with no lead byte in front of them. +// +// They are unrecoverable rather than early: the byte that would have led them +// has already gone past, and nothing can be prepended to a fragment after the +// fact. Holding them would stall the stream for good, and emitting them puts +// invalid UTF-8 on the wire, so the caller DROPS them. Losing a byte keeps the +// stream alive; emitting one ends it, and takes the final_result still to come +// with it. +std::size_t utf8_orphan_prefix_length(const std::string &text) { + std::size_t index = 0; + while (index < text.size() && + (static_cast(text[index]) & 0xC0) == 0x80) { + ++index; + } + return index; +} + +// Length of the longest prefix of `text` that does NOT end inside a multi-byte +// UTF-8 sequence, i.e. the most that can go on the wire without splitting a +// character in half. +// +// Only the TRAILING sequence is examined; the leading side is +// utf8_orphan_prefix_length's job. Bytes in the MIDDLE are neither one's +// business: a family that emitted a malformed sequence inside its own text +// cannot be repaired here without deleting part of that transcript, and the +// stream is lost anyway, because final_result.text carries the same bytes +// through the same proto3 string field. +// +// Anything that can never complete is reported as complete, so it goes out +// rather than being held forever: a lead byte the encoding does not define, and +// a run of five or more continuation bytes, are both passed through. A tracker +// that stalled on undecodable input would turn one bad byte into a permanently +// silent stream, which is worse than the bad byte. +std::size_t utf8_complete_prefix_length(const std::string &text) { + std::size_t index = text.size(); + std::size_t continuations = 0; + while (index > 0 && continuations < 4) { + const auto byte = static_cast(text[index - 1]); + if ((byte & 0xC0) == 0x80) { + --index; + ++continuations; + continue; + } + std::size_t needed = 1; + if ((byte & 0x80) == 0x00) { + needed = 1; + } else if ((byte & 0xE0) == 0xC0) { + needed = 2; + } else if ((byte & 0xF0) == 0xE0) { + needed = 3; + } else if ((byte & 0xF8) == 0xF0) { + needed = 4; + } else { + // Not a lead byte this encoding defines, so nothing is waiting on + // it and it must not be held. + needed = 1; + } + if (continuations + 1 >= needed) { + return text.size(); + } + // The trailing sequence is short by at least one byte: cut before its + // lead byte and keep the rest for the next fragment. + return index - 1; + } + return text.size(); +} + +} // namespace + +std::string TranscriptDeltaTracker::release(const std::string &fragment) { + // Held-back bytes go in FRONT of whatever arrived next, or the character + // they begin is reassembled in the wrong order. + std::string candidate = pending_ + fragment; + // A fragment must not BEGIN mid-character either. pending_ always starts on + // a lead byte, so this only bites when nothing was held and the caller's + // rules dropped the lead byte somewhere upstream; it is the backstop that + // makes "no delta this class returns is ever invalid UTF-8" true of the + // FRONT as well as the back, independently of those rules being right. + candidate.erase(0, utf8_orphan_prefix_length(candidate)); + const std::size_t cut = utf8_complete_prefix_length(candidate); + pending_ = candidate.substr(cut); + std::string emitted = candidate.substr(0, cut); + assembled_ += emitted; + return emitted; +} + +std::string TranscriptDeltaTracker::observe(const std::string &partial_text) { + if (partial_text.empty()) { + return {}; + } + // The comparisons run against everything KNOWN, delivered plus held back, + // rather than against the delivered text alone. Comparing against the + // delivered text would treat the held-back byte as new on the very next + // report and emit it twice. + const std::string known = assembled_ + pending_; + // An EXACT repeat, and nothing looser. This absorbs the duplicate delivery + // voxtral_realtime produces, which is the ONLY thing rule 2 was ever needed + // for: process_available_stream_chunks hands each event it produces to the + // sink from inside its loop and RETURNS the last of the batch + // (session.cpp:385-386), so that last event arrives twice with byte-equal + // text both times. A duplicate IS an exact repeat, so equality covers it. + // + // It used to discard any report the known text merely STARTED WITH, and that + // cost far more than it bought. Two separate defects came out of it, and + // both were found by randomized traces rather than by reading: + // + // 1. A short INCREMENTAL fragment that happens to be a byte prefix of the + // transcript so far was read as a repeat and dropped, losing text with + // a 200 and no diagnostic. Both incremental families emit fragments + // that small routinely: nemotron_asr cuts at a byte offset + // (decoder.cpp:550) and vibevoice_asr at a common prefix + // (session.cpp:89-105). 9.50% of randomized pure-ASCII traces and + // 29.12% of French ones ended with a corrupted transcript. + // 2. When such a fragment was the LEAD BYTE of a multi-byte character, its + // continuation bytes then arrived alone and began the next delta, which + // is invalid UTF-8, which the Go runtime refuses to unmarshal, which + // ends the stream and the final_result with it. + // + // What the narrowing gives up is the shrinking-hypothesis case: a cumulative + // report SHORTER than what is known is now read as an incremental fragment + // and duplicates those bytes at the client. No pinned family produces one. + // voxtral is the only cumulative reporter, and its hypothesis is + // tokenizer_.decode(streaming_token_ids_) over a vector that is only ever + // push_back'ed (session.cpp:436) and cleared by reset() (session.cpp:257), + // so within a stream it can only grow. Measured: narrowing this changed not + // one byte of 30,000 randomized cumulative traces. + // + // KEPT DELIBERATELY THOUGH NO TEST CAN SEE IT. Once narrowed to equality + // this rule became redundant with rule 3 below: an equal partial has an + // empty suffix, so rule 3 would call release("") and emit nothing either + // way. Deleting it is therefore an equivalent mutation, and the mutation + // harness reports it as a survivor, which is the honest result and not a + // gap in the tests. It stays for two reasons: it states the duplicate + // absorption where the citation for it lives, and it is independent of rule + // 3's condition. A future tightening of rule 3 to, say, require a STRICTLY + // longer partial would otherwise send every duplicate down the incremental + // branch and put the whole transcript on the wire a second time. + if (partial_text == known) { + return {}; + } + if (starts_with(partial_text, known)) { + // Cumulative: the report is the whole transcript so far. + return release(partial_text.substr(known.size())); + } + // Incremental: the fragment is new text to append. + // + // A family that REWRITES its hypothesis lands here too, and the client's + // view is then wrong in a way nothing downstream can fix. nemotron's + // decoder has such a branch (decoder.cpp:552-554): when the new text is not + // an extension of what it already emitted, it emits the whole new text. So + // "the cat sat" followed by "the cat sap" leaves the client holding + // "the cat satthe cat sap", and reconcile then correctly refuses to append + // to a contradicted assembly, which leaves concat(deltas) != final with no + // signal on the wire. This is NOT repaired here, and the reason is that a + // delta stream has no retraction: emitting only the differing suffix would + // read as "sap" appended to "the cat sat", which is a different wrong + // answer, and emitting a correction would need a wire field that does not + // exist. final_result carries the authoritative text either way. It did not + // fire in a 331 delta run, because an RNN-T decode is monotonic in practice. + return release(partial_text); +} + +std::string TranscriptDeltaTracker::reconcile(const std::string &final_text) { + if (final_text.empty() || final_text == assembled_) { + // Nothing further is owed. Held-back bytes are dropped rather than + // flushed: they are not in the authoritative text, so sending them + // would contradict it. + pending_.clear(); + return {}; + } + if (!starts_with(final_text, assembled_)) { + // Contradicted. Nothing sent can be taken back, so nothing more is + // sent; final_result carries the authoritative text. + pending_.clear(); + return {}; + } + // Compared against the DELIVERED text, so the fragment below already + // contains whatever was held back. pending_ is therefore cleared rather + // than prepended, or those bytes would go out twice. + // + // The fragment ends on a character boundary whenever final_text is + // well-formed, which is the normal case and the reason a held-back sequence + // is always flushed here. It is still cut, so a family handing back a final + // text that is itself truncated mid-character cannot put a partial sequence + // on the wire through this path either. + pending_.clear(); + return release(final_text.substr(assembled_.size())); +} + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/stream_delta.h b/backend/cpp/audio-cpp/stream_delta.h new file mode 100644 index 000000000..9becf49fb --- /dev/null +++ b/backend/cpp/audio-cpp/stream_delta.h @@ -0,0 +1,133 @@ +#pragma once + +// Turns whatever a streaming session calls a "partial transcript" into the +// incremental deltas AudioTranscriptionStream is contracted to send. Standard +// library only, so it is tested without an audio.cpp checkout. +// +// THIS UNIT EXISTS BECAUSE THE FAMILIES DISAGREE, and the disagreement is +// invisible at the interface: StreamEvent::partial_text is a Transcript either +// way. Read out of the pinned upstream, one family at a time: +// +// nemotron_asr INCREMENTAL. decoder.cpp emits +// current_text.substr(emitted_text.size()) per non-blank +// token, and only through the stream event SINK, during +// finalize(). process_audio_chunk returns empty events. +// vibevoice_asr INCREMENTAL. process_audio_chunk returns +// text.substr(common_prefix_size(...)); the sink is +// deliberately swapped out around its internal run_single, +// so the fragment arrives once, on the return value. +// higgs_audio_stt INCREMENTAL. Same shape as vibevoice_asr. +// voxtral_realtime CUMULATIVE. partial_text is +// tokenizer_.decode(streaming_token_ids_), the whole +// hypothesis so far. process_available_stream_chunks hands +// every event it produces to the sink from INSIDE its loop +// (session.cpp:385-386) and RETURNS only the last of the +// batch, so the last event of each batch arrives twice and +// the others arrive once. +// +// Applying either convention to the other family corrupts the transcript: read +// a cumulative report as a delta and the client sees the transcript repeated on +// every event; read an incremental fragment as cumulative and the suffix +// arithmetic eats the front of it. So the tracker decides per fragment, from +// what it has already delivered, and the one rule it enforces is that TEXT THE +// CLIENT HAS ALREADY BEEN SENT IS NEVER SENT AGAIN. +// +// The cumulative reading is provably safe for voxtral, which is the family it +// matters for: its decode is a pure concatenation of per-token byte strings +// (tokenizer_text.cpp:171-183), so decode(ids[0..n]) is an unconditional BYTE +// PREFIX of decode(ids[0..n+1]) and one of its reports can never be mistaken +// for an incremental fragment. +// +// UTF-8 IS THE OTHER HALF OF THAT SAME FACT. Because that decode concatenates +// raw token BYTES, a multi-byte character is split across token boundaries, and +// the difference between two consecutive cumulative reports is then a lone +// continuation byte. TranscriptStreamResponse.delta is a proto3 `string`, whose +// wire format REQUIRES valid UTF-8: the C++ runtime serializes an invalid one +// with at most a warning, but the Go runtime refuses to unmarshal it, and the +// client loses every remaining delta AND the final_result. So no fragment this +// class returns ever BEGINS OR ENDS inside a character: an incomplete trailing +// sequence is held back and merged into the next fragment, and a leading orphan +// continuation byte, which nothing can ever complete, is dropped. +// +// It is NOT a voxtral-only concern, which is what the first attempt at this +// assumed. The incremental families split characters by the same arithmetic: +// nemotron_asr's decoder cuts at a BYTE offset (decoder.cpp:550) and +// vibevoice_asr's common_prefix_size compares BYTES (session.cpp:80-86). Nor is +// European text the worst case: a Japanese transcript, whose every character is +// three bytes, carried at least one invalid delta in 33.74% of traces until +// rule 2 below learned to leave an incomplete fragment alone. + +#include + +namespace audiocpp_backend { + +class TranscriptDeltaTracker { +public: + // Takes one StreamEvent::partial_text and returns the fragment to put on + // the wire, empty when there is nothing new. + // + // The rules, in order, all of them against everything KNOWN (delivered + // plus held back), never against the delivered text alone: + // 1. An empty partial says nothing. + // 2. A partial IDENTICAL to the known text is a repeat: nothing is + // emitted. Identical, not merely a prefix of it. That absorbs voxtral's + // repeat of the last event in each batch, which is the only duplicate + // any pinned family produces and which carries byte-equal text both + // times. Discarding a mere PREFIX used to swallow an incremental + // fragment that coincided with the start of the transcript, corrupting + // the text silently and, when that fragment was a character's lead + // byte, killing the stream outright; see the note at the rule in the + // implementation. + // 3. A partial that EXTENDS the known text is a cumulative report: only + // its new suffix is emitted. + // 4. Anything else is an incremental fragment: it is emitted whole and + // appended. + // + // Rule 3 is the one judgement call, since a fragment that happens to begin + // with the entire transcript so far is indistinguishable from a cumulative + // report. It is read as cumulative because every cumulative family produces + // that shape on EVERY event, while an incremental family produces it only + // when one fragment repeats everything before it, which no tokenizer output + // does in practice. + // + // What comes back is the fragment MINUS any incomplete trailing UTF-8 + // sequence, which is carried into the next call, and minus any leading + // orphan continuation byte, which is dropped. So an empty return can also + // mean "the only new bytes were half a character", and the caller needs no + // knowledge of that: writing nothing is exactly right. + std::string observe(const std::string &partial_text); + + // Reconciles against TaskResult::text_output, which is authoritative, and + // returns the fragment that makes appending every delta equal it. + // + // This is what makes the OFFLINE FALLBACK a single line rather than its own + // branch: with no partials observed, the assembly is empty and the whole + // final text comes back as one delta. + // + // It is also what FLUSHES a held-back UTF-8 sequence, and it can always do + // so: the final text is complete, so the fragment from the last delivered + // byte to its end ends on a character boundary. + // + // A final text that CONTRADICTS what was already sent returns empty. A + // fragment on the wire cannot be retracted, so the alternative would be to + // send the transcript a second time and let the client hold it twice. + // final_result carries the authoritative text either way. + std::string reconcile(const std::string &final_text); + + // Everything the client has been sent, concatenated. Held-back bytes are + // deliberately NOT included: this is what the client holds, not what the + // tracker knows. + const std::string &assembled() const noexcept { return assembled_; } + +private: + // Appends the emittable prefix of `fragment` to assembled_ and returns it, + // keeping any incomplete trailing UTF-8 sequence in pending_. + std::string release(const std::string &fragment); + + std::string assembled_; + // An incomplete trailing UTF-8 sequence, computed but not sent. Always a + // proper prefix of one character, so at most three bytes. + std::string pending_; +}; + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/stream_delta_test.cpp b/backend/cpp/audio-cpp/stream_delta_test.cpp new file mode 100644 index 000000000..1ceca4e2d --- /dev/null +++ b/backend/cpp/audio-cpp/stream_delta_test.cpp @@ -0,0 +1,486 @@ +// Unit tests for stream_delta. Standard library only. The harness compiles this +// as a single translation unit, so the implementation is included directly. +// +// The traces below are transcribed from the pinned upstream sessions rather +// than invented, because the whole reason this unit exists is that the four +// streaming ASR families do NOT agree on what partial_text means. + +#include "stream_delta.cpp" + +#include +#include +#include + +static int failures = 0; + +static void check(bool ok, const std::string &name) { + if (!ok) { + failures++; + fprintf(stderr, "FAIL: %s\n", name.c_str()); + } else { + fprintf(stderr, "ok: %s\n", name.c_str()); + } +} + +static void check_eq(const std::string &got, const std::string &want, + const std::string &name) { + check(got == want, name + " (got \"" + got + "\" want \"" + want + "\")"); +} + +using audiocpp_backend::TranscriptDeltaTracker; + +// Feeds a whole trace and returns what a client appending every emitted +// fragment would end up holding, which is the only property that matters. +static std::string client_view(TranscriptDeltaTracker &tracker, + const std::vector &partials, + std::vector *emitted = nullptr) { + std::string view; + for (const auto &partial : partials) { + const std::string fragment = tracker.observe(partial); + if (emitted != nullptr && !fragment.empty()) { + emitted->push_back(fragment); + } + view += fragment; + } + return view; +} + +// nemotron_asr: decoder.cpp emits current_text.substr(emitted_text.size()) on +// every non-blank token, i.e. INCREMENTAL fragments, through the stream event +// sink during finalize(). +static void test_incremental_family() { + TranscriptDeltaTracker tracker; + std::vector emitted; + const std::string view = + client_view(tracker, {"Local", " AI", " now", " speaks."}, &emitted); + check_eq(view, "Local AI now speaks.", "incremental deltas concatenate"); + check(emitted.size() == 4, "incremental family emits one fragment per partial"); + check_eq(tracker.assembled(), "Local AI now speaks.", + "incremental family assembles the whole transcript"); + check_eq(tracker.reconcile("Local AI now speaks."), "", + "a final result the deltas already cover adds nothing"); +} + +// voxtral_realtime: process_one_stream_chunk sets partial_text to +// tokenizer_.decode(streaming_token_ids_), the WHOLE accumulated hypothesis, +// and process_available_stream_chunks then hands the SAME event to both the +// sink and the caller, so every partial arrives twice. +static void test_cumulative_family_with_duplicate_delivery() { + TranscriptDeltaTracker tracker; + std::vector emitted; + const std::string view = client_view( + tracker, {"Local", "Local", "Local AI", "Local AI", "Local AI now", + "Local AI now"}, + &emitted); + check_eq(view, "Local AI now", "cumulative partials are not repeated to the client"); + check(emitted.size() == 3, + "the duplicate delivery of each cumulative event emits nothing twice"); + check_eq(emitted.empty() ? "" : emitted[0], "Local", "first cumulative fragment"); + check_eq(emitted.size() < 2 ? "" : emitted[1], " AI", "second cumulative fragment"); + check_eq(emitted.size() < 3 ? "" : emitted[2], " now", "third cumulative fragment"); +} + +// The rule that separates the two: a report IDENTICAL to everything known is a +// repeat and is never sent again. +// +// Only an identical one. Rule 2 used to discard any report the known text merely +// STARTED WITH, and that cost more than it bought: see +// test_a_short_fragment_is_not_mistaken_for_a_repeat. +static void test_an_exact_repeat_is_never_resent() { + TranscriptDeltaTracker tracker; + check_eq(tracker.observe("hello world"), "hello world", "first fragment"); + check_eq(tracker.observe("hello world"), "", "an identical repeat emits nothing"); + check_eq(tracker.observe("hello world"), "", "and a third delivery emits nothing"); + check_eq(tracker.assembled(), "hello world", "the assembly is unchanged by repeats"); +} + +// THE TRADE, pinned so it is a decision rather than a surprise. A CUMULATIVE +// report that SHRINKS is no longer absorbed: it is read as an incremental +// fragment and duplicates a few bytes at the client. +// +// No pinned family produces one. voxtral_realtime is the only cumulative +// reporter, and its hypothesis is tokenizer_.decode(streaming_token_ids_) over a +// vector that is only ever push_back'ed (session.cpp:436) and cleared by reset() +// (session.cpp:257), so within a stream it grows and never shrinks. The +// duplicate delivery rule 2 really exists for is an EXACT repeat, which the test +// above still covers. +static void test_a_shrinking_hypothesis_is_read_as_incremental() { + TranscriptDeltaTracker tracker; + check_eq(tracker.observe("hello world"), "hello world", "first fragment"); + check_eq(tracker.observe("hello"), "hello", + "a shortened hypothesis is now read as an incremental fragment"); + check_eq(tracker.assembled(), "hello worldhello", + "which duplicates those bytes at the client: the accepted cost"); +} + +static void test_empty_partials_are_ignored() { + TranscriptDeltaTracker tracker; + check_eq(tracker.observe(""), "", "an empty partial emits nothing"); + check_eq(tracker.observe("a"), "a", "a real partial after an empty one still emits"); + check_eq(tracker.observe(""), "", "a later empty partial emits nothing"); + check_eq(tracker.assembled(), "a", "empty partials do not disturb the assembly"); +} + +// The offline fallback: a family with no streaming ASR runs once, so nothing is +// ever observed and the reconciliation IS the single delta the RPC promises. +static void test_offline_fallback_is_one_delta() { + TranscriptDeltaTracker tracker; + check_eq(tracker.reconcile("the whole transcript"), "the whole transcript", + "with no partials the final text is emitted whole"); + check_eq(tracker.assembled(), "the whole transcript", + "the reconciliation is recorded as delivered"); + check_eq(tracker.reconcile("the whole transcript"), "", + "reconciling twice does not duplicate"); +} + +// A streaming family whose partials stopped short of the final text: the tail +// is emitted so that appending every delta still equals final_result.text. +static void test_reconcile_emits_the_tail() { + TranscriptDeltaTracker tracker; + check_eq(tracker.observe("Local AI"), "Local AI", "partial arrives"); + check_eq(tracker.reconcile("Local AI now speaks."), " now speaks.", + "the untold tail of the final text is emitted"); + check_eq(tracker.assembled(), "Local AI now speaks.", "tail is recorded"); +} + +// Divergence. nemotron's decoder has a rewrite branch: when the new hypothesis +// is NOT an extension of what it already emitted, it emits the whole new text. +// Nothing can retract a fragment already written to the wire, so the tracker +// must not try: it emits nothing further and leaves final_result authoritative. +static void test_divergent_final_text_is_not_appended() { + TranscriptDeltaTracker tracker; + check_eq(tracker.observe("the cat"), "the cat", "first hypothesis"); + check_eq(tracker.reconcile("the dog"), "", + "a final text that contradicts the deltas is not appended to them"); + check_eq(tracker.assembled(), "the cat", + "a contradicted assembly is left as it was actually sent"); +} + +static void test_empty_final_text() { + TranscriptDeltaTracker tracker; + check_eq(tracker.observe("something"), "something", "partial arrives"); + check_eq(tracker.reconcile(""), "", "an empty final text emits nothing"); + check_eq(tracker.assembled(), "something", "an empty final text changes nothing"); +} + +// A whitespace-only fragment is real text: the space between two words is +// exactly what an incremental family delivers on its own. +static void test_whitespace_fragments_survive() { + TranscriptDeltaTracker tracker; + check_eq(tracker.observe("one"), "one", "word"); + check_eq(tracker.observe(" "), " ", "a bare separator is emitted"); + check_eq(tracker.observe("two"), "two", "next word"); + check_eq(tracker.assembled(), "one two", "separator is kept in the assembly"); +} + +// The ambiguity this unit cannot resolve, pinned so that a future reader sees +// the choice rather than rediscovering it: a fragment that EXTENDS everything +// delivered so far is read as a cumulative report, because that is what every +// cumulative family produces on every event, while an incremental family +// producing one is the rare coincidence of a fragment repeating the whole +// transcript so far. +static void test_prefix_extension_is_read_as_cumulative() { + TranscriptDeltaTracker tracker; + check_eq(tracker.observe("I"), "I", "first fragment"); + check_eq(tracker.observe("I'm"), "'m", + "a fragment extending the assembly is treated as a cumulative report"); + check_eq(tracker.assembled(), "I'm", "cumulative reading assembles once"); +} + +// -------------------------------------------------------------------------- +// UTF-8 boundaries +// +// TranscriptStreamResponse.delta is a proto3 `string`, and the wire format +// REQUIRES a string field to be valid UTF-8. The C++ runtime serializes an +// invalid one with at most a warning; the Go runtime refuses to unmarshal it, +// so the client loses every delta AND the final_result still to come. +// +// Not hypothetical. voxtral_realtime reports the whole hypothesis as +// tokenizer_.decode(streaming_token_ids_), and that decode is a pure +// concatenation of raw token BYTES (tokenizer_text.cpp:171-183), so a +// multi-byte character is split across token boundaries and the cumulative +// difference between two consecutive reports is a lone continuation byte. +// -------------------------------------------------------------------------- + +// True when `text` is well-formed UTF-8. Written out here rather than reused +// from the implementation on purpose: a test that shares the implementation's +// idea of a boundary cannot catch the implementation's idea being wrong. +static bool is_valid_utf8(const std::string &text) { + size_t i = 0; + while (i < text.size()) { + const auto lead = static_cast(text[i]); + size_t length = 0; + if ((lead & 0x80) == 0x00) { + length = 1; + } else if ((lead & 0xE0) == 0xC0) { + length = 2; + } else if ((lead & 0xF0) == 0xE0) { + length = 3; + } else if ((lead & 0xF8) == 0xF0) { + length = 4; + } else { + return false; + } + if (i + length > text.size()) { + return false; + } + for (size_t k = 1; k < length; ++k) { + if ((static_cast(text[i + k]) & 0xC0) != 0x80) { + return false; + } + } + i += length; + } + return true; +} + +// Written as byte escapes so the test does not depend on the encoding of this +// source file. +static const std::string kEAcute = "\xC3\xA9"; // 2 bytes +static const std::string kEuro = "\xE2\x82\xAC"; // 3 bytes +static const std::string kEmoji = "\xF0\x9F\x8E\xA7"; // 4 bytes + +// A cumulative family advancing its hypothesis one BYTE at a time, which is +// what voxtral_realtime does across a multi-byte character. +static void test_cumulative_split_multibyte_character() { + TranscriptDeltaTracker tracker; + std::vector emitted; + const std::string full = "5" + kEuro; + std::vector partials; + for (size_t n = 1; n <= full.size(); ++n) { + partials.push_back(full.substr(0, n)); + } + const std::string view = client_view(tracker, partials, &emitted); + + check_eq(view, full, "a byte-at-a-time cumulative report still assembles"); + for (size_t i = 0; i < emitted.size(); ++i) { + check(is_valid_utf8(emitted[i]), + "cumulative fragment " + std::to_string(i) + " is valid UTF-8"); + } + check_eq(tracker.reconcile(full), "", "the final text adds nothing"); +} + +// An incremental family splitting a character across two fragments. +static void test_incremental_split_multibyte_character() { + TranscriptDeltaTracker tracker; + std::vector emitted; + const std::string view = client_view( + tracker, {"caf" + kEAcute.substr(0, 1), kEAcute.substr(1), " au lait"}, + &emitted); + + check_eq(view, "caf" + kEAcute + " au lait", + "an incremental split character still assembles"); + for (size_t i = 0; i < emitted.size(); ++i) { + check(is_valid_utf8(emitted[i]), + "incremental fragment " + std::to_string(i) + " is valid UTF-8"); + } + check(emitted.size() == 3, "one fragment out per partial, none swallowed"); + if (emitted.size() == 3) { + check_eq(emitted[0], "caf", "the lead byte of the character is held back"); + check_eq(emitted[1], kEAcute, + "the held byte is merged into the next fragment, not sent alone"); + check_eq(emitted[2], " au lait", "the rest follows unchanged"); + } +} + +// A 4 byte character split three ways, so the held-back buffer has to survive +// more than one round. +static void test_four_byte_character_split_three_ways() { + TranscriptDeltaTracker tracker; + std::vector emitted; + const std::string view = + client_view(tracker, + {"listen " + kEmoji.substr(0, 1), kEmoji.substr(1, 2), + kEmoji.substr(3), " now"}, + &emitted); + check_eq(view, "listen " + kEmoji + " now", + "a 4 byte character survives three splits"); + for (size_t i = 0; i < emitted.size(); ++i) { + check(is_valid_utf8(emitted[i]), + "4 byte fragment " + std::to_string(i) + " is valid UTF-8"); + } +} + +// The held-back bytes must reach the client. reconcile can always flush them, +// because the final text is complete by construction. +static void test_reconcile_flushes_a_held_back_sequence() { + TranscriptDeltaTracker tracker; + const std::string first = tracker.observe("done" + kEuro.substr(0, 2)); + check_eq(first, "done", "the incomplete trailing sequence is held back"); + check(is_valid_utf8(first), "what was emitted is valid UTF-8"); + const std::string tail = tracker.reconcile("done" + kEuro); + check_eq(tail, kEuro, "reconcile flushes the completed character"); + check(is_valid_utf8(tail), "the flushed tail is valid UTF-8"); + check_eq(tracker.assembled(), "done" + kEuro, "the client holds the whole text"); +} + +// Held-back bytes are not lost track of: the next cumulative report emits the +// whole character rather than only the bytes that just arrived. +static void test_held_bytes_join_the_next_fragment() { + TranscriptDeltaTracker tracker; + check_eq(tracker.observe("a" + kEuro.substr(0, 1)), "a", + "only the complete prefix goes out"); + check_eq(tracker.observe("a" + kEuro), kEuro, + "the next cumulative report emits the whole character at once"); + check_eq(tracker.assembled(), "a" + kEuro, "assembly is correct"); +} + +// A whole multi-byte character arriving at once must NOT be held back: holding +// a complete sequence would stall every stream by one character. +static void test_a_complete_character_is_not_held() { + TranscriptDeltaTracker tracker; + check_eq(tracker.observe("x" + kEuro), "x" + kEuro, + "a fragment ending on a boundary is emitted immediately"); + check_eq(tracker.observe("x" + kEuro + kEmoji), kEmoji, + "and so is the next one"); +} + +// Bytes that can never complete must not be held forever: a stray continuation +// byte or an invalid lead is passed through rather than stalling the stream. +// Repairing a family's malformed output is not something a delta tracker can do +// without altering the transcript. +static void test_undecodable_bytes_are_not_held_forever() { + TranscriptDeltaTracker tracker; + check_eq(tracker.observe(std::string("ok\x80")), std::string("ok\x80"), + "a stray continuation byte is passed through, not held"); + check_eq(tracker.observe(std::string("ok\x80") + "next"), "next", + "the stream continues"); + + // A lead byte the encoding does not define (0xF8 and above). Nothing can + // ever complete it, so holding it would stall the stream for good. + TranscriptDeltaTracker invalid_lead; + check_eq(invalid_lead.observe(std::string("ok\xFE")), std::string("ok\xFE"), + "an undefined lead byte is passed through, not held"); + check_eq(invalid_lead.observe(std::string("ok\xFE") + "more"), "more", + "the stream continues past an undefined lead byte"); + + // The same byte followed by continuation bytes, which is the shape that + // looks most like a real sequence waiting to be completed. + TranscriptDeltaTracker invalid_run; + check_eq(invalid_run.observe(std::string("\xFE\x80\x80")), std::string("\xFE\x80\x80"), + "an undefined lead with continuations is passed through"); + + // Five continuation bytes with no lead in sight. They are DROPPED, not + // held: nothing can ever precede them, so holding would stall the stream + // for good, and emitting them would put invalid UTF-8 on the wire. See + // test_a_fragment_never_begins_mid_character. + TranscriptDeltaTracker orphans; + check_eq(orphans.observe(std::string("\x80\x80\x80\x80\x80")), "", + "a run of orphan continuation bytes is dropped, not emitted"); + check_eq(orphans.observe("after"), "after", + "and the stream continues past them"); +} + +// THE SECOND HALF OF THE SAME BUG, and the one that survived fix round 1. +// +// Rule 2 discards a fragment the known text already starts with. When that +// fragment is the LEAD BYTE of a NEW character it looks exactly like a repeat of +// an earlier character beginning with the same byte, so it was discarded and +// never held. Its continuation bytes then arrived on their own and began the +// next delta, which is invalid UTF-8 at the FRONT, and utf8_complete_prefix_length +// only ever inspected the TRAILING sequence. +// +// Reachable from shipping families, not synthetic: nemotron_asr's +// decoder.cpp:550 cuts at a BYTE offset (current_text.substr(emitted_text.size())) +// and vibevoice_asr's common_prefix_size (session.cpp:80-86) compares BYTES, so +// both split characters mid-sequence. The trace below is exactly how they split +// "ssee" spelled with the German sharp s, an e-acute, a euro sign and an o-grave, +// three of which begin with the same 0xC3 lead byte. +static void test_a_repeated_lead_byte_is_not_swallowed() { + TranscriptDeltaTracker tracker; + std::vector emitted; + const std::string sharp_s = "\xC3\x9F"; // U+00DF + const std::string e_acute = "\xC3\xA9"; // U+00E9 + const std::string euro = "\xE2\x82\xAC"; // U+20AC + const std::string o_grave = "\xC3\xB2"; // U+00F2 + const std::string full = sharp_s + e_acute + euro + o_grave; + + const std::string view = client_view(tracker, + {sharp_s.substr(0, 1), sharp_s.substr(1), + e_acute.substr(0, 1), + e_acute.substr(1) + euro, + o_grave.substr(0, 1), o_grave.substr(1)}, + &emitted); + + for (size_t i = 0; i < emitted.size(); ++i) { + check(is_valid_utf8(emitted[i]), + "repeated-lead fragment " + std::to_string(i) + " is valid UTF-8"); + } + check_eq(view, full, "no character is lost to a repeated lead byte"); + check_eq(tracker.reconcile(full), "", "the final text adds nothing"); +} + +// The same shape one layer down, as a backstop: a fragment that BEGINS with +// orphan continuation bytes must never go on the wire, whatever produced it. +// Dropping bytes keeps the stream alive; emitting them ends it, and takes the +// final_result that was still to come with it. +static void test_a_fragment_never_begins_mid_character() { + TranscriptDeltaTracker tracker; + const std::string first = tracker.observe(std::string("\xA9") + "rest"); + check(is_valid_utf8(first), "a leading orphan continuation byte is not emitted"); + check_eq(first, "rest", "the rest of the fragment still goes out"); + + TranscriptDeltaTracker all_orphans; + check_eq(all_orphans.observe(std::string("\x82\xAC")), "", + "a fragment that is nothing but orphans emits nothing"); + check_eq(all_orphans.assembled(), "", + "and nothing is recorded as delivered"); + check_eq(all_orphans.observe("after"), "after", "the stream continues"); +} + +// PURE ASCII, no multi-byte character anywhere, and the transcript still comes +// out wrong: a short incremental fragment that happens to be a byte prefix of +// everything known was read as an already-delivered repeat and discarded. +// +// This is the shape both incremental families produce. nemotron_asr emits +// current_text.substr(emitted_text.size()) per non-blank token (decoder.cpp:550) +// and vibevoice_asr emits text.substr(common_prefix_size(...)) (session.cpp:89-105), +// so a one-character fragment is ordinary output, and any of the transcript's +// own leading characters will eventually arrive as one. +// +// Measured over 5,000 randomized traces per transcript before this was fixed: +// 9.50% of pure-ASCII traces and 29.12% of French ones ended with the client +// holding something other than final_result.text, with a 200 and no diagnostic. +static void test_a_short_fragment_is_not_mistaken_for_a_repeat() { + TranscriptDeltaTracker tracker; + std::vector emitted; + const std::string full = "pure ascii transcript"; + const std::string view = client_view( + tracker, {"pure ", "ascii ", "trans", "c", "ri", "p", "t"}, &emitted); + + check_eq(view, full, "the whole transcript reaches the client"); + check_eq(tracker.assembled(), full, "and the assembly agrees with it"); + check_eq(tracker.reconcile(full), "", + "the final text adds nothing, because nothing was lost"); + check(emitted.size() == 7, "every fragment produced exactly one delta"); +} + +int main() { + test_incremental_family(); + test_cumulative_family_with_duplicate_delivery(); + test_an_exact_repeat_is_never_resent(); + test_a_shrinking_hypothesis_is_read_as_incremental(); + test_a_short_fragment_is_not_mistaken_for_a_repeat(); + test_empty_partials_are_ignored(); + test_offline_fallback_is_one_delta(); + test_reconcile_emits_the_tail(); + test_divergent_final_text_is_not_appended(); + test_empty_final_text(); + test_whitespace_fragments_survive(); + test_prefix_extension_is_read_as_cumulative(); + test_cumulative_split_multibyte_character(); + test_incremental_split_multibyte_character(); + test_four_byte_character_split_three_ways(); + test_reconcile_flushes_a_held_back_sequence(); + test_held_bytes_join_the_next_fragment(); + test_a_complete_character_is_not_held(); + test_undecodable_bytes_are_not_held_forever(); + test_a_repeated_lead_byte_is_not_swallowed(); + test_a_fragment_never_begins_mid_character(); + if (failures) { + fprintf(stderr, "%d check(s) failed\n", failures); + return 1; + } + fprintf(stderr, "all stream_delta checks passed\n"); + return 0; +} diff --git a/backend/cpp/audio-cpp/streaming_driver_ctest.cpp b/backend/cpp/audio-cpp/streaming_driver_ctest.cpp new file mode 100644 index 000000000..e25616419 --- /dev/null +++ b/backend/cpp/audio-cpp/streaming_driver_ctest.cpp @@ -0,0 +1,976 @@ +// Tests for the streaming drivers in loaded_model: begin_stream, +// run_streaming_pull, run_streaming_audio and run_streaming_live, plus +// resolve_model_path, which lives in the same engine-linked unit. +// +// Engine-linked, so this runs through ctest rather than +// backend/cpp/run-unit-tests.sh. It builds no model and loads no file: a +// LoadedModel::Session is a plain struct holding a pointer to an engine +// interface, so a fake session exercises the drivers directly, which is the +// only way to assert the STATE OBLIGATION (prepare, then start_stream, on every +// stream) without a GPU and a gigabyte of weights. + +#include "inference_lane.h" +#include "loaded_model.h" + +#include "engine/framework/runtime/session.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace rt = engine::runtime; + +static int failures = 0; + +static void check(bool ok, const std::string &name) { + if (!ok) { + failures++; + fprintf(stderr, "FAIL: %s\n", name.c_str()); + } else { + fprintf(stderr, "ok: %s\n", name.c_str()); + } +} + +static void check_eq(const std::string &got, const std::string &want, + const std::string &name) { + check(got == want, name + " (got \"" + got + "\" want \"" + want + "\")"); +} + +static std::string join(const std::vector &parts) { + std::string out; + for (const auto &part : parts) { + if (!out.empty()) { + out += "|"; + } + out += part; + } + return out; +} + +// -------------------------------------------------------------------------- +// Fakes +// -------------------------------------------------------------------------- + +// Consumes audio chunks, like every streaming ASR family. +// +// It deliberately does NOT override start_stream, so every test that drives it +// also pins the claim the drivers rely on: IStreamingVoiceTaskSession's BASE +// start_stream is a call to reset(). If upstream ever changes that base, the +// replay test below fails rather than the backend silently continuing the +// previous stream. +class FakeAudioSession : public rt::IStreamingVoiceTaskSession { +public: + std::vector calls; + std::vector chunks; + rt::StreamingPolicy policy; + // Set to make process_audio_chunk throw on the nth call (1-based). + int throw_on_chunk = 0; + // Cumulative partial text, which is voxtral_realtime's convention. + bool report_partials = true; + + std::string family() const override { return "fake_audio"; } + rt::VoiceTaskKind task_kind() const override { return rt::VoiceTaskKind::Asr; } + rt::RunMode run_mode() const override { return rt::RunMode::Streaming; } + + void prepare(const rt::SessionPreparationRequest &request) override { + calls.push_back("prepare"); + prepared_ = true; + prepared_rate_ = request.audio.has_value() ? request.audio->sample_rate : 0; + } + + rt::StreamingPolicy streaming_policy() const override { return policy; } + + void set_stream_event_sink(rt::StreamEventCallback sink) override { + calls.push_back(sink ? "sink+" : "sink-"); + sink_ = std::move(sink); + } + + void reset() override { + if (!prepared_) { + // Exactly what silero_vad does, and the reason prepare() has to come + // first rather than being folded into session_for. + throw std::runtime_error("fake: prepare() must be called before reset()"); + } + calls.push_back("reset"); + seen_frames_ = 0; + seen_chunks_ = 0; + text_.clear(); + } + + rt::StreamEvent process_audio_chunk(const rt::AudioChunk &chunk) override { + calls.push_back("chunk"); + chunks.push_back(chunk); + ++seen_chunks_; + if (throw_on_chunk == seen_chunks_) { + throw std::runtime_error("fake: chunk failure"); + } + const int channels = chunk.channels > 0 ? chunk.channels : 1; + seen_frames_ += static_cast(chunk.samples.size()) / channels; + text_ += "w" + std::to_string(seen_chunks_); + rt::StreamEvent event; + if (report_partials) { + event.partial_text = rt::Transcript{text_, "en"}; + } + return event; + } + + rt::TaskResult finalize() override { + calls.push_back("finalize"); + rt::TaskResult result; + result.text_output = rt::Transcript{ + text_ + "/frames=" + std::to_string(seen_frames_), "en"}; + return result; + } + + // Emits through the SINK the way nemotron_asr does, from inside the final + // step rather than from process_audio_chunk. + void emit_through_sink(const std::string &fragment) { + if (!sink_) { + return; + } + rt::StreamEvent event; + event.partial_text = rt::Transcript{fragment, "en"}; + sink_(event); + } + + bool sink_installed() const { return static_cast(sink_); } + int prepared_rate() const { return prepared_rate_; } + +private: + rt::StreamEventCallback sink_; + bool prepared_ = false; + int prepared_rate_ = 0; + std::int64_t seen_frames_ = 0; + int seen_chunks_ = 0; + std::string text_; +}; + +// nemotron_asr's shape: partials arrive only through the sink, and only from +// inside the finalize step. +class SinkOnlyAudioSession : public FakeAudioSession { +public: + SinkOnlyAudioSession() { report_partials = false; } + + rt::TaskResult finalize() override { + emit_through_sink("late "); + emit_through_sink("partial"); + return FakeAudioSession::finalize(); + } +}; + +// Pulls events, like every streaming TTS family. Overrides start_stream the way +// the seven real families do, calling reset() first. +class FakePullSession : public rt::IStreamingVoiceTaskSession { +public: + std::vector calls; + std::size_t event_count = 3; + bool final_on_second = false; + + std::string family() const override { return "fake_pull"; } + rt::VoiceTaskKind task_kind() const override { return rt::VoiceTaskKind::Tts; } + rt::RunMode run_mode() const override { return rt::RunMode::Streaming; } + + void prepare(const rt::SessionPreparationRequest &) override { + calls.push_back("prepare"); + prepared_ = true; + } + + rt::StreamingPolicy streaming_policy() const override { + rt::StreamingPolicy policy; + policy.input = rt::StreamingInputKind::None; + policy.output = rt::StreamingOutputKind::PullEvents; + return policy; + } + + void start_stream(const rt::TaskRequest &request) override { + calls.push_back("start_stream"); + (void)request; + reset(); + } + + void set_stream_event_sink(rt::StreamEventCallback sink) override { + calls.push_back(sink ? "sink+" : "sink-"); + sink_ = std::move(sink); + } + + void reset() override { + if (!prepared_) { + throw std::runtime_error("fake: prepare() must be called before reset()"); + } + calls.push_back("reset"); + emitted_ = 0; + } + + std::optional next_stream_event() override { + if (emitted_ >= event_count) { + return std::nullopt; + } + rt::StreamEvent event; + rt::AudioBuffer audio; + audio.sample_rate = 24000; + audio.channels = 1; + audio.samples.assign(4, 0.25F); + // named_audio_outputs, NOT audio_output: this is where supertonic, + // omnivoice and voxcpm2 all put their streamed chunks. + event.named_audio_outputs.push_back( + {"chunk_" + std::to_string(emitted_), std::move(audio), {}}); + ++emitted_; + if (final_on_second && emitted_ == 2) { + event.is_final = true; + } + calls.push_back("pull"); + return event; + } + + rt::StreamEvent process_audio_chunk(const rt::AudioChunk &) override { + throw std::runtime_error("fake_pull consumes no audio"); + } + + rt::TaskResult finalize() override { + calls.push_back("finalize"); + rt::TaskResult result; + rt::AudioBuffer merged; + merged.sample_rate = 24000; + merged.channels = 1; + merged.samples.assign(4 * emitted_, 0.25F); + result.audio_output = std::move(merged); + return result; + } + + bool sink_installed() const { return static_cast(sink_); } + +private: + rt::StreamEventCallback sink_; + bool prepared_ = false; + std::size_t emitted_ = 0; +}; + +// -------------------------------------------------------------------------- +// Helpers +// -------------------------------------------------------------------------- + +static audiocpp_backend::LoadedModel::Session +streaming_session(rt::IStreamingVoiceTaskSession &fake, + audiocpp_backend::Task task) { + audiocpp_backend::LoadedModel::Session session; + session.task = task; + session.mode = audiocpp_backend::Mode::Streaming; + session.streaming = &fake; + return session; +} + +static rt::TaskRequest audio_request(int sample_rate, int channels, + std::int64_t frames) { + rt::TaskRequest request; + rt::AudioBuffer audio; + audio.sample_rate = sample_rate; + audio.channels = channels; + audio.samples.assign(static_cast(frames * channels), 0.5F); + request.audio_input = std::move(audio); + return request; +} + +// -------------------------------------------------------------------------- +// Tests +// -------------------------------------------------------------------------- + +static void test_begin_stream_prepares_then_starts() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakePullSession fake; + const auto session = streaming_session(fake, audiocpp_backend::Task::Tts); + + rt::TaskRequest request; + audiocpp_backend::begin_stream(session, request, entry); + + check_eq(join(fake.calls), "prepare|start_stream|reset", + "begin_stream prepares before it starts, and start_stream resets"); +} + +// The base implementation of start_stream IS a reset(). FakeAudioSession does +// not override start_stream, so this is that guarantee, read out of the pinned +// header rather than assumed. +static void test_base_start_stream_resets() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakeAudioSession fake; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + + audiocpp_backend::begin_stream(session, audio_request(16000, 1, 10), entry); + check_eq(join(fake.calls), "prepare|reset", + "the interface's own start_stream resets the session"); +} + +static void test_begin_stream_refuses_a_non_streaming_session() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + audiocpp_backend::LoadedModel::Session session; + session.mode = audiocpp_backend::Mode::Offline; + + bool threw_capability = false; + try { + rt::TaskRequest request; + audiocpp_backend::begin_stream(session, request, entry); + } catch (const audiocpp_backend::CapabilityError &) { + threw_capability = true; + } catch (const std::exception &) { + } + check(threw_capability, + "begin_stream on an offline session throws CapabilityError, not a null deref"); +} + +// THE ONE THIS TASK IS ABOUT. A streaming session is cached, so the second +// stream gets the object the first one left behind. Two identical runs against +// the SAME session must produce identical output. +static void test_a_refetched_session_replays_identically() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakeAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 512; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + const auto request = audio_request(16000, 1, 1536); + + std::vector first_fragments; + const auto first = audiocpp_backend::run_streaming_audio( + session, request, *request.audio_input, + [&](const rt::StreamEvent &event) { + if (event.partial_text.has_value()) { + first_fragments.push_back(event.partial_text->text); + } + }, + entry); + + std::vector second_fragments; + const auto second = audiocpp_backend::run_streaming_audio( + session, request, *request.audio_input, + [&](const rt::StreamEvent &event) { + if (event.partial_text.has_value()) { + second_fragments.push_back(event.partial_text->text); + } + }, + entry); + + check_eq(join(second_fragments), join(first_fragments), + "a re-fetched streaming session replays the same partials"); + check_eq(second.text_output.has_value() ? second.text_output->text : "", + first.text_output.has_value() ? first.text_output->text : "", + "a re-fetched streaming session replays the same final text"); + check_eq(first.text_output.has_value() ? first.text_output->text : "", + "w1w2w3/frames=1536", + "the first run saw exactly the audio it was given"); + // Not a tautology: without the reset the second run reports six words and + // 3072 frames, and both checks above fail. + check_eq(join(first_fragments), "w1|w1w2|w1w2w3", "cumulative partials"); +} + +static void test_run_streaming_audio_installs_and_clears_the_sink() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + SinkOnlyAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 1024; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + const auto request = audio_request(16000, 1, 1024); + + std::vector fragments; + const auto result = audiocpp_backend::run_streaming_audio( + session, request, *request.audio_input, + [&](const rt::StreamEvent &event) { + if (event.partial_text.has_value()) { + fragments.push_back(event.partial_text->text); + } + }, + entry); + + check_eq(join(fragments), "late |partial", + "a family that reports only through the sink is not silent"); + check(!fake.sink_installed(), + "the sink is cleared before returning, so the cached session holds no " + "reference to the caller's frame"); + check_eq(join(fake.calls), "sink+|prepare|reset|chunk|finalize|sink-", + "the sink is installed before the stream begins and cleared after it ends"); + check(result.text_output.has_value(), "the final result still comes back"); +} + +static void test_the_sink_is_cleared_when_the_stream_throws() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakeAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 512; + fake.throw_on_chunk = 1; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + const auto request = audio_request(16000, 1, 1024); + + bool threw = false; + try { + audiocpp_backend::run_streaming_audio( + session, request, *request.audio_input, + [](const rt::StreamEvent &) {}, entry); + } catch (const std::exception &) { + threw = true; + } + check(threw, "a failing chunk propagates"); + check(!fake.sink_installed(), + "the sink is cleared on the exception path too"); +} + +static void test_chunking_honours_the_policy_sample_count() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakeAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 16000; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + // 2.5 chunks, so the last one is short. + const auto request = audio_request(16000, 1, 40000); + + audiocpp_backend::run_streaming_audio(session, request, *request.audio_input, + [](const rt::StreamEvent &) {}, entry); + + check(fake.chunks.size() == 3, "40000 frames at 16000 per chunk is three chunks"); + if (fake.chunks.size() == 3) { + check(fake.chunks[0].samples.size() == 16000, "first chunk is full"); + check(fake.chunks[1].samples.size() == 16000, "second chunk is full"); + check(fake.chunks[2].samples.size() == 8000, "last chunk is the remainder"); + check(fake.chunks[0].start_sample == 0, "first chunk starts at zero"); + check(fake.chunks[1].start_sample == 16000, "second chunk start index"); + check(fake.chunks[2].start_sample == 32000, "third chunk start index"); + check(fake.chunks[0].sample_rate == 16000, "chunk carries the buffer's rate"); + check(fake.chunks[0].channels == 1, "chunk carries the buffer's channel count"); + } +} + +// A buffer whose float count is not a whole number of frames is REFUSED rather +// than truncated. The integer division would otherwise drop the tail floats +// from the fed audio, and therefore from the transcript, with no diagnostic. +static void test_a_partial_trailing_frame_is_refused() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakeAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 100; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + + // 501 floats across 2 channels: 250 whole frames and one stray float. + rt::TaskRequest request; + rt::AudioBuffer audio; + audio.sample_rate = 48000; + audio.channels = 2; + audio.samples.assign(501, 0.5F); + request.audio_input = std::move(audio); + + bool threw_config = false; + try { + audiocpp_backend::run_streaming_audio(session, request, *request.audio_input, + [](const rt::StreamEvent &) {}, entry); + } catch (const audiocpp_backend::ConfigError &) { + threw_config = true; + } catch (const std::exception &) { + } + check(threw_config, + "a buffer that is not a whole number of frames is refused with ConfigError"); + check(fake.calls.empty(), + "the refusal precedes every call into the session, so no half-started " + "stream is left on the cached one"); +} + +// higgs_audio_stt states its window in seconds and leaves the sample count at +// zero, so this branch is a real family's path rather than a defensive one. +static void test_chunking_falls_back_to_the_policy_seconds() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakeAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 0; + fake.policy.preferred_audio_chunk_seconds = 4.0; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + const auto request = audio_request(16000, 1, 96000); // 6 s + + audiocpp_backend::run_streaming_audio(session, request, *request.audio_input, + [](const rt::StreamEvent &) {}, entry); + + check(fake.chunks.size() == 2, "6 s at a 4 s window is two chunks"); + if (fake.chunks.size() == 2) { + check(fake.chunks[0].samples.size() == 64000, "first window is 4 s"); + check(fake.chunks[1].samples.size() == 32000, "second window is the 2 s remainder"); + } +} + +static void test_chunking_falls_back_to_the_interface_default() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakeAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 0; + fake.policy.preferred_audio_chunk_seconds = 0.0; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + const auto request = audio_request(16000, 1, 1024); + + audiocpp_backend::run_streaming_audio(session, request, *request.audio_input, + [](const rt::StreamEvent &) {}, entry); + + check(fake.chunks.size() == 2, "a policy naming no window uses the interface's 512"); + if (!fake.chunks.empty()) { + check(fake.chunks[0].samples.size() == 512, "default window is 512 frames"); + } +} + +// A zero sample rate must not turn a seconds-only policy into a zero-length +// chunk, which would loop forever. +static void test_a_seconds_policy_with_no_rate_falls_through() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakeAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 0; + fake.policy.preferred_audio_chunk_seconds = 4.0; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + const auto request = audio_request(0, 1, 1024); + + audiocpp_backend::run_streaming_audio(session, request, *request.audio_input, + [](const rt::StreamEvent &) {}, entry); + check(fake.chunks.size() == 2, "a rateless buffer still chunks at the default 512"); +} + +// FRAMES, not floats. vibevoice_asr refuses a chunk whose sample count is not +// divisible by its channel count, and offsets every span it reports by the +// chunk's start_sample. +static void test_stereo_chunks_are_frame_aligned() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakeAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 300; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + const auto request = audio_request(48000, 2, 750); + + audiocpp_backend::run_streaming_audio(session, request, *request.audio_input, + [](const rt::StreamEvent &) {}, entry); + + check(fake.chunks.size() == 3, "750 frames at 300 frames per chunk is three chunks"); + for (const auto &chunk : fake.chunks) { + check(chunk.samples.size() % 2 == 0, "every stereo chunk is a whole number of frames"); + } + if (fake.chunks.size() == 3) { + check(fake.chunks[0].samples.size() == 600, "300 stereo frames is 600 floats"); + check(fake.chunks[1].start_sample == 300, + "start_sample counts frames, not floats"); + check(fake.chunks[2].samples.size() == 300, "the remainder is 150 frames"); + } +} + +static void test_pull_drains_every_event_and_installs_no_sink() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakePullSession fake; + fake.event_count = 3; + const auto session = streaming_session(fake, audiocpp_backend::Task::Tts); + + std::vector ids; + rt::TaskRequest request; + const auto result = audiocpp_backend::run_streaming_pull( + session, request, + [&](const rt::StreamEvent &event) { + for (const auto &named : event.named_audio_outputs) { + ids.push_back(named.id); + } + }, + entry); + + check_eq(join(ids), "chunk_0|chunk_1|chunk_2", "every pulled event reaches the caller"); + check(!fake.sink_installed(), + "no stream event sink is installed on the pull path, so voxcpm2 cannot " + "deliver every chunk twice"); + check_eq(join(fake.calls), "prepare|start_stream|reset|pull|pull|pull|finalize", + "prepare, start, drain, finish"); + check(result.audio_output.has_value(), "the merged result comes back"); + check(result.audio_output.has_value() && result.audio_output->samples.size() == 12, + "the merged result is the whole synthesis, not a tail"); +} + +static void test_pull_stops_on_a_final_event() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakePullSession fake; + fake.event_count = 5; + fake.final_on_second = true; + const auto session = streaming_session(fake, audiocpp_backend::Task::Tts); + + int events = 0; + rt::TaskRequest request; + audiocpp_backend::run_streaming_pull( + session, request, [&](const rt::StreamEvent &) { ++events; }, entry); + + check(events == 2, "an event marked final ends the pull loop"); +} + +static void test_pull_refuses_a_non_streaming_session() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + audiocpp_backend::LoadedModel::Session session; + + bool threw_capability = false; + try { + rt::TaskRequest request; + audiocpp_backend::run_streaming_pull( + session, request, [](const rt::StreamEvent &) {}, entry); + } catch (const audiocpp_backend::CapabilityError &) { + threw_capability = true; + } catch (const std::exception &) { + } + check(threw_capability, "run_streaming_pull refuses a session with no streaming half"); +} + +// prepare() runs on EVERY stream, not once per session: the preparation request +// is derived from the request (audio contract, text, voice), so a second stream +// at a different rate would otherwise run against the first one's contract. +static void test_prepare_tracks_the_request_not_the_session() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakeAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 4096; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + + const auto first = audio_request(16000, 1, 4096); + audiocpp_backend::run_streaming_audio(session, first, *first.audio_input, + [](const rt::StreamEvent &) {}, entry); + check(fake.prepared_rate() == 16000, "the first stream prepares at its own rate"); + + const auto second = audio_request(44100, 1, 4096); + audiocpp_backend::run_streaming_audio(session, second, *second.audio_input, + [](const rt::StreamEvent &) {}, entry); + check(fake.prepared_rate() == 44100, + "the second stream prepares at ITS rate, not the first one's"); +} + +// -------------------------------------------------------------------------- +// run_streaming_live +// -------------------------------------------------------------------------- + +// Hands the driver a fixed list of wire frames, the way a client's audio +// callback would, and then closes. +static std::function &)> +frames_from(const std::vector &sizes) { + auto index = std::make_shared(0); + auto list = std::make_shared>(sizes); + return [index, list](std::vector &out) { + if (*index >= list->size()) { + return false; + } + out.assign((*list)[*index], 0.5F); + ++*index; + return true; + }; +} + +static rt::TaskRequest live_request(int sample_rate, int channels) { + rt::TaskRequest request; + rt::AudioBuffer contract; + contract.sample_rate = sample_rate; + contract.channels = channels; + request.audio_input = std::move(contract); // no samples: none exist yet + return request; +} + +// The wire's frame size is a property of the client's audio callback. The +// family's window is a statement about what it can decode. The driver feeds the +// second, not the first. +static void test_live_buffers_wire_frames_into_policy_windows() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakeAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 1600; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + + // Ten 512-sample frames: 5120 samples, i.e. three full 1600 windows and a + // 320 sample tail. + const auto result = audiocpp_backend::run_streaming_live( + session, live_request(16000, 1), + frames_from(std::vector(10, 512)), + [](const rt::StreamEvent &) {}, entry); + + check(fake.chunks.size() == 4, + "5120 wire samples at a 1600 frame window is three windows and a tail"); + if (fake.chunks.size() == 4) { + check(fake.chunks[0].samples.size() == 1600, "first window is full"); + check(fake.chunks[2].samples.size() == 1600, "third window is full"); + check(fake.chunks[3].samples.size() == 320, "the tail is what was left"); + check(fake.chunks[0].start_sample == 0, "the first window starts at zero"); + check(fake.chunks[1].start_sample == 1600, "start_sample counts frames"); + check(fake.chunks[3].start_sample == 4800, "the tail is offset by all of it"); + check(fake.chunks[0].sample_rate == 16000, "the chunk carries the session rate"); + } + check(result.text_output.has_value() && + result.text_output->text == "w1w2w3w4/frames=5120", + "every wire sample reaches the family exactly once"); +} + +// nemotron_asr's shape: no partials from process_audio_chunk, every one of them +// through the sink from inside finalize. +static void test_live_installs_and_clears_the_sink() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + SinkOnlyAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 512; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + + std::vector fragments; + audiocpp_backend::run_streaming_live( + session, live_request(16000, 1), frames_from({512}), + [&](const rt::StreamEvent &event) { + if (event.partial_text.has_value()) { + fragments.push_back(event.partial_text->text); + } + }, + entry); + + check_eq(join(fragments), "late |partial", + "a family that reports only through the sink is not silent live either"); + check(!fake.sink_installed(), + "the sink is cleared before returning, so the cached session holds no " + "reference to this call's frame"); + check_eq(join(fake.calls), "sink+|prepare|reset|chunk|finalize|sink-", + "sink installed before the stream begins, cleared after it ends"); +} + +// A client that opens a session and closes it without speaking. finalize is NOT +// called: nemotron_asr throws "finalize requires streamed audio", and an empty +// transcript is the truthful answer to transcribing nothing. +static void test_live_with_no_audio_never_finalizes() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakeAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 512; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + + const auto result = audiocpp_backend::run_streaming_live( + session, live_request(16000, 1), frames_from({}), + [](const rt::StreamEvent &) {}, entry); + + check_eq(join(fake.calls), "sink+|prepare|reset|sink-", + "an empty live stream begins and ends without a chunk or a finalize"); + check(!result.text_output.has_value(), + "an empty live stream reports no transcript rather than an error"); +} + +// A tail shorter than a window is still fed. Without this the last fragment of +// speech never reaches the model, and nothing says so. +static void test_live_feeds_a_short_tail() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakeAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 16000; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + + audiocpp_backend::run_streaming_live(session, live_request(16000, 1), + frames_from({100, 200}), + [](const rt::StreamEvent &) {}, entry); + + check(fake.chunks.size() == 1, + "300 samples against a 16000 frame window is one short chunk, not none"); + if (!fake.chunks.empty()) { + check(fake.chunks[0].samples.size() == 300, "the tail carries everything fed"); + } +} + +// A live request carries no samples, so the CONTRACT is the only thing that says +// what rate the frames are in, and prepare() needs it: nemotron_asr's streaming +// prepare throws without one. +static void test_live_prepares_at_the_contract_rate() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakeAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 512; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + + audiocpp_backend::run_streaming_live(session, live_request(16000, 1), + frames_from({512}), + [](const rt::StreamEvent &) {}, entry); + check(fake.prepared_rate() == 16000, + "the empty contract buffer still carries the rate into prepare()"); + + bool threw_config = false; + try { + rt::TaskRequest bare; // no audio_input at all + audiocpp_backend::run_streaming_live(session, bare, frames_from({512}), + [](const rt::StreamEvent &) {}, entry); + } catch (const audiocpp_backend::ConfigError &) { + threw_config = true; + } catch (const std::exception &) { + } + check(threw_config, "a live request with no audio contract is refused"); +} + +// The pull function is the gRPC read, and a request the handler has to refuse +// mid-stream unwinds through the driver. The sink must not survive it. +static void test_live_clears_the_sink_when_the_puller_throws() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakeAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 512; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + + bool threw = false; + try { + audiocpp_backend::run_streaming_live( + session, live_request(16000, 1), + [](std::vector &) -> bool { + throw std::runtime_error("fake: the client vanished"); + }, + [](const rt::StreamEvent &) {}, entry); + } catch (const std::exception &) { + threw = true; + } + check(threw, "a failing pull propagates"); + check(!fake.sink_installed(), "the sink is cleared on the pull's exception path"); +} + +// Two live streams over the SAME cached session must not run into each other. +static void test_live_replays_identically_on_a_refetched_session() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + FakeAudioSession fake; + fake.policy.preferred_audio_chunk_samples = 512; + const auto session = streaming_session(fake, audiocpp_backend::Task::Asr); + + const auto first = audiocpp_backend::run_streaming_live( + session, live_request(16000, 1), frames_from({512, 512}), + [](const rt::StreamEvent &) {}, entry); + const auto second = audiocpp_backend::run_streaming_live( + session, live_request(16000, 1), frames_from({512, 512}), + [](const rt::StreamEvent &) {}, entry); + + check_eq(second.text_output.has_value() ? second.text_output->text : "", + first.text_output.has_value() ? first.text_output->text : "", + "a re-fetched live session replays the same transcript"); + // Not a tautology: without the reset the second run reports w1..w4 and 2048 + // frames. + check_eq(first.text_output.has_value() ? first.text_output->text : "", + "w1w2/frames=1024", "the first live run saw exactly what was fed"); +} + +static void test_live_refuses_a_non_streaming_session() { + audiocpp_backend::InferenceLane lane("test"); + audiocpp_backend::LaneEntry entry(lane, 0); + audiocpp_backend::LoadedModel::Session session; + session.mode = audiocpp_backend::Mode::Offline; + + bool threw_capability = false; + try { + audiocpp_backend::run_streaming_live(session, live_request(16000, 1), + frames_from({512}), + [](const rt::StreamEvent &) {}, entry); + } catch (const audiocpp_backend::CapabilityError &) { + threw_capability = true; + } catch (const std::exception &) { + } + check(threw_capability, + "run_streaming_live refuses a session with no streaming half"); +} + +// -------------------------------------------------------------------------- +// resolve_model_path +// -------------------------------------------------------------------------- +// +// It lives in loaded_model.cpp and is a pure (dir, file, name) -> string, so it +// is tested here rather than in a standalone unit: that file cannot compile +// without the engine headers. +// +// It is tested at all because it shipped a bug no test could have caught. THE +// SHAPES BELOW ARE THE PRODUCTION SHAPES, not convenient ones, and that +// distinction is the entire point. Task 15 verified the bundled: form with a +// hand-written LoadModel that left ModelFile empty, which is the one shape the +// server never produces: pkg/model/loader.go's LoadModelWithFile always fills +// ModelFile with filepath.Join(ModelPath, model), and core/backend/options.go +// only overrides it for a managed artifact. The first case below is therefore +// the regression test; the other three are what it must not have broken. + +// std::string::ends_with is C++20 and this target is C++17. +static bool ends_with(const std::string &value, const std::string &suffix) { + return value.size() >= suffix.size() && + value.compare(value.size() - suffix.size(), suffix.size(), suffix) == 0; +} + +// THE REGRESSION CASE. What a model YAML saying `model: bundled:silero_vad` +// actually arrives as: Model intact, ModelFile joined onto the models directory. +static void test_bundled_in_model_survives_a_joined_model_file() { + const std::string resolved = audiocpp_backend::resolve_model_path( + "/models", "/models/bundled:silero_vad", "bundled:silero_vad"); + + check(ends_with(resolved, "/assets/silero_vad"), + "bundled: in Model resolves under the package assets dir (got \"" + + resolved + "\")"); + // Checked separately from the suffix because this is the failure that + // shipped: the joined ModelFile came back verbatim and the load died on + // "model path does not exist: /models/bundled:silero_vad". + check(resolved.find("/models/") == std::string::npos, + "bundled: in Model is not resolved against the models directory (got \"" + + resolved + "\")"); +} + +// Task 15's shape: the form in ModelFile with Model empty. It worked before the +// fix and must keep working. +static void test_bundled_in_model_file_still_resolves() { + const std::string resolved = + audiocpp_backend::resolve_model_path("", "bundled:marblenet_vad", ""); + + check(ends_with(resolved, "/assets/marblenet_vad"), + "bundled: in ModelFile still resolves under the package assets dir (got \"" + + resolved + "\")"); +} + +// The ordinary case, and the one the bundled: lookup must not capture: a real +// artifact path in ModelFile with a plain name in Model. +static void test_a_plain_name_resolves_to_the_model_file() { + const std::string resolved = audiocpp_backend::resolve_model_path( + "/models", "/models/chatterbox-q8_0.gguf", "chatterbox-q8_0.gguf"); + + check_eq(resolved, "/models/chatterbox-q8_0.gguf", + "a plain name resolves to the absolute ModelFile"); +} + +// A relative ModelFile is still joined onto ModelPath. The fix does not touch +// this branch, which is why it is pinned: the bundled: lookup now runs before it +// and has to fall through for every non-bundled input. +static void test_a_relative_model_file_joins_the_model_path() { + const std::string resolved = audiocpp_backend::resolve_model_path( + "/models", "sub/nemotron-asr-q8_0.gguf", "nemotron-asr"); + + check_eq(resolved, "/models/sub/nemotron-asr-q8_0.gguf", + "a relative ModelFile joins the models directory"); +} + +int main() { + test_bundled_in_model_survives_a_joined_model_file(); + test_bundled_in_model_file_still_resolves(); + test_a_plain_name_resolves_to_the_model_file(); + test_a_relative_model_file_joins_the_model_path(); + test_begin_stream_prepares_then_starts(); + test_base_start_stream_resets(); + test_begin_stream_refuses_a_non_streaming_session(); + test_a_refetched_session_replays_identically(); + test_run_streaming_audio_installs_and_clears_the_sink(); + test_the_sink_is_cleared_when_the_stream_throws(); + test_chunking_honours_the_policy_sample_count(); + test_a_partial_trailing_frame_is_refused(); + test_chunking_falls_back_to_the_policy_seconds(); + test_chunking_falls_back_to_the_interface_default(); + test_a_seconds_policy_with_no_rate_falls_through(); + test_stereo_chunks_are_frame_aligned(); + test_pull_drains_every_event_and_installs_no_sink(); + test_pull_stops_on_a_final_event(); + test_pull_refuses_a_non_streaming_session(); + test_prepare_tracks_the_request_not_the_session(); + test_live_buffers_wire_frames_into_policy_windows(); + test_live_installs_and_clears_the_sink(); + test_live_with_no_audio_never_finalizes(); + test_live_feeds_a_short_tail(); + test_live_prepares_at_the_contract_rate(); + test_live_clears_the_sink_when_the_puller_throws(); + test_live_replays_identically_on_a_refetched_session(); + test_live_refuses_a_non_streaming_session(); + if (failures) { + fprintf(stderr, "%d check(s) failed\n", failures); + return 1; + } + fprintf(stderr, "all streaming driver checks passed\n"); + return 0; +} diff --git a/backend/cpp/audio-cpp/transcript_assembly.cpp b/backend/cpp/audio-cpp/transcript_assembly.cpp new file mode 100644 index 000000000..5349ec790 --- /dev/null +++ b/backend/cpp/audio-cpp/transcript_assembly.cpp @@ -0,0 +1,215 @@ +#include "transcript_assembly.h" + +#include "audio_units.h" + +#include +#include +#include + +namespace audiocpp_backend { +namespace { + +std::int64_t midpoint(const Span &span) { + return span.start_sample + (span.end_sample - span.start_sample) / 2; +} + +bool contains(const Span &span, std::int64_t sample) { + return sample >= span.start_sample && sample < span.end_sample; +} + +std::int64_t overlap(const Span &a, const Span &b) { + const std::int64_t begin = std::max(a.start_sample, b.start_sample); + const std::int64_t end = std::min(a.end_sample, b.end_sample); + return end > begin ? end - begin : 0; +} + +// Joins a segment's words into that segment's text. +// +// THE SEPARATOR IS NOT ALWAYS A SPACE, and getting it wrong is visible to every +// caller rather than cosmetic: core/http/endpoints/openai/transcription.go +// routes response_format text, srt, vtt and lrc through +// schema.TranscriptionResponse, which builds the entire body out of +// Segments[].Text and never reads the top-level text. For those four formats +// the segment text IS the response. +// +// Two producer conventions have to be told apart: +// +// whole words "Some", "call", "me" -> join with a space +// subword pieces "So", "me", " call" -> concatenate +// +// The second is SentencePiece, where a word boundary is carried as a LEADING +// SPACE on the piece; nemotron_asr emits one entry per token in exactly that +// form. Space-joining those produced "So me call me na ture ,", which is +// what response_format=text returned while the correct sentence sat unread in +// the top-level field. Concatenating them reproduces text_output exactly. +// +// The convention is read off the words themselves, because nothing else in the +// result declares it. One leading space anywhere is enough to decide: a +// whole-word producer has no reason to emit one, and a subword producer emits +// one at every word boundary, so the two populations do not overlap. A producer +// that mixed both conventions inside one segment could not be served correctly +// by any single separator; this picks concatenation for it. +// +// This does NOT touch the top-level text, which stays text_output verbatim. The +// rule that forbids deriving the transcript from the segments is about the +// direction segments -> text. Segment text has no source other than its words +// and is necessarily derived. +std::string join_words(const std::vector &words) { + const bool subword_pieces = + std::any_of(words.begin(), words.end(), [](const OutWord &word) { + return !word.text.empty() && word.text.front() == ' '; + }); + std::string out; + for (const auto &word : words) { + if (word.text.empty()) { + continue; + } + if (!subword_pieces && !out.empty()) { + out += " "; + } + out += word.text; + } + return out; +} + +std::string speaker_for(const Span &segment, + const std::vector &turns) { + std::string best; + std::int64_t best_overlap = 0; + for (const auto &turn : turns) { + const std::int64_t shared = overlap(segment, turn.span); + if (shared > best_overlap) { + best_overlap = shared; + best = turn.speaker; + } + } + return best; +} + +// The chosen segmentation. labels is empty unless the spans were sourced from +// the speaker turns themselves, in which case it is parallel to spans and holds +// the label each span arrived with. +struct SegmentSource { + std::vector spans; + std::vector labels; +}; + +// Chooses the segment spans, per the documented precedence. +SegmentSource choose_segment_spans(const std::string &text_output, + const std::vector &speech_segments, + const std::vector &speaker_turns, + const std::vector &words) { + if (!speech_segments.empty()) { + return {speech_segments, {}}; + } + if (!speaker_turns.empty()) { + // The labels are carried out rather than re-derived by overlap later. A + // turn wholly contained in another speaker's turn overlaps its own span + // completely, which is the largest overlap possible, so it can only tie + // with the containing turn and would then lose that tie on order. + // sortformer_diar binarizes each speaker's track independently and + // sorts the result by start sample, so the container always comes + // first, and the interjecting speaker would be silently relabelled to + // the speaker it interrupted. + SegmentSource source; + source.spans.reserve(speaker_turns.size()); + source.labels.reserve(speaker_turns.size()); + for (const auto &turn : speaker_turns) { + source.spans.push_back(turn.span); + source.labels.push_back(turn.speaker); + } + return source; + } + if (!words.empty()) { + Span covering = words.front().span; + for (const auto &word : words) { + covering.start_sample = + std::min(covering.start_sample, word.span.start_sample); + covering.end_sample = std::max(covering.end_sample, word.span.end_sample); + } + return {{covering}, {}}; + } + if (!text_output.empty()) { + // A zero span rather than a fabricated duration: the model reported no + // timing, and inventing one would be a lie the caller cannot detect. + return {{Span{0, 0}}, {}}; + } + return {}; +} + +// Returns the index of the segment a word belongs to, or the nearest segment +// when the word falls outside all of them. +size_t segment_index_for_word(const std::vector &spans, const Span &word) { + const std::int64_t centre = midpoint(word); + for (size_t i = 0; i < spans.size(); ++i) { + if (contains(spans[i], centre)) { + return i; + } + } + size_t nearest = 0; + std::int64_t best_distance = std::numeric_limits::max(); + for (size_t i = 0; i < spans.size(); ++i) { + const std::int64_t distance = std::llabs(midpoint(spans[i]) - centre); + if (distance < best_distance) { + best_distance = distance; + nearest = i; + } + } + return nearest; +} + +} // namespace + +AssembledTranscript assemble_transcript(const std::string &text_output, + const std::vector &speech_segments, + const std::vector &speaker_turns, + const std::vector &words, + int sample_rate) { + AssembledTranscript assembled; + // THE RULE. Never derived from spans. + assembled.text = text_output; + + const SegmentSource source = + choose_segment_spans(text_output, speech_segments, speaker_turns, words); + const std::vector &spans = source.spans; + if (spans.empty()) { + return assembled; + } + + assembled.segments.resize(spans.size()); + for (size_t i = 0; i < spans.size(); ++i) { + OutSegment &segment = assembled.segments[i]; + segment.id = static_cast(i); + segment.start_ns = samples_to_nanoseconds(spans[i].start_sample, sample_rate); + segment.end_ns = samples_to_nanoseconds(spans[i].end_sample, sample_rate); + // A segment that came from a speaker turn already knows its speaker. + // Only the other three sources have to look one up by overlap. + segment.speaker = source.labels.empty() + ? speaker_for(spans[i], speaker_turns) + : source.labels[i]; + } + + for (const auto &word : words) { + const size_t index = segment_index_for_word(spans, word.span); + OutWord out; + out.start_ns = samples_to_nanoseconds(word.span.start_sample, sample_rate); + out.end_ns = samples_to_nanoseconds(word.span.end_sample, sample_rate); + out.text = word.word; + assembled.segments[index].words.push_back(out); + } + + for (auto &segment : assembled.segments) { + segment.text = join_words(segment.words); + } + + // A single segment with no word timing carries the whole transcript. With + // several segments there is no defensible way to split the text, so their + // per-segment text stays empty and only the top-level text is authoritative. + if (assembled.segments.size() == 1 && assembled.segments[0].words.empty()) { + assembled.segments[0].text = text_output; + } + + return assembled; +} + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/transcript_assembly.h b/backend/cpp/audio-cpp/transcript_assembly.h new file mode 100644 index 000000000..e58fa2b72 --- /dev/null +++ b/backend/cpp/audio-cpp/transcript_assembly.h @@ -0,0 +1,87 @@ +#pragma once + +// Builds LocalAI's TranscriptResult shape from audio.cpp's TaskResult spans. +// Standard library only; result_map.cpp converts the engine types into these +// PODs at the boundary. +// +// THE RULE: the top-level transcript text is text_output verbatim, always. +// audio.cpp carries transcript text in exactly one place, TaskResult.text_output. +// speech_segments, speaker_turns and word_timestamps carry spans and labels but +// no text. Deriving the top-level text by concatenating per-segment text +// therefore yields an empty transcript for every producer that reports segments +// without word timestamps, which includes VibeVoice diarized ASR. + +#include +#include +#include + +namespace audiocpp_backend { + +// Sample-index span, mirroring engine::runtime::TimeSpan. +struct Span { + std::int64_t start_sample = 0; + std::int64_t end_sample = 0; +}; + +struct WordSpan { + Span span; + std::string word; +}; + +struct SpeakerSpan { + Span span; + std::string speaker; +}; + +struct OutWord { + std::int64_t start_ns = 0; + std::int64_t end_ns = 0; + std::string text; +}; + +struct OutSegment { + int id = 0; + std::int64_t start_ns = 0; + std::int64_t end_ns = 0; + std::string text; + std::string speaker; + std::vector words; +}; + +struct AssembledTranscript { + std::string text; + std::vector segments; +}; + +// Segment source, first non-empty wins: +// 1. speech_segments +// 2. speaker_turns +// 3. one segment spanning all words, when words are present +// 4. one zero-span segment carrying the full text, when text is present +// 5. no segments +// +// Words attach to the segment whose range contains their midpoint; a word +// outside every segment attaches to the nearest one by midpoint distance so it +// is never silently dropped. A lone segment with no words carries the full text. +// +// A segment's text is its words joined, and the separator depends on the +// producer's convention: whole words ("Some", "call") are joined with a space, +// while SentencePiece-style subword pieces, which carry the word boundary as a +// LEADING SPACE (" call"), are concatenated. One leading space anywhere in the +// segment selects concatenation. This matters beyond tidiness: response_format +// text, srt, vtt and lrc build their entire body out of the segment text and +// never read the top-level text. +// +// A segment's speaker is the speaker turn with the greatest overlap, except +// when the segments came from the speaker turns themselves (source 2), where +// each segment keeps its own turn's label. Re-deriving it there loses a turn +// nested inside another speaker's turn: the nested turn overlaps its own span +// completely, so it can only tie with the containing turn, which is listed +// first and wins the tie. +AssembledTranscript assemble_transcript(const std::string &text_output, + const std::vector &speech_segments, + const std::vector &speaker_turns, + const std::vector &words, + int sample_rate); + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/transcript_assembly_test.cpp b/backend/cpp/audio-cpp/transcript_assembly_test.cpp new file mode 100644 index 000000000..a0e0ede49 --- /dev/null +++ b/backend/cpp/audio-cpp/transcript_assembly_test.cpp @@ -0,0 +1,532 @@ +// Unit tests for transcript_assembly. Standard library only. The harness +// compiles this as a single translation unit, so both implementations are +// included directly rather than linked. +// +// Every fixture below either mirrors a producer shape actually observed from +// audio.cpp families, and names the families it was checked against, or says in +// its own comment that it is defensive. Do not replace an observed shape with an +// invented one and do not quietly promote a defensive fixture to an observed +// one: an invented shape is what let the earlier attempt ship an empty +// transcript. + +#include "audio_units.cpp" +#include "transcript_assembly.cpp" + +#include +#include +#include +#include + +static int failures = 0; + +static void check(bool ok, const std::string &name) { + if (!ok) { + failures++; + fprintf(stderr, "FAIL: %s\n", name.c_str()); + } else { + fprintf(stderr, "ok: %s\n", name.c_str()); + } +} + +using namespace audiocpp_backend; + +// Indexed access that reports a named failure instead of running off the end. +// std::vector::operator[] past the end is undefined behaviour, so a regression +// that drops a segment would crash the process here and take every later check +// with it. Returning a default element keeps the rest of the suite reporting. +static const OutSegment &segment_at(const AssembledTranscript &out, size_t index, + const std::string &name) { + static const OutSegment missing; + if (index >= out.segments.size()) { + failures++; + fprintf(stderr, "FAIL: %s (segment %zu is missing)\n", name.c_str(), index); + return missing; + } + return out.segments[index]; +} + +static const OutWord &word_at(const OutSegment &segment, size_t index, + const std::string &name) { + static const OutWord missing; + if (index >= segment.words.size()) { + failures++; + fprintf(stderr, "FAIL: %s (word %zu is missing)\n", name.c_str(), index); + return missing; + } + return segment.words[index]; +} + +static const int kRate = 16000; + +// Shape A: word timestamps only. Emitted by nemotron_asr, qwen3_asr and +// qwen3_forced_aligner, all of which set text_output plus word_timestamps and +// leave speech_segments empty. +static void test_words_only() { + const std::vector words = { + {{0, 8000}, "hello"}, + {{8000, 16000}, "world"}, + }; + const auto out = assemble_transcript("hello world", {}, {}, words, kRate); + + check(out.text == "hello world", "text is text_output verbatim"); + check(out.segments.size() == 1, "words with no segments yield one segment"); + const OutSegment &first = segment_at(out, 0, "words only segment"); + check(first.start_ns == 0, "segment starts at the first word"); + check(first.end_ns == 1000000000LL, "segment ends at the last word"); + check(first.words.size() == 2, "both words attached"); + check(word_at(first, 0, "first word").text == "hello", "first word text"); + check(word_at(first, 1, "second word").start_ns == 500000000LL, + "second word start in ns"); + check(first.text == "hello world", "segment text joins its words"); + check(first.id == 0, "ids are zero based"); +} + +// Shape A-whole: the whole-word convention, stated explicitly rather than left +// implicit in the shape A tests. qwen3_forced_aligner emits one entry per WORD +// (processor.cpp parses per-word timestamp tokens), so its pieces carry no +// leading space and must be joined with one. +static void test_whole_words_are_space_joined() { + const std::vector words = { + {{0, 8000}, "Some"}, + {{8000, 16000}, "call"}, + {{16000, 24000}, "me"}, + }; + const auto out = assemble_transcript("Some call me", {}, {}, words, kRate); + + check(segment_at(out, 0, "whole words").text == "Some call me", + "whole words are joined with a single space"); +} + +// Shape A-subword: the SentencePiece convention, where the word boundary is a +// LEADING SPACE on the piece. These are the first eleven word_timestamps +// nemotron_asr actually returned for audio.cpp/assets/resources/sample_16k.wav +// with the q8_0 GGUF, copied verbatim rather than invented, including the lone +// " " piece at index 3. +// +// Space-joining these produced "So me call me na ture , other s call", +// which is not a cosmetic problem: response_format text, srt, vtt and lrc build +// their entire body from the segment text and never read the top-level text, so +// that string WAS the transcription response for those formats. +static void test_subword_pieces_are_concatenated() { + const std::vector words = { + {{15360, 16640}, "So"}, {{15360, 16640}, "me"}, + {{23040, 24320}, " call"}, {{28160, 29440}, " "}, + {{28160, 29440}, "me"}, {{30720, 32000}, " na"}, + {{33280, 34560}, "ture"}, {{35840, 37120}, ","}, + {{38400, 39680}, " other"}, {{40960, 42240}, "s"}, + {{43520, 44800}, " call"}, + }; + const auto out = assemble_transcript( + "Some call me nature, others call me mother nature.", {}, {}, words, kRate); + + check(out.text == "Some call me nature, others call me mother nature.", + "the top-level text is still text_output verbatim"); + check(segment_at(out, 0, "subword pieces").text == + "Some call me nature, others call", + "subword pieces are concatenated, reproducing text_output"); +} + +// One leading space anywhere decides for the whole segment. A subword producer +// emits a boundary space at every word start, so its first piece, which is +// sentence-initial, does not have one; keying off the first piece alone would +// therefore pick the wrong convention on every segment. +static void test_a_single_leading_space_selects_concatenation() { + const std::vector words = { + {{0, 8000}, "al"}, + {{8000, 16000}, "pha"}, + {{16000, 24000}, " beta"}, + }; + const auto out = assemble_transcript("alpha beta", {}, {}, words, kRate); + + check(segment_at(out, 0, "mixed").text == "alpha beta", + "a leading space on a later piece selects concatenation"); +} + +// Shape A': the same producer, but text_output is punctuated and cased while +// the word timestamps are not. qwen3_asr rebuilds text_output from its word +// list only when timestamps are requested, so the two genuinely differ; this +// pins the sole segment's text to its words rather than to the top-level text. +static void test_words_only_with_punctuated_text_output() { + const std::vector words = { + {{0, 8000}, "hello"}, + {{8000, 16000}, "world"}, + }; + const auto out = assemble_transcript("Hello, world!", {}, {}, words, kRate); + + check(out.text == "Hello, world!", "punctuated text_output is untouched"); + check(out.segments.size() == 1, "one segment"); + check(segment_at(out, 0, "punctuated segment").text == "hello world", + "a segment with words takes its text from the words, not text_output"); +} + +// Shape A'': a merged word list whose last word is not the one that ends +// latest. audio.cpp concatenates per-chunk word lists in chunk order +// (append_chunk_word_timestamps in framework/audio/chunking.cpp). It drops a +// word whose global start falls before the chunk's keep span, but it never +// clips a word's end to that boundary, so the last word kept from one chunk can +// outlast the first word kept from the next. The covering span must therefore +// be the extent of every word, not the span from the first to the last. +static void test_covering_span_spans_every_word() { + const std::vector words = { + {{0, 4000}, "a"}, + // Kept from the earlier chunk, ending past the chunk boundary. + {{4000, 10000}, "b"}, + // First word of the next chunk, shorter, so it ends earlier. + {{8000, 9000}, "c"}, + }; + const auto out = assemble_transcript("a b c", {}, {}, words, kRate); + + check(out.segments.size() == 1, "one covering segment"); + check(segment_at(out, 0, "covering segment").end_ns == 625000000LL, + "the covering span reaches the latest word end, not the last word's"); +} + +// Defensive, not observed: no pinned family emits a word with no text. +// nemotron_asr's build_token_timestamps (models/nemotron_asr/decoder.cpp:97) +// skips a token that decodes to an empty chunk before it ever becomes a +// WordTimestamp. join_words guards against one anyway, and an unexercised guard +// is a guard the next reader deletes as dead weight. +static void test_empty_word_contributes_no_separator() { + const std::vector words = { + {{0, 4000}, "alpha"}, + {{4000, 8000}, ""}, + {{8000, 12000}, "beta"}, + }; + const auto out = assemble_transcript("alpha beta", {}, {}, words, kRate); + + check(out.segments.size() == 1, "one segment"); + check(segment_at(out, 0, "sole segment").text == "alpha beta", + "an empty word adds no separator to the segment text"); + check(segment_at(out, 0, "sole segment").words.size() == 3, + "the empty word still reports its span"); +} + +// Shape B: speech segments, no words. Emitted by ASR families that report +// utterance boundaries without word-level timing. +static void test_segments_without_words() { + const std::vector segments = {{0, 16000}, {16000, 32000}}; + const auto out = assemble_transcript("one two three", segments, {}, {}, kRate); + + // The regression: this must NOT be empty. + check(out.text == "one two three", "multi-segment text is not empty"); + check(out.segments.size() == 2, "both segments survive"); + check(segment_at(out, 0, "first segment").end_ns == 1000000000LL, + "first segment ends at 1s"); + check(segment_at(out, 1, "second segment").start_ns == 1000000000LL, + "second segment starts at 1s"); + check(segment_at(out, 0, "first segment").text.empty(), + "per-segment text stays empty when there are no words to split by"); + check(segment_at(out, 1, "second segment").id == 1, "ids increment"); +} + +// Shape C: speech segments plus speaker turns, no words. This is the real +// VibeVoice diarized ASR shape that broke the earlier attempt. +static void test_segments_with_speaker_turns_no_words() { + const std::vector segments = {{0, 16000}, {16000, 32000}}; + const std::vector turns = { + {{0, 16000}, "SPEAKER_00"}, + {{16000, 32000}, "SPEAKER_01"}, + }; + const auto out = assemble_transcript("hi there", segments, turns, {}, kRate); + + check(out.text == "hi there", "diarized multi-segment text is not empty"); + check(out.segments.size() == 2, "two segments"); + check(segment_at(out, 0, "first diarized segment").speaker == "SPEAKER_00", + "first speaker assigned"); + check(segment_at(out, 1, "second diarized segment").speaker == "SPEAKER_01", + "second speaker assigned"); +} + +// Defensive, not observed: speech segments and speaker turns that disagree. +// vibevoice_asr builds each SpeakerTurn with turn.span = speech_segment.span in +// one loop (models/vibevoice_asr/session.cpp:965) and shifts and clips both +// lists identically when merging chunks, so in practice the two lists are 1:1 +// with identical spans. That is exactly why the shape C fixture above cannot +// show which list is the segment source: swapping the precedence there produces +// byte-identical output. This fixture pins the precedence, and it is the shape +// any future family that segments and diarizes separately would produce. +static void test_speech_segments_outrank_speaker_turns() { + const std::vector segments = {{0, 32000}}; + const std::vector turns = { + {{0, 16000}, "SPEAKER_00"}, + {{16000, 32000}, "SPEAKER_01"}, + }; + const auto out = assemble_transcript("hi there", segments, turns, {}, kRate); + + check(out.segments.size() == 1, + "speech segments decide the segmentation, not speaker turns"); + check(segment_at(out, 0, "single utterance").end_ns == 2000000000LL, + "the utterance keeps its own span"); +} + +// Defensive, not observed: no pinned family emits speaker turns and word +// timestamps together. It pins rule 2 against rule 3, which nothing else does: +// a diarized result is segmented by who spoke, and words only fill the turns in. +static void test_speaker_turns_outrank_words() { + const std::vector turns = { + {{0, 16000}, "SPEAKER_00"}, + {{16000, 32000}, "SPEAKER_01"}, + }; + const std::vector words = { + {{0, 8000}, "hi"}, + {{16000, 24000}, "there"}, + }; + const auto out = assemble_transcript("hi there", {}, turns, words, kRate); + + check(out.segments.size() == 2, "the two turns segment the result"); + check(segment_at(out, 0, "turn 0").text == "hi", "first turn takes its word"); + check(segment_at(out, 1, "turn 1").text == "there", "second turn takes its word"); +} + +// Shape D: text only. Emitted by ASR families that report no timing at all, +// such as hviske_asr and citrinet_asr. +static void test_text_only() { + const auto out = assemble_transcript("just text", {}, {}, {}, kRate); + + check(out.text == "just text", "text survives"); + check(out.segments.size() == 1, "a single synthetic segment is emitted"); + const OutSegment &only = segment_at(out, 0, "synthetic segment"); + check(only.start_ns == 0 && only.end_ns == 0, + "synthetic segment has zero span, not a fabricated duration"); + check(only.text == "just text", "the sole segment carries the full text"); +} + +// Shape E: speaker turns only, no speech segments and no text. This is +// sortformer_diar, reached through the Diarize RPC. +static void test_speaker_turns_only() { + const std::vector turns = { + {{0, 24000}, "0"}, + {{24000, 48000}, "1"}, + }; + const auto out = assemble_transcript("", {}, turns, {}, kRate); + + check(out.text.empty(), "no text is reported when the model produced none"); + check(out.segments.size() == 2, "turns become segments"); + check(segment_at(out, 0, "turn 0").speaker == "0", + "speaker label preserved verbatim"); + check(segment_at(out, 1, "turn 1").start_ns == 1500000000LL, + "second turn starts at 1.5s"); +} + +// Shape E', the same producer with one speaker talking over another. +// decode_sortformer_speaker_turns (models/sortformer_diar/postprocess.cpp) +// binarizes each speaker's probability track independently, which is the whole +// point of sortformer, then sorts the turns by start sample. So a turn can be +// wholly contained in another speaker's turn, and the containing turn always +// comes first. A segment sourced from a speaker turn must keep that turn's own +// label: re-deriving it by overlap can only ever tie with the containing turn, +// which then wins on order and silently erases the interjecting speaker. +static void test_nested_speaker_turn_keeps_its_own_label() { + const std::vector turns = { + {{0, 100000}, "speaker_0"}, + {{10000, 20000}, "speaker_1"}, + }; + const auto out = assemble_transcript("", {}, turns, {}, kRate); + + check(out.segments.size() == 2, "both turns become segments"); + check(segment_at(out, 0, "containing turn").speaker == "speaker_0", + "the containing turn keeps its label"); + check(segment_at(out, 1, "nested turn").speaker == "speaker_1", + "a turn nested inside another is not relabelled to the container"); +} + +// Shape F: nothing at all. A model that ran but produced no output must not +// crash or fabricate a segment. +static void test_empty() { + const auto out = assemble_transcript("", {}, {}, {}, kRate); + check(out.text.empty(), "empty stays empty"); + check(out.segments.empty(), "no segments are invented"); +} + +// Shape H: speech segments with no text and no words at all. This is the VAD +// path, silero_vad and marblenet_vad, which fill speech_segments and never +// touch text_output. It reaches the lone-segment rule with nothing to carry. +static void test_vad_segments_without_text() { + const std::vector segments = {{0, 16000}, {24000, 32000}}; + const auto out = assemble_transcript("", segments, {}, {}, kRate); + + check(out.text.empty(), "VAD reports no text"); + check(out.segments.size() == 2, "both speech regions survive"); + check(segment_at(out, 1, "second speech region").start_ns == 1500000000LL, + "second region starts at 1.5s"); + check(segment_at(out, 0, "first speech region").text.empty(), + "a VAD segment carries no text"); + + const std::vector one = {{0, 16000}}; + const auto single = assemble_transcript("", one, {}, {}, kRate); + check(single.segments.size() == 1, "a single speech region survives"); + check(segment_at(single, 0, "lone speech region").text.empty(), + "a lone VAD segment does not fabricate text"); +} + +// Shape G: segments and words together. Words are assigned by midpoint so a +// word straddling a boundary lands in exactly one segment. +static void test_words_distributed_into_segments() { + const std::vector segments = {{0, 16000}, {16000, 32000}}; + const std::vector words = { + {{0, 4000}, "alpha"}, + {{4000, 8000}, "beta"}, + // Straddles the boundary; midpoint 16000 falls in the second segment. + {{12000, 20000}, "gamma"}, + {{20000, 28000}, "delta"}, + }; + const auto out = assemble_transcript("alpha beta gamma delta", segments, {}, + words, kRate); + + check(out.text == "alpha beta gamma delta", "top level text unchanged"); + check(out.segments.size() == 2, "two segments"); + check(segment_at(out, 0, "first segment").words.size() == 2, + "first segment takes two words"); + check(segment_at(out, 1, "second segment").words.size() == 2, + "second segment takes two words"); + check(segment_at(out, 0, "first segment").text == "alpha beta", + "first segment text"); + check(segment_at(out, 1, "second segment").text == "gamma delta", + "boundary-straddling word lands by midpoint"); +} + +// The midpoint rule is not the same as either endpoint rule. "early" starts in +// the first segment but ends in the second, and "late" the other way round; +// each must land where its midpoint says, which no start-only or end-only rule +// reproduces. +static void test_words_assigned_by_midpoint_not_endpoint() { + const std::vector segments = {{0, 16000}, {16000, 32000}}; + const std::vector words = { + // Midpoint 12000 -> first segment, although it ends in the second. + {{4000, 20000}, "early"}, + // Midpoint 20000 -> second segment, although it starts in the first. + {{12000, 28000}, "late"}, + }; + const auto out = assemble_transcript("early late", segments, {}, words, kRate); + + check(segment_at(out, 0, "first segment").text == "early", + "a word ending past the boundary stays where its midpoint is"); + check(segment_at(out, 1, "second segment").text == "late", + "a word starting before the boundary follows its midpoint"); +} + +// A word outside every segment must still be reachable rather than dropped +// silently, so it attaches to the nearest segment by midpoint distance. +static void test_word_outside_all_segments() { + const std::vector segments = {{0, 16000}}; + const std::vector words = { + {{0, 8000}, "inside"}, + {{40000, 48000}, "outside"}, + }; + const auto out = assemble_transcript("inside outside", segments, {}, words, + kRate); + check(out.segments.size() == 1, "one segment"); + check(segment_at(out, 0, "sole segment").words.size() == 2, + "the stray word is not dropped"); +} + +// The fallback picks the nearest segment, which is not the same as picking the +// first. With one segment the two are indistinguishable, so this uses three and +// puts the stray word past the last one. +static void test_stray_word_goes_to_the_nearest_segment() { + const std::vector segments = {{0, 8000}, {8000, 16000}, {16000, 24000}}; + const std::vector words = { + // Midpoint 44000, nearest the third segment. + {{40000, 48000}, "trailing"}, + }; + const auto out = assemble_transcript("trailing", segments, {}, words, kRate); + + check(segment_at(out, 0, "first segment").words.empty(), + "the stray word does not fall back to the first segment"); + check(segment_at(out, 2, "third segment").text == "trailing", + "the stray word attaches to the nearest segment"); +} + +// "Nearest" is measured from the segment's midpoint, and it is neither "the +// first segment" nor "the last". A leading stray word is the case a +// trailing-only fixture cannot reach: forced-aligner words scored against VAD +// segments produce one, and with only trailing coverage it would land at the +// end of the transcript with the suite green. Here the leading word's nearest +// midpoint is the first segment while its nearest start is the second, and the +// trailing word's nearest midpoint is the third while its nearest end is the +// second, so no endpoint rule reproduces this assignment either. +static void test_stray_word_distance_is_measured_from_the_midpoint() { + const std::vector segments = {{0, 2000}, {8000, 200000}, {300000, 302000}}; + const std::vector words = { + {{4000, 6000}, "lead"}, + {{249000, 251000}, "trail"}, + }; + const auto out = assemble_transcript("lead trail", segments, {}, words, kRate); + + check(out.segments.size() == 3, "three segments"); + check(segment_at(out, 0, "first segment").text == "lead", + "the leading stray word goes to the nearest segment by midpoint"); + check(segment_at(out, 2, "third segment").text == "trail", + "the trailing stray word goes to the nearest segment by midpoint"); + check(segment_at(out, 1, "middle segment").words.empty(), + "the long middle segment claims neither stray word"); +} + +// Speaker assignment uses greatest overlap, not first match, so a turn that +// barely touches a segment does not win over one that covers it. +static void test_speaker_assigned_by_greatest_overlap() { + const std::vector segments = {{8000, 24000}}; + const std::vector turns = { + {{0, 9000}, "brief"}, // overlaps 1000 samples + {{9000, 24000}, "main"} // overlaps 15000 samples + }; + const auto out = assemble_transcript("x", segments, turns, {}, kRate); + check(out.segments.size() == 1, "one segment"); + check(segment_at(out, 0, "sole segment").speaker == "main", + "greatest overlap wins"); +} + +// A segment no turn touches gets no speaker rather than the label of whichever +// turn happened to be listed first. +static void test_segment_without_any_overlapping_turn_has_no_speaker() { + const std::vector segments = {{0, 8000}, {40000, 48000}}; + const std::vector turns = {{{0, 8000}, "SPEAKER_00"}}; + const auto out = assemble_transcript("x", segments, turns, {}, kRate); + + check(segment_at(out, 0, "overlapped segment").speaker == "SPEAKER_00", + "the overlapped segment is labelled"); + check(segment_at(out, 1, "unlabelled segment").speaker.empty(), + "a segment no turn overlaps is left unlabelled"); +} + +static void test_zero_sample_rate_is_safe() { + const std::vector segments = {{0, 16000}}; + const auto out = assemble_transcript("x", segments, {}, {}, 0); + check(out.segments.size() == 1, "a zero sample rate still yields the segment"); + const OutSegment &only = segment_at(out, 0, "sole segment"); + check(only.start_ns == 0 && only.end_ns == 0, + "unknown sample rate yields zero timings rather than garbage"); +} + +int main() { + test_words_only(); + test_whole_words_are_space_joined(); + test_subword_pieces_are_concatenated(); + test_a_single_leading_space_selects_concatenation(); + test_words_only_with_punctuated_text_output(); + test_covering_span_spans_every_word(); + test_empty_word_contributes_no_separator(); + test_segments_without_words(); + test_segments_with_speaker_turns_no_words(); + test_speech_segments_outrank_speaker_turns(); + test_speaker_turns_outrank_words(); + test_text_only(); + test_speaker_turns_only(); + test_nested_speaker_turn_keeps_its_own_label(); + test_empty(); + test_vad_segments_without_text(); + test_words_distributed_into_segments(); + test_words_assigned_by_midpoint_not_endpoint(); + test_word_outside_all_segments(); + test_stray_word_goes_to_the_nearest_segment(); + test_stray_word_distance_is_measured_from_the_midpoint(); + test_speaker_assigned_by_greatest_overlap(); + test_segment_without_any_overlapping_turn_has_no_speaker(); + test_zero_sample_rate_is_safe(); + if (failures) { + fprintf(stderr, "%d check(s) failed\n", failures); + return 1; + } + fprintf(stderr, "all transcript_assembly checks passed\n"); + return 0; +} diff --git a/backend/cpp/audio-cpp/upstream_absence_ctest.cpp b/backend/cpp/audio-cpp/upstream_absence_ctest.cpp new file mode 100644 index 000000000..51ddb3b06 --- /dev/null +++ b/backend/cpp/audio-cpp/upstream_absence_ctest.cpp @@ -0,0 +1,283 @@ +// Asserts the ABSENCES that capability_routing.cpp's five refusal reasons rest +// on, against the engine itself rather than against somebody's reading of it. +// +// Those reasons are prose making checkable claims about a pinned third-party +// checkout: "VoiceTaskKind has no codec entry", "no family advertises spk", +// "miocodec advertises only vc and s2s". Prose rots silently across an +// AUDIO_CPP_VERSION bump, and it rots in the worst possible place, since a +// refusal that states a false fact is worse than a bare UNIMPLEMENTED: it will +// be believed. One of the four claims this backend was planned against was +// already false when it was written ("streaming exists for tts and asr only" is +// contradicted by silero_vad). This test is what turns the next such change +// from a silent lie on the wire into a build failure. +// +// Pinned at audio.cpp e800d435d130dc776baf6f3e6129bb62b1495c89. What follows +// was true of that commit; a bump is exactly when it needs to be re-run. +// +// It links engine_runtime and queries make_default_registry(), touching only +// include/engine/framework/**, like every other unit in this backend. It loads +// no model and reads no file: advertise_loaders() is the path-free catalog +// upstream publishes for --list-loaders. +// +// FOUR CAVEATS, so nobody reads more into a green run than it earns: +// +// 1. It cannot assert the AudioToAudioStream reason. That one contrasts +// LocalAI's OpenAI-Realtime contract (conversation, system prompt, tool +// loop) with what audio.cpp's s2s families actually do, and semantics are +// not a queryable property. What IS asserted is the enumerable half: that +// s2s is advertised by exactly miocodec and vevo2, which is the clause the +// message names by hand. +// 2. It queries the LOADER catalog, while grpc-server.cpp reads the LOADED +// model's own capabilities(). The two agree today, cross-checked on the +// wire: this test asserts miocodec advertises {vc/offline, s2s/offline}, +// and a live LoadModel of miocodec-q8_0.gguf reports exactly +// "vc/offline, s2s/offline". A family whose loaded capabilities diverged +// from its advertisement would slip past, but nothing loads without a +// model file and a ctest cannot depend on one. +// 3. Capabilities need not come from a loader at all. +// src/framework/model_spec/metadata.cpp's advertised_capabilities() builds +// a CapabilitySet from a spec's "capabilities"/"tasks"/"modes" keys, which +// is a second route by which a bump could falsify claim 3. Today no +// shipped model_specs/*.json carries a top-level "tasks" key and no loader +// calls that function, so the route is dead. It is covered anyway to the +// extent that a loader adopting it would surface through advertise_loaders +// like any other capability, which is why every assertion below queries +// advertised capabilities rather than loader source. +// 4. An absence test passes trivially when the query is broken, so +// test_catalog_is_populated below is a POSITIVE control and is not +// optional. It proves the registry is non-empty and that the query does +// find a capability it should, before any absence is believed. + +#include "engine/framework/runtime/registry.h" +#include "engine/framework/runtime/session.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +int failures = 0; + +void check(bool ok, const std::string &name) { + if (!ok) { + failures++; + fprintf(stderr, "FAIL: %s\n", name.c_str()); + } else { + fprintf(stderr, "ok: %s\n", name.c_str()); + } +} + +using engine::runtime::LoaderAdvertisement; +using engine::runtime::RunMode; +using engine::runtime::VoiceTaskKind; + +// Built once. make_default_registry constructs every loader, which is the +// expensive part, and none of the assertions mutates it. +const std::vector &catalog() { + static const std::vector kCatalog = + engine::runtime::make_default_registry().advertise_loaders(); + return kCatalog; +} + +bool advertises(const LoaderAdvertisement &loader, VoiceTaskKind task, + RunMode mode) { + for (const auto &capability : loader.capabilities.supported_tasks) { + if (capability.task != task) { + continue; + } + return std::find(capability.modes.begin(), capability.modes.end(), mode) != + capability.modes.end(); + } + return false; +} + +bool advertises_any_mode(const LoaderAdvertisement &loader, VoiceTaskKind task) { + return advertises(loader, task, RunMode::Offline) || + advertises(loader, task, RunMode::Streaming); +} + +std::string describe(const LoaderAdvertisement &loader) { + std::string out; + for (const auto &capability : loader.capabilities.supported_tasks) { + for (const RunMode mode : capability.modes) { + if (!out.empty()) { + out += ", "; + } + out += engine::runtime::to_string(capability.task); + out += "/"; + out += engine::runtime::to_string(mode); + } + } + return out.empty() ? "nothing" : out; +} + +// THE POSITIVE CONTROL. Every other test here asserts an absence, and an +// absence is what a broken query returns for everything. If make_default_registry +// ever returns an empty registry, or advertise_loaders stops populating modes, +// this is the test that fails instead of the suite going quietly green while +// asserting nothing. +void test_catalog_is_populated() { + check(catalog().size() >= 20, + "the default registry advertises a plausible number of families (" + + std::to_string(catalog().size()) + ")"); + + bool found_streaming_asr = false; + for (const auto &loader : catalog()) { + if (advertises(loader, VoiceTaskKind::Asr, RunMode::Streaming)) { + found_streaming_asr = true; + break; + } + } + check(found_streaming_asr, + "the query finds a capability that IS advertised (asr/streaming)"); +} + +// The AudioEncode and AudioDecode premise. Not "no family does codec" but the +// stronger "the task kind does not exist", which is what makes those two RPCs +// unroutable rather than merely unserved. +// +// Note what is NOT asserted: model_spec/schema.cpp's task whitelist DOES accept +// the string "codec", so a spec declaring it validates and then fails here. That +// is a hole in upstream's own validation, and it is deliberately kept off the +// wire; the refusal rests on this parser, which is the thing routing would have +// to go through. +void test_no_codec_task_kind() { + bool threw = false; + try { + (void)engine::runtime::parse_voice_task_kind("codec"); + } catch (const std::exception &) { + threw = true; + } + check(threw, "no codec task kind: parse_voice_task_kind(\"codec\") throws"); + + // The control for the line above: a real name must NOT throw, or the test + // would pass against a parser that rejected everything. + bool tts_threw = false; + try { + (void)engine::runtime::parse_voice_task_kind("tts"); + } catch (const std::exception &) { + tts_threw = true; + } + check(!tts_threw, "parse_voice_task_kind accepts a real name (\"tts\")"); +} + +// The VoiceEmbed premise. SpeakerRecognition IS in the enum, and TitaNet and +// ECAPA-TDNN exist as internal conditioning encoders; what is missing is any +// registered family advertising the task, which is what the message says. +void test_no_family_advertises_speaker_recognition() { + std::string offenders; + for (const auto &loader : catalog()) { + if (advertises_any_mode(loader, VoiceTaskKind::SpeakerRecognition)) { + if (!offenders.empty()) { + offenders += ", "; + } + offenders += loader.family; + } + } + check(offenders.empty(), + "no family advertises spk" + + (offenders.empty() ? std::string() : " (found: " + offenders + ")")); +} + +// The AudioTransformStream premise, and the one that has to be scoped exactly. +// The broad claim "upstream streams tts and asr only" is FALSE: silero_vad +// advertises vad/streaming. The claim that holds is that none of the four tasks +// AudioTransform routes to is advertised streaming by anybody. +// +// NEGATIVE CONTROL: add VoiceTaskKind::Tts to kTransformTasks and this test must +// fail, naming the tts families. A run where that edit stays green means the +// query is broken and every absence above is worthless. +void test_no_streaming_for_the_transform_tasks() { + const VoiceTaskKind kTransformTasks[] = { + VoiceTaskKind::SourceSeparation, + VoiceTaskKind::VoiceConversion, + VoiceTaskKind::Svc, + VoiceTaskKind::SpeechToSpeech, + }; + + std::string offenders; + for (const VoiceTaskKind task : kTransformTasks) { + for (const auto &loader : catalog()) { + if (advertises(loader, task, RunMode::Streaming)) { + if (!offenders.empty()) { + offenders += ", "; + } + offenders += loader.family; + offenders += "/"; + offenders += engine::runtime::to_string(task); + } + } + } + check(offenders.empty(), + "no streaming for sep/vc/svc/s2s" + + (offenders.empty() ? std::string() : " (found: " + offenders + ")")); +} + +// The clause the AudioEncode message names by hand: miocodec carries a Codec tag +// in upstream's README, and its loader advertises only vc and s2s. Asserted +// exactly, not as a subset, so a bump that ADDS a codec capability to miocodec +// fails here rather than leaving the message stale. +void test_miocodec_advertises_exactly_vc_and_s2s() { + const LoaderAdvertisement *miocodec = nullptr; + for (const auto &loader : catalog()) { + if (loader.family == "miocodec") { + miocodec = &loader; + break; + } + } + if (miocodec == nullptr) { + check(false, "miocodec is a registered family"); + return; + } + check(describe(*miocodec) == "vc/offline, s2s/offline", + "miocodec advertises exactly vc/offline, s2s/offline (got: " + + describe(*miocodec) + ")"); +} + +// The "declared only by" clause in the AudioToAudioStream message. Exact set, +// for the same reason as above: a third s2s family would make the message stale +// without making it obviously wrong. +void test_speech_to_speech_is_exactly_miocodec_and_vevo2() { + std::set families; + for (const auto &loader : catalog()) { + if (advertises_any_mode(loader, VoiceTaskKind::SpeechToSpeech)) { + families.insert(loader.family); + } + } + const std::set expected = {"miocodec", "vevo2"}; + std::string found; + for (const auto &family : families) { + if (!found.empty()) { + found += ", "; + } + found += family; + } + check(families == expected, + "s2s is advertised by exactly miocodec and vevo2 (got: " + + (found.empty() ? "nothing" : found) + ")"); +} + +} // namespace + +int main() { + test_catalog_is_populated(); + test_no_codec_task_kind(); + test_no_family_advertises_speaker_recognition(); + test_no_streaming_for_the_transform_tasks(); + test_miocodec_advertises_exactly_vc_and_s2s(); + test_speech_to_speech_is_exactly_miocodec_and_vevo2(); + if (failures) { + fprintf(stderr, + "%d upstream absence check(s) failed. A refusal message in " + "capability_routing.cpp now states something that is not true " + "of the pinned audio.cpp; fix the message, not this test.\n", + failures); + return 1; + } + fprintf(stderr, "all upstream absence checks passed\n"); + return 0; +} diff --git a/backend/cpp/audio-cpp/wav_header.cpp b/backend/cpp/audio-cpp/wav_header.cpp new file mode 100644 index 000000000..9fb4fb4f9 --- /dev/null +++ b/backend/cpp/audio-cpp/wav_header.cpp @@ -0,0 +1,63 @@ +#include "wav_header.h" + +#include +#include + +namespace audiocpp_backend { +namespace { + +void append_u32(std::string &out, std::uint32_t value) { + out.push_back(static_cast(value & 0xFF)); + out.push_back(static_cast((value >> 8) & 0xFF)); + out.push_back(static_cast((value >> 16) & 0xFF)); + out.push_back(static_cast((value >> 24) & 0xFF)); +} + +void append_u16(std::string &out, std::uint16_t value) { + out.push_back(static_cast(value & 0xFF)); + out.push_back(static_cast((value >> 8) & 0xFF)); +} + +// The unknown-length sentinel, in both the RIFF and the data chunk size. +constexpr std::uint32_t kStreamingSize = 0xFFFFFFFFu; +constexpr std::uint16_t kBitsPerSample = 16; +constexpr std::uint32_t kPcmFmtChunkSize = 16; +constexpr std::uint16_t kFormatTagPcm = 1; +constexpr int kMaxChannels = 65535; + +} // namespace + +std::string streaming_wav_header(int sample_rate, int channels) { + int clamped_channels = channels > 0 ? channels : 1; + if (clamped_channels > kMaxChannels) { + clamped_channels = kMaxChannels; + } + const auto channel_count = static_cast(clamped_channels); + const auto rate = static_cast(sample_rate > 0 ? sample_rate : 0); + const auto block_align = + static_cast(channel_count * (kBitsPerSample / 8)); + // uint32 arithmetic on purpose: 384 kHz by 8 channels is 6.1 MB/s, which + // does not fit the uint16 block align it is derived from. + const std::uint32_t byte_rate = rate * static_cast(block_align); + static_assert(std::numeric_limits::max() >= 0xFFFFFFFFu, + "the streaming sentinel must be representable"); + + std::string header; + header.reserve(44); + header += "RIFF"; + append_u32(header, kStreamingSize); // unknown total length + header += "WAVE"; + header += "fmt "; + append_u32(header, kPcmFmtChunkSize); + append_u16(header, kFormatTagPcm); + append_u16(header, channel_count); + append_u32(header, rate); + append_u32(header, byte_rate); + append_u16(header, block_align); + append_u16(header, kBitsPerSample); + header += "data"; + append_u32(header, kStreamingSize); // unknown payload length + return header; +} + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/wav_header.h b/backend/cpp/audio-cpp/wav_header.h new file mode 100644 index 000000000..42dff40b6 --- /dev/null +++ b/backend/cpp/audio-cpp/wav_header.h @@ -0,0 +1,39 @@ +#pragma once + +// Builds the 44 byte canonical WAV header that precedes a streamed PCM body. +// Standard library only. +// +// TTSStream chunks travel in Reply.audio and the FIRST chunk must be this +// header, or an HTTP client has no format to decode the PCM with. Because the +// total length is unknown while the model is still generating, both size fields +// carry 0xFFFFFFFF; that is the convention backend/go/vibevoice-cpp established +// (govibevoicecpp.go, TTSStream) and the one core/backend/tts.go writes when it +// synthesises a header itself. +// +// WHY THIS BACKEND SENDS THE HEADER RATHER THAN LETTING GO DO IT. +// core/backend/tts.go's ModelTTSStream will emit a header of its own, but only +// when the FIRST Reply carries a non-empty `message` field holding a JSON blob +// with a sample_rate. This backend sends audio and never sets `message`, so +// that branch never fires and there is exactly one header on the wire: this +// one. Do not start setting Reply.message on this RPC without deleting the +// header below, or every stream gains a second header 44 bytes into the PCM. +// +// The header this produces is byte-identical to the one pkg/audio.WAVHeader +// serialises, which is what core/http/endpoints/openai/realtime_model.go +// assumes when it reads the sample rate out of byte offset 24 of the first +// callback. + +#include + +namespace audiocpp_backend { + +// 16-bit PCM, little endian, interleaved. +// +// `channels` is clamped to at least 1 and at most 65535, so a garbage channel +// count can never write a zero block align, which is what a reader divides the +// data size by. `sample_rate` is clamped to at least 0 rather than wrapped: +// a zero rate is visibly wrong to whoever reads the header, whereas the +// 4294967295 an unsigned conversion of -1 would write looks like a real field. +std::string streaming_wav_header(int sample_rate, int channels); + +} // namespace audiocpp_backend diff --git a/backend/cpp/audio-cpp/wav_header_test.cpp b/backend/cpp/audio-cpp/wav_header_test.cpp new file mode 100644 index 000000000..f86ddfe00 --- /dev/null +++ b/backend/cpp/audio-cpp/wav_header_test.cpp @@ -0,0 +1,163 @@ +// Unit tests for wav_header. Standard library only. The harness compiles this +// as a single translation unit, so the implementation is included directly. + +#include "wav_header.cpp" + +#include +#include +#include + +static int failures = 0; + +static void check(bool ok, const std::string &name) { + if (!ok) { + failures++; + fprintf(stderr, "FAIL: %s\n", name.c_str()); + } else { + fprintf(stderr, "ok: %s\n", name.c_str()); + } +} + +using audiocpp_backend::streaming_wav_header; + +static std::uint32_t read_u32(const std::string &data, size_t offset) { + return static_cast(static_cast(data[offset])) | + (static_cast(static_cast(data[offset + 1])) << 8) | + (static_cast(static_cast(data[offset + 2])) << 16) | + (static_cast(static_cast(data[offset + 3])) << 24); +} + +static std::uint16_t read_u16(const std::string &data, size_t offset) { + return static_cast( + static_cast(static_cast(data[offset])) | + (static_cast(static_cast(data[offset + 1])) << 8)); +} + +static std::string to_hex(const std::string &data) { + static const char *digits = "0123456789abcdef"; + std::string out; + out.reserve(data.size() * 2); + for (const char byte : data) { + const auto value = static_cast(byte); + out.push_back(digits[value >> 4]); + out.push_back(digits[value & 0x0F]); + } + return out; +} + +static void test_header_layout() { + const std::string header = streaming_wav_header(24000, 1); + + check(header.size() == 44, "canonical 44 byte header"); + check(header.compare(0, 4, "RIFF") == 0, "RIFF magic"); + check(header.compare(8, 4, "WAVE") == 0, "WAVE magic"); + check(header.compare(12, 4, "fmt ") == 0, "fmt chunk id"); + check(header.compare(36, 4, "data") == 0, "data chunk id"); + + check(read_u32(header, 16) == 16, "PCM fmt chunk is 16 bytes"); + check(read_u16(header, 20) == 1, "format tag 1 is PCM"); + check(read_u16(header, 22) == 1, "mono channel count"); + check(read_u32(header, 24) == 24000, "sample rate"); + // byte rate = rate * channels * bytes per sample + check(read_u32(header, 28) == 24000 * 1 * 2, "byte rate"); + check(read_u16(header, 32) == 2, "block align for mono 16 bit"); + check(read_u16(header, 34) == 16, "16 bits per sample"); +} + +// The field-wise checks above can all pass while the fields sit in the wrong +// ORDER, since several of them hold the same value. This pins the whole 44 byte +// string against a literal transcribed from the layout in pkg/audio/audio.go, +// which is the struct binary.Write serializes for every Go LocalAI backend. +// Independent of the implementation: it was written out by hand rather than +// captured from a run. +static void test_exact_bytes_match_the_go_layout() { + const std::string expected = + "52494646" // "RIFF" + "ffffffff" // chunk size: streaming sentinel + "57415645" // "WAVE" + "666d7420" // "fmt " + "10000000" // subchunk1 size 16 + "0100" // audio format 1 (PCM) + "0100" // channels 1 + "c05d0000" // sample rate 24000 + "80bb0000" // byte rate 48000 + "0200" // block align 2 + "1000" // bits per sample 16 + "64617461" // "data" + "ffffffff"; // subchunk2 size: streaming sentinel + check(to_hex(streaming_wav_header(24000, 1)) == expected, + "byte for byte match with the canonical mono 24 kHz header"); +} + +// The whole point: a streaming header cannot know the final length, so both +// size fields are the sentinel. A client that sees a real size stops early. +static void test_streaming_sentinels() { + const std::string header = streaming_wav_header(16000, 1); + check(read_u32(header, 4) == 0xFFFFFFFFu, "RIFF chunk size is the sentinel"); + check(read_u32(header, 40) == 0xFFFFFFFFu, "data chunk size is the sentinel"); +} + +static void test_stereo() { + const std::string header = streaming_wav_header(44100, 2); + check(read_u16(header, 22) == 2, "stereo channel count"); + check(read_u32(header, 28) == 44100 * 2 * 2, "stereo byte rate"); + check(read_u16(header, 32) == 4, "block align for stereo 16 bit"); + check(header.size() == 44, "stereo header is still 44 bytes"); +} + +static void test_degenerate_inputs() { + // A zero or negative channel count must not produce a header that divides + // by zero downstream; clamp to mono. + const std::string zero_channels = streaming_wav_header(16000, 0); + check(read_u16(zero_channels, 22) == 1, "zero channels clamps to mono"); + check(read_u16(zero_channels, 32) == 2, "zero channels still block aligns as mono"); + check(read_u32(zero_channels, 28) == 32000, "zero channels byte rate is the mono one"); + + const std::string negative_channels = streaming_wav_header(16000, -3); + check(read_u16(negative_channels, 22) == 1, "negative channels clamps to mono"); + + // A channel count past the field's range must saturate rather than wrap: + // 65536 truncated to uint16 is 0, and a zero channel count writes a zero + // block align, which is what a reader divides the data size by. + const std::string too_many = streaming_wav_header(16000, 65536); + check(read_u16(too_many, 22) == 65535, "an out of range channel count saturates"); + check(read_u16(too_many, 32) != 0, "an out of range channel count never writes a zero block align"); + + // A non-positive rate is written as zero rather than wrapping through the + // unsigned conversion: 0 is visibly wrong to whoever reads the header, + // 4294967295 looks like a plausible field nobody checks. + const std::string zero_rate = streaming_wav_header(0, 1); + check(read_u32(zero_rate, 24) == 0, "zero sample rate stays zero"); + check(read_u32(zero_rate, 28) == 0, "zero sample rate yields a zero byte rate"); + const std::string negative_rate = streaming_wav_header(-48000, 1); + check(read_u32(negative_rate, 24) == 0, "negative sample rate is clamped to zero"); + check(read_u32(negative_rate, 28) == 0, "negative sample rate yields a zero byte rate"); + + // The sentinels are unconditional. A degenerate rate must not turn the + // stream into one a client thinks it can measure. + check(read_u32(zero_rate, 4) == 0xFFFFFFFFu, "degenerate input keeps the RIFF sentinel"); + check(read_u32(zero_rate, 40) == 0xFFFFFFFFu, "degenerate input keeps the data sentinel"); +} + +// Large but legal: 384 kHz 8 channel would overflow a 16 bit byte rate and +// must not overflow the 32 bit one either. +static void test_large_but_legal() { + const std::string header = streaming_wav_header(384000, 8); + check(read_u32(header, 28) == 384000u * 8u * 2u, "high rate multichannel byte rate"); + check(read_u16(header, 32) == 16, "high channel count block align"); +} + +int main() { + test_header_layout(); + test_exact_bytes_match_the_go_layout(); + test_streaming_sentinels(); + test_stereo(); + test_degenerate_inputs(); + test_large_but_legal(); + if (failures) { + fprintf(stderr, "%d check(s) failed\n", failures); + return 1; + } + fprintf(stderr, "all wav_header checks passed\n"); + return 0; +} diff --git a/backend/cpp/bonsai/Makefile b/backend/cpp/bonsai/Makefile index 82ba9660f..5cbbbfffe 100644 --- a/backend/cpp/bonsai/Makefile +++ b/backend/cpp/bonsai/Makefile @@ -1,7 +1,7 @@ # Pinned to the HEAD of the `prism` branch on https://github.com/PrismML-Eng/llama.cpp. # Auto-bumped nightly by .github/workflows/bump_deps.yaml. -BONSAI_VERSION?=7529fdaaf99ffdc5ca71ace9c7409a56b27ad92f +BONSAI_VERSION?=9ca265a57f85f2117942490f421f64a226dd9847 LLAMA_REPO?=https://github.com/PrismML-Eng/llama.cpp CMAKE_ARGS?= @@ -41,6 +41,7 @@ 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 @@ -77,6 +78,7 @@ 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 diff --git a/backend/cpp/bonsai/package.sh b/backend/cpp/bonsai/package.sh index 1fa40abb3..66a9cb73e 100755 --- a/backend/cpp/bonsai/package.sh +++ b/backend/cpp/bonsai/package.sh @@ -24,34 +24,7 @@ if [ -d "$CURDIR/ggml-shared-libs" ]; then fi # 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 -else - echo "Error: Could not detect architecture" - exit 1 -fi +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries based on BUILD_TYPE GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh" diff --git a/backend/cpp/bonsai/run.sh b/backend/cpp/bonsai/run.sh index 4c49f9e40..e59c5954e 100755 --- a/backend/cpp/bonsai/run.sh +++ b/backend/cpp/bonsai/run.sh @@ -40,6 +40,27 @@ else if [ -d "$CURDIR/lib/hipblaslt/library" ]; then export HIPBLASLT_TENSILE_LIBPATH="$CURDIR"/lib/hipblaslt/library fi + # Backends built for Intel GPUs carry a copy of the Intel graphics driver, + # and libze_loader is only there in those builds. Level Zero looks for a + # driver on its own, so point it at the copy that came with this backend: it + # was built against the same C library, while the machine's own driver may + # not have been, and loading that one can crash on start. + # + # Anything the user set is left alone, so a machine with a graphics card + # newer than the driver carried here can still be told to use its own. + # Nothing is said about OpenCL: no OpenCL driver is carried, so anything we + # set there would leave OpenCL worse off than the machine's own setup. + if [ -e "$CURDIR/lib/libze_loader.so.1" ]; then + if [ -e "$CURDIR/lib/libze_intel_gpu.so.1" ] && [ -z "${ZE_ENABLE_ALT_DRIVERS:-}" ]; then + export ZE_ENABLE_ALT_DRIVERS="$CURDIR"/lib/libze_intel_gpu.so.1 + fi + # Ask the driver how much graphics memory is free. Without this, the + # backend reads zero on an integrated graphics chip, because such a chip + # shares the system memory instead of having its own. + if [ -z "${ZES_ENABLE_SYSMAN:-}" ]; then + export ZES_ENABLE_SYSMAN=1 + fi + fi fi # If there is a lib/ld.so, use it diff --git a/backend/cpp/ds4/CMakeLists.txt b/backend/cpp/ds4/CMakeLists.txt index 10d999729..0535a8a44 100644 --- a/backend/cpp/ds4/CMakeLists.txt +++ b/backend/cpp/ds4/CMakeLists.txt @@ -69,7 +69,15 @@ target_include_directories(hw_grpc_proto PUBLIC ${CMAKE_CURRENT_BINARY_DIR}) set(DS4_OBJS "${DS4_DIR}/ds4.o") if(DS4_GPU STREQUAL "cuda") - list(APPEND DS4_OBJS "${DS4_DIR}/ds4_cuda.o") + list(APPEND DS4_OBJS + "${DS4_DIR}/ds4_cuda.o" + "${DS4_DIR}/cuda/mmq/ds4_ggml_stubs.o" + "${DS4_DIR}/cuda/mmq/ds4_mmq.o" + "${DS4_DIR}/cuda/mmq/ds4_mmq_d2r.o" + "${DS4_DIR}/cuda/mmq/quantize.o" + "${DS4_DIR}/cuda/mmq/mmid.o" + "${DS4_DIR}/cuda/mmq/mmvq.o" + "${DS4_DIR}/cuda/mmq/ds4_repack.o") elseif(DS4_GPU STREQUAL "metal") list(APPEND DS4_OBJS "${DS4_DIR}/ds4_metal.o") elseif(DS4_GPU STREQUAL "cpu") diff --git a/backend/cpp/ds4/Makefile b/backend/cpp/ds4/Makefile index 19929789d..c5480badd 100644 --- a/backend/cpp/ds4/Makefile +++ b/backend/cpp/ds4/Makefile @@ -1,10 +1,10 @@ # ds4 backend Makefile. # -# Upstream pin lives below as DS4_VERSION?=0a7ad776b9068348e6cb09df8cafa9cadd285298 +# Upstream pin lives below as DS4_VERSION?=b0309611041655f4e45671cfd9c9886aff161406 # (.github/bump_deps.sh) can find and update it - matches the # llama-cpp / ik-llama-cpp / turboquant convention. -DS4_VERSION?=0a7ad776b9068348e6cb09df8cafa9cadd285298 +DS4_VERSION?=b0309611041655f4e45671cfd9c9886aff161406 DS4_REPO?=https://github.com/antirez/ds4 CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST)))) @@ -23,7 +23,9 @@ CMAKE_ARGS ?= -DCMAKE_BUILD_TYPE=Release # are shared by every GPU mode, so append them unconditionally below. ifeq ($(BUILD_TYPE),cublas) CMAKE_ARGS += -DDS4_GPU=cuda - DS4_OBJ_TARGET := ds4.o ds4_cuda.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o + DS4_OBJ_TARGET := ds4.o ds4_cuda.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o \ + cuda/mmq/ds4_ggml_stubs.o cuda/mmq/ds4_mmq.o cuda/mmq/ds4_mmq_d2r.o \ + cuda/mmq/quantize.o cuda/mmq/mmid.o cuda/mmq/mmvq.o cuda/mmq/ds4_repack.o else ifeq ($(UNAME_S),Darwin) CMAKE_ARGS += -DDS4_GPU=metal DS4_OBJ_TARGET := ds4.o ds4_metal.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o @@ -55,7 +57,7 @@ ds4: # the right per-platform compile flags (Objective-C/Metal on Darwin, nvcc on Linux+CUDA). ds4/ds4.o: ds4 ifeq ($(BUILD_TYPE),cublas) - +$(MAKE) -C ds4 ds4.o ds4_cuda.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o + +$(MAKE) -C ds4 $(DS4_OBJ_TARGET) else ifeq ($(UNAME_S),Darwin) +$(MAKE) -C ds4 ds4.o ds4_metal.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o else diff --git a/backend/cpp/ds4/package.sh b/backend/cpp/ds4/package.sh index de8daed25..2769aa53c 100755 --- a/backend/cpp/ds4/package.sh +++ b/backend/cpp/ds4/package.sh @@ -17,13 +17,7 @@ if [ "$UNAME_S" = "Darwin" ]; then exit 0 fi -if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then - cp -arfLv /lib64/ld-linux-x86-64.so.2 "$PACKAGE_DIR/lib/ld.so" -elif [ -f "/lib/ld-linux-aarch64.so.1" ]; then - cp -arfLv /lib/ld-linux-aarch64.so.1 "$PACKAGE_DIR/lib/ld.so" -else - echo "package.sh: unknown architecture" >&2; exit 1 -fi +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Bundle the complete dependency closure for both executables. In particular, # grpc-server links the distro gRPC/protobuf/absl stack; copying only the core diff --git a/backend/cpp/ik-llama-cpp/Makefile b/backend/cpp/ik-llama-cpp/Makefile index ce8c6c754..da3a47d78 100644 --- a/backend/cpp/ik-llama-cpp/Makefile +++ b/backend/cpp/ik-llama-cpp/Makefile @@ -1,5 +1,5 @@ -IK_LLAMA_VERSION?=0a4e10c7fb65d2dd5a4afb78339c7d373a8cdfaa +IK_LLAMA_VERSION?=cf1aa57e1a0fabfd015831718fc99d1aec01ada5 LLAMA_REPO?=https://github.com/ikawrakow/ik_llama.cpp CMAKE_ARGS?= diff --git a/backend/cpp/ik-llama-cpp/package.sh b/backend/cpp/ik-llama-cpp/package.sh index 56d430563..c18c271e3 100644 --- a/backend/cpp/ik-llama-cpp/package.sh +++ b/backend/cpp/ik-llama-cpp/package.sh @@ -15,34 +15,7 @@ cp -avrf $CURDIR/ik-llama-cpp-* $CURDIR/package/ cp -rfv $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 -else - echo "Error: Could not detect architecture" - exit 1 -fi +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries based on BUILD_TYPE # The GPU library packaging script will detect BUILD_TYPE and copy appropriate GPU libraries diff --git a/backend/cpp/llama-cpp/CMakeLists.txt b/backend/cpp/llama-cpp/CMakeLists.txt index 47852d400..f861713ee 100644 --- a/backend/cpp/llama-cpp/CMakeLists.txt +++ b/backend/cpp/llama-cpp/CMakeLists.txt @@ -110,4 +110,9 @@ if(LLAMA_GRPC_BUILD_TESTS) target_link_libraries(parent_watch_test PRIVATE Threads::Threads) target_compile_features(parent_watch_test PRIVATE cxx_std_17) add_test(NAME parent_watch_test COMMAND parent_watch_test) + + add_executable(passthrough_options_test passthrough_options_test.cpp passthrough_options.h) + target_include_directories(passthrough_options_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) + target_compile_features(passthrough_options_test PRIVATE cxx_std_17) + add_test(NAME passthrough_options_test COMMAND passthrough_options_test) endif() diff --git a/backend/cpp/llama-cpp/Makefile b/backend/cpp/llama-cpp/Makefile index 73de6f757..c2d9096b5 100644 --- a/backend/cpp/llama-cpp/Makefile +++ b/backend/cpp/llama-cpp/Makefile @@ -1,5 +1,5 @@ -LLAMA_VERSION?=0d47ea7427463093e69128bf2c2f9cd06b3ee5b3 +LLAMA_VERSION?=221f0f6356efe2260023208365705ec5d5a7c8f5 LLAMA_REPO?=https://github.com/ggerganov/llama.cpp CMAKE_ARGS?= diff --git a/backend/cpp/llama-cpp/disable-score-task.sh b/backend/cpp/llama-cpp/disable-score-task.sh new file mode 100644 index 000000000..164d4a57f --- /dev/null +++ b/backend/cpp/llama-cpp/disable-score-task.sh @@ -0,0 +1,43 @@ +#!/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 " >&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" diff --git a/backend/cpp/llama-cpp/grpc-server.cpp b/backend/cpp/llama-cpp/grpc-server.cpp index 02435310e..67d5ef11f 100644 --- a/backend/cpp/llama-cpp/grpc-server.cpp +++ b/backend/cpp/llama-cpp/grpc-server.cpp @@ -54,6 +54,7 @@ #include "chat-auto-parser.h" #include "llama_compat.h" // fork-skew switches, generated by prepare.sh #include "message_content.h" +#include "passthrough_options.h" #include #include #include @@ -152,40 +153,6 @@ 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 slot_loop_inflight{0}; -static std::atomic 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& self; - conflict_guard(const char* rpc, std::atomic& self_, std::atomic& 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 shutdown_handler; static std::atomic_flag is_terminating = ATOMIC_FLAG_INIT; @@ -613,6 +580,8 @@ static void params_parse(server_context& /*ctx_server*/, const backend::ModelOpt // Raw upstream llama-server flags collected from any option entry that // starts with '-'. Applied once after the loop via common_params_parse. std::vector extra_argv; + bool passthrough_main_gpu_layers = false; + bool passthrough_draft_gpu_layers = false; // O_DIRECT intent from the `direct_io` option. Upstream folded // use_mmap/use_mlock/use_direct_io into a single common_params::load_mode @@ -732,6 +701,22 @@ 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 { @@ -1204,6 +1189,17 @@ static void params_parse(server_context& /*ctx_server*/, const backend::ModelOpt flag.c_str()); } else { extra_argv.push_back(flag); + passthrough_main_gpu_layers = + passthrough_main_gpu_layers || + flag == "-ngl" || + flag == "--gpu-layers" || + flag == "--n-gpu-layers"; + passthrough_draft_gpu_layers = + passthrough_draft_gpu_layers || + flag == "--spec-draft-ngl" || + flag == "-ngld" || + flag == "--gpu-layers-draft" || + flag == "--n-gpu-layers-draft"; // Preserve the whole value after the first ':' so embedded // colons (e.g. host:port) survive strtok's truncation of optval. auto colon = opt.find(':'); @@ -1367,6 +1363,14 @@ static void params_parse(server_context& /*ctx_server*/, const backend::ModelOpt // (n_parallel -> -1, use_color). Snapshot n_parallel so an unrelated // passthrough flag can't silently clobber LocalAI's resolved value. const int saved_n_parallel = params.n_parallel; + // Newer upstream parsers assert that these fields still contain their + // negative initialization sentinels. LocalAI resolves them from the + // model request before applying passthrough options, so stage the + // sentinels and restore the values unless a raw flag overrides them. + const auto saved_gpu_layers = + llama_grpc::prepare_passthrough_gpu_layers( + params.n_gpu_layers, + params.speculative.draft.n_gpu_layers); std::vector argv; std::string prog = "llama-server"; @@ -1390,8 +1394,25 @@ static void params_parse(server_context& /*ctx_server*/, const backend::ModelOpt if (params.n_parallel == -1) { params.n_parallel = saved_n_parallel; } + llama_grpc::restore_passthrough_gpu_layers( + params.n_gpu_layers, + params.speculative.draft.n_gpu_layers, + saved_gpu_layers, + passthrough_main_gpu_layers, + passthrough_draft_gpu_layers); } +#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. @@ -1478,6 +1499,16 @@ 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); @@ -1680,7 +1711,6 @@ 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()); @@ -2249,7 +2279,6 @@ 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; @@ -2783,7 +2812,6 @@ 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,7 +2921,6 @@ 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(); @@ -2970,37 +2997,16 @@ public: // Score returns the model's joint log-probability of each candidate // continuation given a shared prompt. // - // 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. + // 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. grpc::Status Score(ServerContext* context, const backend::ScoreRequest* request, backend::ScoreResponse* response) override { auto auth = checkAuth(context); if (!auth.ok()) return auth; @@ -3009,40 +3015,21 @@ 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 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 @@ -3051,6 +3038,15 @@ public: std::vector 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> cand_tokens(request->candidates_size()); + std::vector cand_divergence(request->candidates_size(), 0); + + // candidates that actually have tokens to score + std::vector included; + for (int ci = 0; ci < request->candidates_size(); ci++) { const std::string & candidate_text = request->candidates(ci); @@ -3067,9 +3063,135 @@ public: break; } } + divergence = std::min(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(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 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 shared_logprobs; + std::vector> 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(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); + cs->set_num_tokens(cand_len > 0 ? cand_len : 0); if (cand_len <= 0) { cs->set_log_prob(0.0); if (request->length_normalize()) { @@ -3077,101 +3199,57 @@ public: } continue; } - 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"); + + // 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 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 ((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) + ")"); + 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]); } - // 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); + if ((int32_t) lp.size() != cand_len) { return grpc::Status(grpc::StatusCode::INTERNAL, - "llama_decode failed during Score: " + std::to_string(decode_err)); + "Score: result for candidate " + std::to_string(ci) + " is missing token logprobs"); } - // Sum log-probabilities of the actual candidate tokens. double total_log_prob = 0.0; for (int32_t k = 0; k < cand_len; k++) { - // 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); + const float token_log_prob = lp[k]; + if (std::isnan(token_log_prob)) { return grpc::Status(grpc::StatusCode::INTERNAL, - "llama_get_logits_ith returned null at position " + std::to_string(logit_pos)); + "Score: incomplete result for candidate " + std::to_string(ci) + + " at token " + std::to_string(k)); } - 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; + total_log_prob += (double) token_log_prob; if (request->include_token_logprobs()) { backend::TokenLogProb * tlp = cs->add_tokens(); - std::string piece = common_token_to_piece(lctx, target_token); - tlp->set_token(piece); + tlp->set_token(common_token_to_piece(vocab, cand_tokens[ci][divergence + k])); tlp->set_log_prob(token_log_prob); } } cs->set_log_prob(total_log_prob); - if (request->length_normalize() && cand_len > 0) { + if (request->length_normalize()) { 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 { @@ -3182,7 +3260,6 @@ 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; @@ -3202,9 +3279,23 @@ public: return grpc::Status::OK; } + grpc::Status Detokenize(ServerContext* context, const backend::DetokenizeRequest* request, backend::DetokenizeResponse* response) override { + auto auth = checkAuth(context); + if (!auth.ok()) return auth; + if (params_base.model.path.empty()) { + return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded"); + } + + std::string content; + for (const auto token : request->tokens()) { + content.append(common_token_to_piece(ctx_server.get_llama_context(), token)); + } + response->set_content(content); + return grpc::Status::OK; + } + 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(); diff --git a/backend/cpp/llama-cpp/package.sh b/backend/cpp/llama-cpp/package.sh index 5d2b18c5b..a36a279ae 100755 --- a/backend/cpp/llama-cpp/package.sh +++ b/backend/cpp/llama-cpp/package.sh @@ -31,34 +31,7 @@ if [ -d "$CURDIR/ggml-shared-libs" ]; then fi # 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 -else - echo "Error: Could not detect architecture" - exit 1 -fi +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries based on BUILD_TYPE # The GPU library packaging script will detect BUILD_TYPE and copy appropriate GPU libraries diff --git a/backend/cpp/llama-cpp/passthrough_options.h b/backend/cpp/llama-cpp/passthrough_options.h new file mode 100644 index 000000000..b2c897d8f --- /dev/null +++ b/backend/cpp/llama-cpp/passthrough_options.h @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: MIT + +#pragma once + +namespace llama_grpc { + +struct passthrough_gpu_layers_state { + int main; + int draft; +}; + +inline passthrough_gpu_layers_state prepare_passthrough_gpu_layers( + int & main_gpu_layers, + int & draft_gpu_layers) { + const passthrough_gpu_layers_state saved{ + main_gpu_layers, + draft_gpu_layers, + }; + main_gpu_layers = -1; + draft_gpu_layers = -1; + return saved; +} + +inline void restore_passthrough_gpu_layers( + int & main_gpu_layers, + int & draft_gpu_layers, + passthrough_gpu_layers_state saved, + bool main_overridden = false, + bool draft_overridden = false) { + if (!main_overridden && main_gpu_layers == -1) { + main_gpu_layers = saved.main; + } + if (!draft_overridden && draft_gpu_layers == -1) { + draft_gpu_layers = saved.draft; + } +} + +} // namespace llama_grpc diff --git a/backend/cpp/llama-cpp/passthrough_options_test.cpp b/backend/cpp/llama-cpp/passthrough_options_test.cpp new file mode 100644 index 000000000..c78825e5c --- /dev/null +++ b/backend/cpp/llama-cpp/passthrough_options_test.cpp @@ -0,0 +1,68 @@ +// SPDX-License-Identifier: MIT + +#include + +#include "passthrough_options.h" + +static int failures = 0; + +static void check(bool ok, const char * name) { + if (!ok) { + ++failures; + std::fprintf(stderr, "FAIL: %s\n", name); + } +} + +static void test_stages_resolved_gpu_layers_for_upstream_parser() { + int main_gpu_layers = 99; + int draft_gpu_layers = 12; + + const auto saved = llama_grpc::prepare_passthrough_gpu_layers( + main_gpu_layers, draft_gpu_layers); + + check(main_gpu_layers == -1, "main GPU layers use parser sentinel"); + check(draft_gpu_layers == -1, "draft GPU layers use parser sentinel"); + + llama_grpc::restore_passthrough_gpu_layers( + main_gpu_layers, draft_gpu_layers, saved); + + check(main_gpu_layers == 99, "main GPU layers restored"); + check(draft_gpu_layers == 12, "draft GPU layers restored"); +} + +static void test_keeps_explicit_passthrough_overrides() { + int main_gpu_layers = 99; + int draft_gpu_layers = 12; + + const auto saved = llama_grpc::prepare_passthrough_gpu_layers( + main_gpu_layers, draft_gpu_layers); + + main_gpu_layers = 4; + draft_gpu_layers = 2; + llama_grpc::restore_passthrough_gpu_layers( + main_gpu_layers, draft_gpu_layers, saved); + + check(main_gpu_layers == 4, "main passthrough override retained"); + check(draft_gpu_layers == 2, "draft passthrough override retained"); +} + +static void test_keeps_explicit_auto_passthrough_overrides() { + int main_gpu_layers = 99; + int draft_gpu_layers = 12; + + const auto saved = llama_grpc::prepare_passthrough_gpu_layers( + main_gpu_layers, draft_gpu_layers); + + llama_grpc::restore_passthrough_gpu_layers( + main_gpu_layers, draft_gpu_layers, saved, true, true); + + check(main_gpu_layers == -1, "main auto passthrough override retained"); + check(draft_gpu_layers == -1, "draft auto passthrough override retained"); +} + +int main() { + test_stages_resolved_gpu_layers_for_upstream_parser(); + test_keeps_explicit_passthrough_overrides(); + test_keeps_explicit_auto_passthrough_overrides(); + return failures == 0 ? 0 : 1; +} diff --git a/backend/cpp/llama-cpp/patches/0001-add-minimax-m3-chat-parser.patch b/backend/cpp/llama-cpp/patches/0001-add-minimax-m3-chat-parser.patch deleted file mode 100644 index 2655102c7..000000000 --- a/backend/cpp/llama-cpp/patches/0001-add-minimax-m3-chat-parser.patch +++ /dev/null @@ -1,225 +0,0 @@ -# MiniMax-M3 chat-template parser, vendored from upstream llama.cpp PR #24523. -# -# Upstream has since merged the *model* half of #24523 (LLM_ARCH_MINIMAX_M3, -# src/models/minimax-m3.cpp, the gguf-py constants and conversion/minimax.py), so -# only the chat half is carried here: M3's namespace token "]<]minimax[>[" collides -# with the autoparser's markup delimiters, so common/chat.cpp needs a dedicated -# template detection + PEG parser that upstream does not have yet. -# -# Rebased against LLAMA_VERSION 0d47ea7427463093e69128bf2c2f9cd06b3ee5b3, which also -# renamed common_chat_params::thinking_end_tag to thinking_end_tags (a vector). -# LLAMA_VERSION is auto-bumped nightly; if a bump rejects this patch, re-vendor from -# #24523 — or, once the chat half merges upstream, delete this file. -# See https://github.com/mudler/LocalAI/issues/10820 and PR #10837. -diff --git a/common/chat.cpp b/common/chat.cpp -index 7a6e7238c..2dd015a2e 100644 ---- a/common/chat.cpp -+++ b/common/chat.cpp -@@ -2121,6 +2121,191 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha - return data; - } - -+static common_chat_params common_chat_params_init_minimax_m3(const common_chat_template & tmpl, -+ const autoparser::generation_params & inputs) { -+ common_chat_params data; -+ -+ data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); -+ data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); -+ data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; -+ data.supports_thinking = true; -+ data.thinking_start_tag = ""; -+ data.thinking_end_tags = {""}; -+ -+ // M3 prefixes every tool tag with the namespace token "]<]minimax[>["; -+ // params use the parameter name as the tag (...). -+ const std::string NS = "]<]minimax[>["; -+ const std::string THINK_START = ""; -+ const std::string THINK_END = ""; -+ const std::string FC_START = NS + ""; -+ const std::string FC_END = NS + ""; -+ const std::string INVOKE_END = NS + ""; -+ -+ data.preserved_tokens = { -+ NS, -+ "", -+ "", -+ THINK_START, -+ THINK_END, -+ }; -+ -+ auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); -+ auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); -+ auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; -+ auto include_grammar = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE); -+ -+ const std::string GEN_PROMPT = data.generation_prompt; -+ -+ if (inputs.has_continuation()) { -+ const auto & msg = inputs.continue_msg; -+ -+ data.generation_prompt = GEN_PROMPT + THINK_START + msg.reasoning_content; -+ if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { -+ data.generation_prompt += THINK_END + msg.render_content(); -+ } -+ -+ data.prompt += data.generation_prompt; -+ } -+ -+ auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { -+ auto generation_prompt = p.literal(GEN_PROMPT); -+ auto end = p.end(); -+ -+ auto reasoning = p.eps(); -+ // M3 can emit a bare (no opener) after tool results; keep the opener optional. -+ if (extract_reasoning && inputs.enable_thinking) { -+ reasoning = p.optional(p.optional(p.literal(THINK_START)) + p.reasoning(p.until(THINK_END)) + THINK_END); -+ } else if (extract_reasoning) { -+ reasoning = p.optional(p.optional(p.literal(THINK_START)) + p.until(THINK_END) + p.literal(THINK_END)); -+ } -+ -+ if (has_response_format) { -+ auto response_format = p.rule("response-format", -+ p.literal("```json") + p.space() + -+ p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)) + -+ p.space() + p.literal("```")); -+ return generation_prompt + reasoning + response_format + end; -+ } -+ -+ if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { -+ return generation_prompt + reasoning + p.content(p.rest()) + end; -+ } -+ -+ auto tool_choice = p.choice(); -+ foreach_function(inputs.tools, [&](const json & tool) { -+ const auto & function = tool.at("function"); -+ std::string name = function.at("name"); -+ auto params = function.contains("parameters") ? function.at("parameters") : json::object(); -+ const auto & props = params.contains("properties") ? params.at("properties") : json::object(); -+ -+ std::set required; -+ if (params.contains("required")) { -+ params.at("required").get_to(required); -+ } -+ -+ auto schema_info = common_schema_info(); -+ schema_info.resolve_refs(params); -+ -+ std::vector required_parsers; -+ std::vector optional_parsers; -+ for (const auto & [param_name, param_schema] : props.items()) { -+ bool is_required = required.find(param_name) != required.end(); -+ bool is_string = schema_info.resolves_to_string(param_schema); -+ -+ const std::string p_close = NS + ""; -+ -+ auto arg = p.tool_arg( -+ p.tool_arg_open( -+ p.literal(NS + "<") + -+ p.tool_arg_name(p.literal(param_name)) + -+ p.literal(">")) + -+ (is_string -+ ? p.ac(p.tool_arg_string_value(p.until(p_close)) + -+ p.tool_arg_close(p.literal(p_close)), p_close) -+ : p.tool_arg_json_value(p.schema(p.json(), -+ "tool-" + name + "-arg-" + param_name + "-schema", -+ param_schema, false)) + -+ p.tool_arg_close(p.literal(p_close)))); -+ -+ auto named_arg = p.rule("tool-" + name + "-arg-" + param_name, arg); -+ if (is_required) { -+ required_parsers.push_back(named_arg); -+ } else { -+ optional_parsers.push_back(named_arg); -+ } -+ } -+ -+ common_peg_parser args_seq = p.eps(); -+ for (size_t i = 0; i < required_parsers.size(); i++) { -+ if (i > 0) { -+ args_seq = args_seq + p.space(); -+ } -+ args_seq = args_seq + required_parsers[i]; -+ } -+ -+ if (!optional_parsers.empty()) { -+ common_peg_parser any_opt = p.choice(); -+ for (const auto & opt : optional_parsers) { -+ any_opt |= opt; -+ } -+ args_seq = args_seq + p.repeat(p.space() + any_opt, 0, -1); -+ } -+ -+ common_peg_parser invoke_body = args_seq; -+ auto func_parser = p.tool( -+ p.tool_open(p.literal(NS + "")) + -+ p.space() + invoke_body + p.space() + -+ p.tool_close(p.literal(INVOKE_END))); -+ -+ tool_choice |= p.rule("tool-" + name, func_parser); -+ }); -+ -+ auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; -+ -+ common_peg_parser tool_calls = p.eps(); -+ if (inputs.parallel_tool_calls) { -+ tool_calls = p.trigger_rule("tool-call", -+ p.literal(FC_START) + p.space() + tool_choice + -+ p.zero_or_more(p.space() + tool_choice) + p.space() + p.literal(FC_END)); -+ } else { -+ tool_calls = p.trigger_rule("tool-call", -+ p.literal(FC_START) + p.space() + tool_choice + p.space() + p.literal(FC_END)); -+ } -+ -+ if (!require_tools) { -+ tool_calls = p.optional(tool_calls); -+ } -+ -+ auto content_before_tools = p.content(p.until(FC_START)); -+ return generation_prompt + reasoning + content_before_tools + tool_calls + end; -+ }); -+ -+ data.parser = parser.save(); -+ -+ if (include_grammar) { -+ data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED)); -+ data.grammar = build_grammar([&](const common_grammar_builder & builder) { -+ foreach_function(inputs.tools, [&](const json & tool) { -+ const auto & function = tool.at("function"); -+ auto schema = function.contains("parameters") ? function.at("parameters") : json::object(); -+ builder.resolve_refs(schema); -+ }); -+ if (has_response_format) { -+ auto schema = inputs.json_schema; -+ builder.resolve_refs(schema); -+ } -+ parser.build_grammar(builder, data.grammar_lazy); -+ }); -+ -+ data.grammar_triggers = { -+ { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, FC_START }, -+ }; -+ } -+ -+ return data; -+} -+ - // Cohere2 MoE (a.k.a. "North Code") parser. - // - // The assistant turn is fully marker-wrapped: -@@ -2707,6 +2892,15 @@ std::optional common_chat_try_specialized_template( - return common_chat_params_init_gigachat_v3(tmpl, params); - } - -+ // MiniMax-M3: the namespace token "]<]minimax[>[" collides with the autoparser's -+ // markup delimiters, so detect the template and use a dedicated parser. -+ if (src.find("]<]minimax[>[") != std::string::npos && -+ src.find("") != std::string::npos && -+ src.find("(1, std::min(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(1, std::min(n_batch, n_outputs)); + } +@@ -202,6 +211,26 @@ struct server_slot { + + std::vector generated_token_probs; + ++ // SERVER_TASK_TYPE_SCORE: shared-prefix token logprobs harvested ++ // incrementally across batch views (NaN = not yet produced) ++ std::vector 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> 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(); ++ 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(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 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 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 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 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> 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; diff --git a/backend/cpp/llama-cpp/prepare.sh b/backend/cpp/llama-cpp/prepare.sh index e99561179..b86924d2d 100644 --- a/backend/cpp/llama-cpp/prepare.sh +++ b/backend/cpp/llama-cpp/prepare.sh @@ -2,6 +2,7 @@ set -e + ## Patches ## Apply patches from the `patches` directory. Runs under set -e so a @@ -24,6 +25,9 @@ cp -r grpc-server.cpp llama.cpp/tools/grpc-server/ # unit test (compiled only when -DLLAMA_GRPC_BUILD_TESTS=ON). cp -r message_content.h llama.cpp/tools/grpc-server/ cp -r message_content_test.cpp llama.cpp/tools/grpc-server/ +# Generic passthrough parser staging and its standalone regression test. +cp -r passthrough_options.h llama.cpp/tools/grpc-server/ +cp -r passthrough_options_test.cpp llama.cpp/tools/grpc-server/ # Parent-death watcher (included by grpc-server.cpp) and its standalone unit # test (run via backend/cpp/run-unit-tests.sh; also buildable under ctest). cp -r parent_watch.h llama.cpp/tools/grpc-server/ @@ -58,4 +62,3 @@ else echo "add_subdirectory(grpc-server)" >> llama.cpp/tools/CMakeLists.txt fi set -e - diff --git a/backend/cpp/llama-cpp/run.sh b/backend/cpp/llama-cpp/run.sh index 1ccc1a37b..4c7ad19c5 100755 --- a/backend/cpp/llama-cpp/run.sh +++ b/backend/cpp/llama-cpp/run.sh @@ -12,10 +12,11 @@ grep -e "flags" /proc/cpuinfo | head -1 BINARY=llama-cpp-fallback -# CPU images (x86, arm64, darwin) ship a single llama-cpp-cpu-all built with ggml +# CPU images and most x86 GPU images ship a single llama-cpp-cpu-all built with ggml # CPU_ALL_VARIANTS: ggml's backend registry dlopens the best libggml-cpu-*.so for this -# host, so no shell-side AVX probing. GPU images (cublas/sycl/vulkan/hipblas) ship only -# llama-cpp-fallback (the accelerator does the compute), so fall back to it when absent. +# host, so no shell-side AVX probing. GPU arm64 images still ship llama-cpp-fallback +# until their builder toolchains support ggml's complete arm variant matrix, and so do +# the SYCL images, whose icpx compiler hangs on the sapphirerapids variant. if [ -e "$CURDIR"/llama-cpp-cpu-all ]; then BINARY=llama-cpp-cpu-all fi @@ -42,6 +43,27 @@ else if [ -d "$CURDIR/lib/hipblaslt/library" ]; then export HIPBLASLT_TENSILE_LIBPATH="$CURDIR"/lib/hipblaslt/library fi + # Backends built for Intel GPUs carry a copy of the Intel graphics driver, + # and libze_loader is only there in those builds. Level Zero looks for a + # driver on its own, so point it at the copy that came with this backend: it + # was built against the same C library, while the machine's own driver may + # not have been, and loading that one can crash on start. + # + # Anything the user set is left alone, so a machine with a graphics card + # newer than the driver carried here can still be told to use its own. + # Nothing is said about OpenCL: no OpenCL driver is carried, so anything we + # set there would leave OpenCL worse off than the machine's own setup. + if [ -e "$CURDIR/lib/libze_loader.so.1" ]; then + if [ -e "$CURDIR/lib/libze_intel_gpu.so.1" ] && [ -z "${ZE_ENABLE_ALT_DRIVERS:-}" ]; then + export ZE_ENABLE_ALT_DRIVERS="$CURDIR"/lib/libze_intel_gpu.so.1 + fi + # Ask the driver how much graphics memory is free. Without this, + # llama.cpp reads zero on an integrated graphics chip, because such a + # chip shares the system memory instead of having its own. + if [ -z "${ZES_ENABLE_SYSMAN:-}" ]; then + export ZES_ENABLE_SYSMAN=1 + fi + fi fi # If there is a lib/ld.so, use it @@ -55,4 +77,4 @@ echo "Using binary: $BINARY" exec "$CURDIR"/$BINARY "$@" # We should never reach this point, however just in case we do, run fallback -exec "$CURDIR"/llama-cpp-fallback "$@" \ No newline at end of file +exec "$CURDIR"/llama-cpp-fallback "$@" diff --git a/backend/cpp/privacy-filter/package.sh b/backend/cpp/privacy-filter/package.sh index dd839c70e..258ba9b5d 100755 --- a/backend/cpp/privacy-filter/package.sh +++ b/backend/cpp/privacy-filter/package.sh @@ -11,13 +11,7 @@ cp -rfv "$CURDIR/run.sh" "$CURDIR/package/" # The dynamic loader, renamed to lib/ld.so so run.sh can invoke it explicitly # (makes the image independent of the host's glibc layout). -if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then - cp -arfLv /lib64/ld-linux-x86-64.so.2 "$CURDIR/package/lib/ld.so" -elif [ -f "/lib/ld-linux-aarch64.so.1" ]; then - cp -arfLv /lib/ld-linux-aarch64.so.1 "$CURDIR/package/lib/ld.so" -else - echo "package.sh: unknown architecture" >&2; exit 1 -fi +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Bundle the binary's transitive shared deps (libstdc++, libgomp, and the apt # grpc++/protobuf/absl stack) by walking ldd — robust to whichever of those are diff --git a/backend/cpp/run-unit-tests.sh b/backend/cpp/run-unit-tests.sh index 603d13a91..a87faf563 100755 --- a/backend/cpp/run-unit-tests.sh +++ b/backend/cpp/run-unit-tests.sh @@ -37,9 +37,19 @@ if [ -z "$JSON_INC" ]; then fi # Active source dirs only - exclude per-variant build copies, dev snapshots and -# the vendored upstream llama.cpp tree. +# the vendored upstream checkouts. +# +# Every upstream checkout needs its own -not -path, and audio.cpp is safe TODAY +# only by luck: its 44 tests all put "test" at the FRONT (17 test-*.cpp, 27 +# test_*.cpp, zero *_test.cpp), so the glob below misses every one of them. +# Nothing enforces that. This gate runs on every PR for every backend, and it +# compiles each match as a STANDALONE translation unit with nothing but +# nlohmann/json on the include path, so the day upstream adds or renames one +# test to *_test.cpp the whole gate goes red repo-wide, on an Apache-2.0 file +# nobody here wrote. Exclude it now rather than diagnose that later. mapfile -t tests < <(find "$ROOT" -name '*_test.cpp' \ -not -path '*/llama.cpp/*' \ + -not -path '*/audio.cpp/*' \ -not -path '*-build/*' \ -not -path '*-dev/*' \ -not -path '*fallback*' | sort) diff --git a/backend/cpp/turboquant/Makefile b/backend/cpp/turboquant/Makefile index ac2de1e15..162d09fd1 100644 --- a/backend/cpp/turboquant/Makefile +++ b/backend/cpp/turboquant/Makefile @@ -1,7 +1,7 @@ # Pinned to the HEAD of feature/turboquant-kv-cache on https://github.com/TheTom/llama-cpp-turboquant. # Auto-bumped nightly by .github/workflows/bump_deps.yaml. -TURBOQUANT_VERSION?=c26cbdffcf6fc9b7430cd6b117757e9a3f70b7ea +TURBOQUANT_VERSION?=8a891f4b566efdbd3cea92fafee3227a0a267683 LLAMA_REPO?=https://github.com/TheTom/llama-cpp-turboquant CMAKE_ARGS?= @@ -47,6 +47,7 @@ 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 @@ -84,6 +85,7 @@ 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 diff --git a/backend/cpp/turboquant/package.sh b/backend/cpp/turboquant/package.sh index c4559a68d..c6ef6884d 100755 --- a/backend/cpp/turboquant/package.sh +++ b/backend/cpp/turboquant/package.sh @@ -24,34 +24,7 @@ if [ -d "$CURDIR/ggml-shared-libs" ]; then fi # 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 -else - echo "Error: Could not detect architecture" - exit 1 -fi +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries based on BUILD_TYPE GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh" diff --git a/backend/cpp/turboquant/run.sh b/backend/cpp/turboquant/run.sh index 84db6985a..44f1f942f 100755 --- a/backend/cpp/turboquant/run.sh +++ b/backend/cpp/turboquant/run.sh @@ -12,9 +12,12 @@ grep -e "flags" /proc/cpuinfo | head -1 BINARY=turboquant-fallback -# x86/arm64 ship a single turboquant-cpu-all built with ggml CPU_ALL_VARIANTS: ggml's +# CPU images and most x86 GPU images ship a single turboquant-cpu-all built with ggml +# CPU_ALL_VARIANTS: ggml's # backend registry dlopens the best libggml-cpu-*.so for this host, so no shell-side -# probing. ROCm ships only turboquant-fallback, so fall back to it when cpu-all is absent. +# probing. GPU arm64 images still ship turboquant-fallback until their builder toolchains +# support ggml's complete arm variant matrix, and so do the SYCL images, whose icpx +# compiler hangs on the sapphirerapids variant. if [ -e "$CURDIR"/turboquant-cpu-all ]; then BINARY=turboquant-cpu-all fi @@ -40,6 +43,27 @@ else if [ -d "$CURDIR/lib/hipblaslt/library" ]; then export HIPBLASLT_TENSILE_LIBPATH="$CURDIR"/lib/hipblaslt/library fi + # Backends built for Intel GPUs carry a copy of the Intel graphics driver, + # and libze_loader is only there in those builds. Level Zero looks for a + # driver on its own, so point it at the copy that came with this backend: it + # was built against the same C library, while the machine's own driver may + # not have been, and loading that one can crash on start. + # + # Anything the user set is left alone, so a machine with a graphics card + # newer than the driver carried here can still be told to use its own. + # Nothing is said about OpenCL: no OpenCL driver is carried, so anything we + # set there would leave OpenCL worse off than the machine's own setup. + if [ -e "$CURDIR/lib/libze_loader.so.1" ]; then + if [ -e "$CURDIR/lib/libze_intel_gpu.so.1" ] && [ -z "${ZE_ENABLE_ALT_DRIVERS:-}" ]; then + export ZE_ENABLE_ALT_DRIVERS="$CURDIR"/lib/libze_intel_gpu.so.1 + fi + # Ask the driver how much graphics memory is free. Without this, the + # backend reads zero on an integrated graphics chip, because such a chip + # shares the system memory instead of having its own. + if [ -z "${ZES_ENABLE_SYSMAN:-}" ]; then + export ZES_ENABLE_SYSMAN=1 + fi + fi fi # If there is a lib/ld.so, use it diff --git a/backend/go/acestep-cpp/package.sh b/backend/go/acestep-cpp/package.sh index 5fecf3455..2e37c81e5 100755 --- a/backend/go/acestep-cpp/package.sh +++ b/backend/go/acestep-cpp/package.sh @@ -17,40 +17,7 @@ cp -fv $CURDIR/libgoacestepcpp-*.dylib $CURDIR/package/ 2>/dev/null || true 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/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/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/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/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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries based on BUILD_TYPE # The GPU library packaging script will detect BUILD_TYPE and copy appropriate GPU libraries diff --git a/backend/go/ced/package.sh b/backend/go/ced/package.sh index ff20d727f..db5241143 100755 --- a/backend/go/ced/package.sh +++ b/backend/go/ced/package.sh @@ -22,34 +22,7 @@ if ! ls "$CURDIR"/package/lib/libced.* >/dev/null 2>&1; then exit 1 fi -if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then - 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 - 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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh" if [ -f "$GPU_LIB_SCRIPT" ]; then diff --git a/backend/go/cloud-proxy/provider_anthropic.go b/backend/go/cloud-proxy/provider_anthropic.go index d86f2ab8e..aa39e1ccb 100644 --- a/backend/go/cloud-proxy/provider_anthropic.go +++ b/backend/go/cloud-proxy/provider_anthropic.go @@ -32,7 +32,9 @@ import ( type anthropicRequest struct { Model string `json:"model"` MaxTokens int32 `json:"max_tokens"` - System string `json:"system,omitempty"` + // 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"` Messages []anthropicMessage `json:"messages"` Stream bool `json:"stream,omitempty"` Temperature *float64 `json:"temperature,omitempty"` @@ -52,9 +54,30 @@ type anthropicMessage struct { } type anthropicTool struct { - Name string `json:"name"` - Description string `json:"description,omitempty"` - InputSchema json.RawMessage `json:"input_schema"` + 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"` } // anthropicToolChoice mirrors the four shapes Anthropic accepts: @@ -81,8 +104,9 @@ 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"` + ToolUseID string `json:"tool_use_id,omitempty"` + ResultContent string `json:"content,omitempty"` + CacheControl *anthropicCacheControl `json:"cache_control,omitempty"` } type anthropicResponse struct { @@ -156,6 +180,11 @@ 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() { @@ -189,15 +218,54 @@ func buildAnthropicRequest(opts *pb.PredictOptions, cfg *proxyConfig, stream boo }) } } - req.System = strings.Join(systemParts, "\n\n") + // 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 + } + } 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, diff --git a/backend/go/cloud-proxy/provider_anthropic_test.go b/backend/go/cloud-proxy/provider_anthropic_test.go index 6119c97cf..2ed9a2127 100644 --- a/backend/go/cloud-proxy/provider_anthropic_test.go +++ b/backend/go/cloud-proxy/provider_anthropic_test.go @@ -328,3 +328,62 @@ 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"`)) +} diff --git a/backend/go/cloud-proxy/proxy.go b/backend/go/cloud-proxy/proxy.go index 4541ed772..f6e3a9f1b 100644 --- a/backend/go/cloud-proxy/proxy.go +++ b/backend/go/cloud-proxy/proxy.go @@ -48,6 +48,7 @@ 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 { @@ -106,6 +107,7 @@ 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(), diff --git a/backend/go/crispasr/Makefile b/backend/go/crispasr/Makefile index 64dc95cbb..c24edb9cf 100644 --- a/backend/go/crispasr/Makefile +++ b/backend/go/crispasr/Makefile @@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1) # CrispASR version (release tag) CRISPASR_REPO?=https://github.com/CrispStrobe/CrispASR -CRISPASR_VERSION?=306faee45fab641d54f9f941f075de1e9c0d3278 +CRISPASR_VERSION?=21901d3f7c23554f072964828363e49ddbc2dc68 SO_TARGET?=libgocrispasr.so CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF diff --git a/backend/go/crispasr/gocrispasr.go b/backend/go/crispasr/gocrispasr.go index 1fd64b08e..be431165d 100644 --- a/backend/go/crispasr/gocrispasr.go +++ b/backend/go/crispasr/gocrispasr.go @@ -67,7 +67,16 @@ const defaultTTSSampleRate = 24000 // resampling, so the WAV header must match it. Returns ok=false for non-piper // models (key absent) or an unreadable file, letting the caller fall back to // defaultTTSSampleRate. -func piperSampleRate(modelPath string) (int, bool) { +func piperSampleRate(modelPath string) (rate int, ok bool) { + // A malformed metadata length can make gguf-parser-go panic before it can + // return an error. Keep a bad voice file from crash-looping the backend. + defer func() { + if recover() != nil { + rate = 0 + ok = false + } + }() + // Only scalar architecture keys are read, so skip the large array metadata // (phoneme map) and mmap the header - same rationale as pkg/vram's reader. f, err := gguf.ParseGGUFFile(modelPath, gguf.UseMMap(), gguf.SkipLargeMetadata()) @@ -78,7 +87,7 @@ func piperSampleRate(modelPath string) (int, bool) { if !ok || kv.ValueType != gguf.GGUFMetadataValueTypeUint32 { return 0, false } - rate := int(kv.ValueUint32()) + rate = int(kv.ValueUint32()) if rate <= 0 { return 0, false } diff --git a/backend/go/crispasr/gocrispasr_samplerate_test.go b/backend/go/crispasr/gocrispasr_samplerate_test.go index 6b0cf726b..c36e7a75a 100644 --- a/backend/go/crispasr/gocrispasr_samplerate_test.go +++ b/backend/go/crispasr/gocrispasr_samplerate_test.go @@ -3,6 +3,7 @@ package main import ( "bytes" "encoding/binary" + "math" "os" "path/filepath" @@ -102,6 +103,24 @@ var _ = Describe("piper sample rate", func() { _, ok := piperSampleRate(p) Expect(ok).To(BeFalse()) }) + + It("returns ok=false instead of panicking on a malformed string length", func() { + p := filepath.Join(GinkgoT().TempDir(), "malformed.gguf") + var b bytes.Buffer + b.WriteString("GGUF") + Expect(binary.Write(&b, binary.LittleEndian, uint32(3))).To(Succeed()) + Expect(binary.Write(&b, binary.LittleEndian, uint64(0))).To(Succeed()) + Expect(binary.Write(&b, binary.LittleEndian, uint64(1))).To(Succeed()) + key := "general.name" + Expect(binary.Write(&b, binary.LittleEndian, uint64(len(key)))).To(Succeed()) + b.WriteString(key) + Expect(binary.Write(&b, binary.LittleEndian, ggufTypeString)).To(Succeed()) + Expect(binary.Write(&b, binary.LittleEndian, uint64(math.MaxInt64))).To(Succeed()) + Expect(os.WriteFile(p, b.Bytes(), 0o644)).To(Succeed()) + + _, ok := piperSampleRate(p) + Expect(ok).To(BeFalse()) + }) }) // End-to-end through the built .so. Gated on CRISPASR_PIPER_MODEL_PATH (a diff --git a/backend/go/crispasr/package.sh b/backend/go/crispasr/package.sh index 9b89dad1b..32a9e5764 100755 --- a/backend/go/crispasr/package.sh +++ b/backend/go/crispasr/package.sh @@ -17,40 +17,7 @@ cp -fv $CURDIR/libgocrispasr-*.dylib $CURDIR/package/ 2>/dev/null || true 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/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/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/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/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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Bundle espeak-ng (+ its libpcaudio/libsonic runtime deps) and its voice data so # the piper TTS backend can phonemize non-English text. CrispASR dlopens diff --git a/backend/go/depth-anything-cpp/Makefile b/backend/go/depth-anything-cpp/Makefile index e142607ab..fe9071631 100644 --- a/backend/go/depth-anything-cpp/Makefile +++ b/backend/go/depth-anything-cpp/Makefile @@ -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?=f4e17dea695dd12ae76bea98ba58030996b98118 +DEPTHANYTHING_VERSION?=2028b47ac75a8659c6a9aa617baf09be193eb55f ifeq ($(NATIVE),false) CMAKE_ARGS+=-DGGML_NATIVE=OFF @@ -82,6 +82,16 @@ else VARIANT_TARGETS = libdepthanythingcpp-fallback.dylib endif +## Builds the native engine variants and stops short of the Go binary. The +## variants depend only on sources/depth-anything.cpp, a clone pinned by +## DEPTHANYTHING_VERSION, so nothing in this target can observe a change +## elsewhere in the LocalAI tree. Dockerfile.golang calls it from a layer that +## copies in this Makefile and nothing else, which keeps the multi-minute ggml +## compile in the registry layer cache across builds whose only change is on the +## Go side. See .agents/ci-caching.md. +.PHONY: engine +engine: $(VARIANT_TARGETS) + depth-anything-cpp: main.go godepthanythingcpp.go $(VARIANT_TARGETS) CGO_ENABLED=0 $(GOCMD) build -tags "$(GO_TAGS)" -o depth-anything-cpp ./ diff --git a/backend/go/depth-anything-cpp/package.sh b/backend/go/depth-anything-cpp/package.sh index 5bbd5559b..d4d8e74b4 100755 --- a/backend/go/depth-anything-cpp/package.sh +++ b/backend/go/depth-anything-cpp/package.sh @@ -16,36 +16,7 @@ cp -avf $CURDIR/depth-anything-cpp $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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries based on BUILD_TYPE GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh" diff --git a/backend/go/face-detect/package.sh b/backend/go/face-detect/package.sh index 36ffa8993..a63393deb 100644 --- a/backend/go/face-detect/package.sh +++ b/backend/go/face-detect/package.sh @@ -25,34 +25,7 @@ cp -avf "$CURDIR"/libfacedetect.so* "$CURDIR/package/lib/" 2>/dev/null || { # Detect architecture and copy the core runtime libs libfacedetect.so links # against, plus the matching dynamic loader as lib/ld.so. -if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then - 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 - 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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries (CUDA/ROCm/Intel/Vulkan loader + ICDs + drivers) based on # BUILD_TYPE so the backend can reach the GPU without the runtime base image diff --git a/backend/go/localvqe/package.sh b/backend/go/localvqe/package.sh index 9f9f2533d..fc88f396d 100755 --- a/backend/go/localvqe/package.sh +++ b/backend/go/localvqe/package.sh @@ -21,34 +21,7 @@ cp -P $CURDIR/libggml*.dylib $CURDIR/package/ 2>/dev/null || true cp -fv $CURDIR/run.sh $CURDIR/package/ # Detect architecture and copy appropriate libraries -if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then - 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 - 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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries based on BUILD_TYPE GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh" diff --git a/backend/go/locate-anything-cpp/package.sh b/backend/go/locate-anything-cpp/package.sh index 1e6cbee80..c2f41806a 100755 --- a/backend/go/locate-anything-cpp/package.sh +++ b/backend/go/locate-anything-cpp/package.sh @@ -16,36 +16,7 @@ cp -avf $CURDIR/locate-anything-cpp $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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries based on BUILD_TYPE GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh" diff --git a/backend/go/moss-transcribe-cpp/package.sh b/backend/go/moss-transcribe-cpp/package.sh index 6cbf24b09..b8afb8fec 100755 --- a/backend/go/moss-transcribe-cpp/package.sh +++ b/backend/go/moss-transcribe-cpp/package.sh @@ -28,34 +28,7 @@ fi # Detect architecture and copy the core runtime libs libmoss-transcribe.so links # against, plus the matching dynamic loader as lib/ld.so. -if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then - 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 - 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 — system frameworks linked dynamically, no bundled libs needed" -else - echo "Error: Could not detect architecture" - exit 1 -fi +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries (CUDA/ROCm/Intel/Vulkan loader + ICDs + drivers) based # on BUILD_TYPE so the backend can reach the GPU without the runtime base image diff --git a/backend/go/moss-tts-cpp/package.sh b/backend/go/moss-tts-cpp/package.sh index e41d817b0..ba7ea33cb 100644 --- a/backend/go/moss-tts-cpp/package.sh +++ b/backend/go/moss-tts-cpp/package.sh @@ -17,40 +17,7 @@ cp -fv $CURDIR/libgomosstts-cpp-*.dylib $CURDIR/package/ 2>/dev/null || true 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/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/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/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/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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries based on BUILD_TYPE GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh" diff --git a/backend/go/omnivoice-cpp/package.sh b/backend/go/omnivoice-cpp/package.sh index 97a8d7809..5ff3f735e 100755 --- a/backend/go/omnivoice-cpp/package.sh +++ b/backend/go/omnivoice-cpp/package.sh @@ -17,40 +17,7 @@ cp -fv $CURDIR/libgomnivoicecpp-*.dylib $CURDIR/package/ 2>/dev/null || true 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/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/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/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/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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries based on BUILD_TYPE GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh" diff --git a/backend/go/opus/package.sh b/backend/go/opus/package.sh index 1e1aaeabf..32af77649 100644 --- a/backend/go/opus/package.sh +++ b/backend/go/opus/package.sh @@ -28,31 +28,7 @@ if command -v pkg-config >/dev/null 2>&1 && pkg-config --exists opus; then fi # Detect architecture and copy appropriate libraries -if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then - 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/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 - 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/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 — system libraries linked dynamically, no bundled loader needed" -else - echo "Warning: Could not detect architecture for system library bundling" -fi +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" echo "Packaging completed successfully" ls -liah $CURDIR/package/ diff --git a/backend/go/parakeet-cpp/Makefile b/backend/go/parakeet-cpp/Makefile index 80512a2b1..fce2f3f2c 100644 --- a/backend/go/parakeet-cpp/Makefile +++ b/backend/go/parakeet-cpp/Makefile @@ -1,6 +1,6 @@ # parakeet-cpp backend Makefile. # -# Upstream pin lives below as PARAKEET_VERSION?=1da853421de9710cbe894a0110711de5a0516486 +# Upstream pin lives below as PARAKEET_VERSION?=1bfbebfaaf493866f49597cd3b7901959d395c60 # (.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?=1da853421de9710cbe894a0110711de5a0516486 +PARAKEET_VERSION?=1bfbebfaaf493866f49597cd3b7901959d395c60 PARAKEET_REPO?=https://github.com/mudler/parakeet.cpp GOCMD?=go diff --git a/backend/go/parakeet-cpp/package.sh b/backend/go/parakeet-cpp/package.sh index af8e6b9e1..0454aa334 100755 --- a/backend/go/parakeet-cpp/package.sh +++ b/backend/go/parakeet-cpp/package.sh @@ -28,34 +28,7 @@ fi # Detect architecture and copy the core runtime libs libparakeet.so links # against, plus the matching dynamic loader as lib/ld.so. -if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then - 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 - 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 — system frameworks linked dynamically, no bundled libs needed" -else - echo "Error: Could not detect architecture" - exit 1 -fi +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries (CUDA/ROCm/Intel/Vulkan loader + ICDs + drivers) # based on BUILD_TYPE so the backend can reach the GPU without the runtime diff --git a/backend/go/piper/package.sh b/backend/go/piper/package.sh index a3f59c95f..1c37ca24b 100755 --- a/backend/go/piper/package.sh +++ b/backend/go/piper/package.sh @@ -16,46 +16,7 @@ cp -rfv $CURDIR/run.sh $CURDIR/package/ cp -rfLv $CURDIR/sources/go-piper/piper-phonemize/pi/lib/* $CURDIR/package/lib/ # Detect architecture and copy appropriate libraries -if [ "$(uname)" = "Darwin" ]; then - # macOS has no glibc loader to bundle. The piper binary links its bundled - # libs (libucd, libespeak-ng, libpiper_phonemize, libonnxruntime) via - # @rpath but ships with no LC_RPATH, so dyld aborts at launch with - # "Library not loaded: @rpath/libucd.dylib ... no LC_RPATH's found". - # Add an @loader_path/lib rpath so @rpath resolves to package/lib/. - echo "Detected macOS; adding @loader_path/lib rpath so bundled libs resolve via @rpath..." - install_name_tool -add_rpath @loader_path/lib "$CURDIR/package/piper" -elif [ -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/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/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/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/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 -else - echo "Error: Could not detect architecture" - exit 1 -fi +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "$CURDIR/package/piper" echo "Packaging completed successfully" ls -liah $CURDIR/package/ diff --git a/backend/go/qwen3-tts-cpp/Makefile b/backend/go/qwen3-tts-cpp/Makefile index 8aa0a6df6..e75eb6cca 100644 --- a/backend/go/qwen3-tts-cpp/Makefile +++ b/backend/go/qwen3-tts-cpp/Makefile @@ -7,6 +7,16 @@ GO_TAGS?= JOBS?=$(shell nproc --ignore=1) # qwentts.cpp version +# +# Held at 35ebe537 rather than tracking latest: abab6b3 hangs in synthesis. +# TTS() never returns from the native call, so tests-qwen3-tts-cpp goes from +# ~5 minutes to the 20 minute Go test timeout. Reproduced on master on +# 2026-08-01 and again on re-run, and the bump PR (#11241) was merged with +# this same check already red. +# +# The regression is in 35ebe537..abab6b3, three upstream commits whose only +# functional change is 26dd8adb, "predictor: unroll the frame into one cgraph +# and sample in standard ops". Restore the bump once that is fixed upstream. QWEN3TTS_REPO?=https://github.com/ServeurpersoCom/qwentts.cpp QWEN3TTS_CPP_VERSION?=35ebe5376b82a0a59d008586d55bbe623d449011 SO_TARGET?=libgoqwen3ttscpp.so diff --git a/backend/go/qwen3-tts-cpp/package.sh b/backend/go/qwen3-tts-cpp/package.sh index 11d4c57c3..1d00dbcfc 100755 --- a/backend/go/qwen3-tts-cpp/package.sh +++ b/backend/go/qwen3-tts-cpp/package.sh @@ -17,40 +17,7 @@ cp -fv $CURDIR/libgoqwen3ttscpp-*.dylib $CURDIR/package/ 2>/dev/null || true 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/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/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/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/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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries based on BUILD_TYPE GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh" diff --git a/backend/go/rfdetr-cpp/Makefile b/backend/go/rfdetr-cpp/Makefile index 29f2842eb..ba06e1244 100644 --- a/backend/go/rfdetr-cpp/Makefile +++ b/backend/go/rfdetr-cpp/Makefile @@ -11,7 +11,7 @@ JOBS?=$(shell nproc --ignore=1) # build; leaving this on `master` always picks up the latest C-API surface # (incl. the per-detection accessor functions used by gorfdetrcpp.go). RFDETR_REPO?=https://github.com/localai-org/rf-detr.cpp.git -RFDETR_VERSION?=65c0ffcc9a9bc9dae38252f63d0417c9845a6cf7 +RFDETR_VERSION?=98d0f381b832ef08a608b65c7dd78db066ed8b9a ifeq ($(NATIVE),false) CMAKE_ARGS+=-DGGML_NATIVE=OFF diff --git a/backend/go/rfdetr-cpp/package.sh b/backend/go/rfdetr-cpp/package.sh index 17319bf27..7ed4ea4fa 100755 --- a/backend/go/rfdetr-cpp/package.sh +++ b/backend/go/rfdetr-cpp/package.sh @@ -16,36 +16,7 @@ cp -avf $CURDIR/rfdetr-cpp $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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries based on BUILD_TYPE GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh" diff --git a/backend/go/sam3-cpp/package.sh b/backend/go/sam3-cpp/package.sh index a648ee93c..ea1776915 100755 --- a/backend/go/sam3-cpp/package.sh +++ b/backend/go/sam3-cpp/package.sh @@ -16,36 +16,7 @@ cp -avf $CURDIR/sam3-cpp $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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries based on BUILD_TYPE GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh" diff --git a/backend/go/sherpa-onnx/package.sh b/backend/go/sherpa-onnx/package.sh index 5a596e490..90b4fbbc8 100755 --- a/backend/go/sherpa-onnx/package.sh +++ b/backend/go/sherpa-onnx/package.sh @@ -10,34 +10,7 @@ cp -avf $CURDIR/sherpa-onnx $CURDIR/package/ cp -avf $CURDIR/run.sh $CURDIR/package/ cp -rfLv $CURDIR/backend-assets/lib/* $CURDIR/package/lib/ -if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then - 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 - 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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh" if [ -f "$GPU_LIB_SCRIPT" ]; then diff --git a/backend/go/silero-vad/package.sh b/backend/go/silero-vad/package.sh index a96ff4c8b..f19dd55ca 100755 --- a/backend/go/silero-vad/package.sh +++ b/backend/go/silero-vad/package.sh @@ -15,45 +15,7 @@ cp -avf $CURDIR/run.sh $CURDIR/package/ cp -rfLv $CURDIR/backend-assets/lib/* $CURDIR/package/lib/ # Detect architecture and copy appropriate libraries -if [ "$(uname)" = "Darwin" ]; then - # macOS has no glibc loader to bundle. silero-vad links its bundled - # libonnxruntime via @rpath but ships with no LC_RPATH, so dyld can't find - # it at runtime. Add an @loader_path/lib rpath so @rpath resolves to - # package/lib/ (matching the piper darwin fix, #10525). - echo "Detected macOS; adding @loader_path/lib rpath so bundled libs resolve via @rpath..." - install_name_tool -add_rpath @loader_path/lib "$CURDIR/package/silero-vad" -elif [ -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/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/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/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/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 -else - echo "Error: Could not detect architecture" - exit 1 -fi +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "$CURDIR/package/silero-vad" echo "Packaging completed successfully" ls -liah $CURDIR/package/ diff --git a/backend/go/stablediffusion-ggml/Makefile b/backend/go/stablediffusion-ggml/Makefile index f7daa6c15..241151774 100644 --- a/backend/go/stablediffusion-ggml/Makefile +++ b/backend/go/stablediffusion-ggml/Makefile @@ -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?=2d0385ba85af358f7115dda608a63eafd9de7ffd +STABLEDIFFUSION_GGML_VERSION?=c6beeef35526c6dc94b74a7fb69f9d2e6a2a7a12 CMAKE_ARGS+=-DGGML_MAX_NAME=128 diff --git a/backend/go/stablediffusion-ggml/gosd.go b/backend/go/stablediffusion-ggml/gosd.go index 219b78470..e1567bd18 100644 --- a/backend/go/stablediffusion-ggml/gosd.go +++ b/backend/go/stablediffusion-ggml/gosd.go @@ -5,6 +5,7 @@ import ( "os" "path/filepath" "runtime" + "strconv" "strings" "unsafe" @@ -18,6 +19,7 @@ type SDGGML struct { threads int sampleMethod string cfgScale float32 + vaeTiling vaeTiling } var ( @@ -43,6 +45,62 @@ var ( VidGenParamsSetVideoFrames func(params uintptr, n int) ) +type vaeTiling struct { + enabled bool + tileSizeX int + tileSizeY int + hasTileSize bool + targetOverlap float32 + hasOverlap bool +} + +func parseVAETiling(options []string) vaeTiling { + var t vaeTiling + for _, op := range options { + name, value, hasValue := strings.Cut(op, ":") + switch name { + case "vae_tiling": + // A bare flag reads as "on", matching "diffusion_model". The truthy + // spellings are the ones load_model already accepts for its own + // bool options, so an author does not have to remember two + // conventions. + t.enabled = !hasValue || value == "true" || value == "1" + case "vae_tile_size": + if x, y, ok := parseTileSize(value); ok { + t.tileSizeX, t.tileSizeY, t.hasTileSize = x, y, true + } + case "vae_tile_overlap": + if f, err := strconv.ParseFloat(value, 32); err == nil && f >= 0 { + t.targetOverlap, t.hasOverlap = float32(f), true + } + } + } + return t +} + +// parseTileSize accepts "512" for a square tile and "512x384" for a +// rectangular one. +// +// A value it cannot make sense of is reported as absent rather than as a zero. +// The caller only calls the upstream setter when a size was given, so a typo +// leaves the library's own default in place instead of installing a degenerate +// tiling that would fail at generation time. +func parseTileSize(value string) (int, int, bool) { + xs, ys, split := strings.Cut(value, "x") + if !split { + ys = xs + } + x, err := strconv.Atoi(xs) + if err != nil || x <= 0 { + return 0, 0, false + } + y, err := strconv.Atoi(ys) + if err != nil || y <= 0 { + return 0, 0, false + } + return x, y, true +} + // Copied from Purego internal/strings // TODO: We should upstream sending []string func hasSuffix(s, suffix string) bool { @@ -100,6 +158,9 @@ func (sd *SDGGML) Load(opts *pb.ModelOptions) error { } sd.cfgScale = opts.CFGScale + // Read from the unfiltered list: none of the tiling options name a path, so + // the resolution pass above neither rewrites nor drops them. + sd.vaeTiling = parseVAETiling(opts.Options) ret := LoadModel(modelFile, modelPathC, options, opts.Threads, diffusionModel) runtime.KeepAlive(keepAlive) @@ -148,8 +209,22 @@ func (sd *SDGGML) GenerateImage(opts *pb.GenerateImageRequest) error { ImgGenParamsSetPrompts(p, t, negative) ImgGenParamsSetDimensions(p, int(opts.Width), int(opts.Height)) ImgGenParamsSetSeed(p, int64(opts.Seed)) + // Tiling decodes the latent in overlapping tiles, so the VAE compute buffer + // scales with the tile rather than with the image. That is the difference + // between working and failing on any device that caps a single allocation + // (RADV reports a 4GiB maxMemoryAllocationSize, for one) or that simply + // does not have the VRAM for a full-frame decode at high resolution. + // + // Only the setters the operator configured are called, so an unset tile + // size or overlap keeps the library's own default. vaep := ImgGenParamsGetVaeTilingParams(p) - TilingParamsSetEnabled(vaep, false) + TilingParamsSetEnabled(vaep, sd.vaeTiling.enabled) + if sd.vaeTiling.hasTileSize { + TilingParamsSetTileSizes(vaep, sd.vaeTiling.tileSizeX, sd.vaeTiling.tileSizeY) + } + if sd.vaeTiling.hasOverlap { + TilingParamsSetTargetOverlap(vaep, sd.vaeTiling.targetOverlap) + } ret := GenImage(p, int(opts.Step), dst, sd.cfgScale, srcImage, strength, maskImage, refImages, refImagesCount) runtime.KeepAlive(keepAlive) diff --git a/backend/go/stablediffusion-ggml/gosd_test.go b/backend/go/stablediffusion-ggml/gosd_test.go new file mode 100644 index 000000000..3476700c6 --- /dev/null +++ b/backend/go/stablediffusion-ggml/gosd_test.go @@ -0,0 +1,180 @@ +package main + +import ( + "testing" + + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestStableDiffusionGGML(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "stablediffusion-ggml backend test suite") +} + +var _ = DescribeTable("parseVAETiling enablement", + func(options []string, want bool) { + Expect(parseVAETiling(options).enabled).To(Equal(want)) + }, + Entry("explicit true", []string{"vae_tiling:true"}, true), + Entry("one is truthy", []string{"vae_tiling:1"}, true), + // "diffusion_model" is already a bare flag in this option list, so accept + // the same shape rather than making vae_tiling the one option that demands + // a value. + Entry("bare flag", []string{"vae_tiling"}, true), + Entry("explicit false", []string{"vae_tiling:false"}, false), + Entry("anything else is false", []string{"vae_tiling:maybe"}, false), + // Tiling trades a little quality at the tile seams for a much smaller + // compute buffer, so it must not switch itself on for the many models that + // never needed it. + Entry("absent leaves it off", []string{"diffusion_model", "sampler:euler"}, false), + Entry("nil options", []string(nil), false), +) + +var _ = DescribeTable("parseVAETiling tile size", + func(options []string, wantSet bool, wantX, wantY int) { + got := parseVAETiling(options) + Expect(got.hasTileSize).To(Equal(wantSet)) + if wantSet { + Expect(got.tileSizeX).To(Equal(wantX)) + Expect(got.tileSizeY).To(Equal(wantY)) + } + }, + Entry("one number is a square tile", []string{"vae_tile_size:512"}, true, 512, 512), + Entry("rectangular", []string{"vae_tile_size:512x384"}, true, 512, 384), + // Absent must stay absent: the caller only invokes the upstream setter when + // a size was given, so this is what preserves the library's own default + // instead of pushing a zero. + Entry("absent", []string{"vae_tiling:true"}, false, 0, 0), + // A typo must not silently become a zero tile size and break a model that + // would otherwise have worked on the default. + Entry("unparseable", []string{"vae_tile_size:banana"}, false, 0, 0), + Entry("zero rejected", []string{"vae_tile_size:0"}, false, 0, 0), + Entry("negative rejected", []string{"vae_tile_size:-8"}, false, 0, 0), + Entry("half unparseable", []string{"vae_tile_size:512xbanana"}, false, 0, 0), +) + +var _ = DescribeTable("parseVAETiling target overlap", + func(options []string, wantSet bool, want float32) { + got := parseVAETiling(options) + Expect(got.hasOverlap).To(Equal(wantSet)) + if wantSet { + Expect(got.targetOverlap).To(Equal(want)) + } + }, + Entry("fraction", []string{"vae_tile_overlap:0.25"}, true, float32(0.25)), + Entry("zero is a legitimate overlap", []string{"vae_tile_overlap:0"}, true, float32(0)), + Entry("absent", []string{"vae_tiling:true"}, false, float32(0)), + Entry("unparseable", []string{"vae_tile_overlap:banana"}, false, float32(0)), + Entry("negative rejected", []string{"vae_tile_overlap:-0.5"}, false, float32(0)), +) + +// fakeSDLib swaps the purego bindings for recorders so the wiring between the +// parsed options and the upstream tiling setters can be exercised without the +// shared library. The bindings are package-level vars, which is the only seam +// available here; every one the code under test touches must be set or the +// call panics on a nil func. +type fakeSDLib struct { + tilingEnabled bool + tileSizeCalls int + tileSizeX int + tileSizeY int + overlapCalls int + targetOverlap float32 +} + +// install points the bindings at the recorder and restores them afterwards, so +// one spec cannot leak fakes into the next. +func (f *fakeSDLib) install() { + savedImgGenParamsNew := ImgGenParamsNew + savedImgGenParamsSetPrompts := ImgGenParamsSetPrompts + savedImgGenParamsSetDimensions := ImgGenParamsSetDimensions + savedImgGenParamsSetSeed := ImgGenParamsSetSeed + savedImgGenParamsGetVaeTilingParams := ImgGenParamsGetVaeTilingParams + savedTilingParamsSetEnabled := TilingParamsSetEnabled + savedTilingParamsSetTileSizes := TilingParamsSetTileSizes + savedTilingParamsSetTargetOverlap := TilingParamsSetTargetOverlap + savedGenImage := GenImage + savedLoadModel := LoadModel + + DeferCleanup(func() { + ImgGenParamsNew = savedImgGenParamsNew + ImgGenParamsSetPrompts = savedImgGenParamsSetPrompts + ImgGenParamsSetDimensions = savedImgGenParamsSetDimensions + ImgGenParamsSetSeed = savedImgGenParamsSetSeed + ImgGenParamsGetVaeTilingParams = savedImgGenParamsGetVaeTilingParams + TilingParamsSetEnabled = savedTilingParamsSetEnabled + TilingParamsSetTileSizes = savedTilingParamsSetTileSizes + TilingParamsSetTargetOverlap = savedTilingParamsSetTargetOverlap + GenImage = savedGenImage + LoadModel = savedLoadModel + }) + + ImgGenParamsNew = func() uintptr { return 1 } + ImgGenParamsSetPrompts = func(uintptr, string, string) {} + ImgGenParamsSetDimensions = func(uintptr, int, int) {} + ImgGenParamsSetSeed = func(uintptr, int64) {} + ImgGenParamsGetVaeTilingParams = func(uintptr) uintptr { return 2 } + TilingParamsSetEnabled = func(_ uintptr, enabled bool) { f.tilingEnabled = enabled } + TilingParamsSetTileSizes = func(_ uintptr, x, y int) { + f.tileSizeCalls++ + f.tileSizeX, f.tileSizeY = x, y + } + TilingParamsSetTargetOverlap = func(_ uintptr, o float32) { + f.overlapCalls++ + f.targetOverlap = o + } + GenImage = func(uintptr, int, string, float32, string, float32, string, []uintptr, int) int { return 0 } + LoadModel = func(string, string, []uintptr, int32, int) int { return 0 } +} + +var _ = Describe("GenerateImage VAE tiling", func() { + var fake *fakeSDLib + + BeforeEach(func() { + fake = &fakeSDLib{} + fake.install() + }) + + // generate drives the real Load and GenerateImage so the options travel the + // path they travel in production, with only the C boundary faked. + generate := func(options []string) { + sd := &SDGGML{} + Expect(sd.Load(&pb.ModelOptions{Options: options})).To(Succeed()) + Expect(sd.GenerateImage(&pb.GenerateImageRequest{Width: 1024, Height: 1024})).To(Succeed()) + } + + It("enables tiling when the model asks for it", func() { + generate([]string{"vae_tiling:true"}) + + Expect(fake.tilingEnabled).To(BeTrue()) + }) + + // The pre-existing behaviour: every model that never asked for tiling must + // still get it switched off. + It("leaves tiling off by default", func() { + generate([]string{"sampler:euler"}) + + Expect(fake.tilingEnabled).To(BeFalse()) + }) + + It("applies a configured tile size and overlap", func() { + generate([]string{"vae_tiling:true", "vae_tile_size:512x384", "vae_tile_overlap:0.25"}) + + Expect(fake.tileSizeCalls).To(Equal(1)) + Expect(fake.tileSizeX).To(Equal(512)) + Expect(fake.tileSizeY).To(Equal(384)) + Expect(fake.overlapCalls).To(Equal(1)) + Expect(fake.targetOverlap).To(Equal(float32(0.25))) + }) + + // Not calling the setters is what preserves the library's own defaults, so + // unconfigured values must leave them untouched rather than send a zero. + It("leaves unset tiling parameters alone", func() { + generate([]string{"vae_tiling:true"}) + + Expect(fake.tileSizeCalls).To(BeZero()) + Expect(fake.overlapCalls).To(BeZero()) + }) +}) diff --git a/backend/go/stablediffusion-ggml/package.sh b/backend/go/stablediffusion-ggml/package.sh index 922fb71ea..68585516d 100755 --- a/backend/go/stablediffusion-ggml/package.sh +++ b/backend/go/stablediffusion-ggml/package.sh @@ -17,40 +17,7 @@ cp -avf $CURDIR/stablediffusion-ggml $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/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/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/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/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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries based on BUILD_TYPE # The GPU library packaging script will detect BUILD_TYPE and copy appropriate GPU libraries diff --git a/backend/go/supertonic/package.sh b/backend/go/supertonic/package.sh index 678ca5ead..b29ece0d9 100755 --- a/backend/go/supertonic/package.sh +++ b/backend/go/supertonic/package.sh @@ -10,36 +10,7 @@ cp -avf $CURDIR/supertonic $CURDIR/package/ cp -avf $CURDIR/run.sh $CURDIR/package/ cp -rfLv $CURDIR/backend-assets/lib/* $CURDIR/package/lib/ -if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then - 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 - 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 - # macOS: dyld resolves the bundled .dylib via DYLD_LIBRARY_PATH (set in - # run.sh); there is no ld.so loader nor glibc to bundle. - echo "Detected Darwin" -else - echo "Error: Could not detect architecture" - exit 1 -fi +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh" if [ -f "$GPU_LIB_SCRIPT" ]; then diff --git a/backend/go/trellis2cpp/.gitignore b/backend/go/trellis2cpp/.gitignore new file mode 100644 index 000000000..b5cfa6aba --- /dev/null +++ b/backend/go/trellis2cpp/.gitignore @@ -0,0 +1,6 @@ +package/ +sources/ +.cache/ +build-*/ +variants/ +trellis2cpp diff --git a/backend/go/trellis2cpp/Makefile b/backend/go/trellis2cpp/Makefile new file mode 100644 index 000000000..50500edd5 --- /dev/null +++ b/backend/go/trellis2cpp/Makefile @@ -0,0 +1,132 @@ +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?=2f3e6e26edbbaaf8ce93d092f16f46968a366a6a + +# 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 diff --git a/backend/go/trellis2cpp/glb_mesh.go b/backend/go/trellis2cpp/glb_mesh.go new file mode 100644 index 000000000..144c9243a --- /dev/null +++ b/backend/go/trellis2cpp/glb_mesh.go @@ -0,0 +1,252 @@ +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 +} diff --git a/backend/go/trellis2cpp/glb_mesh_test.go b/backend/go/trellis2cpp/glb_mesh_test.go new file mode 100644 index 000000000..8ba4f9264 --- /dev/null +++ b/backend/go/trellis2cpp/glb_mesh_test.go @@ -0,0 +1,74 @@ +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")) + }) +}) diff --git a/backend/go/trellis2cpp/main.go b/backend/go/trellis2cpp/main.go new file mode 100644 index 000000000..14938bdd6 --- /dev/null +++ b/backend/go/trellis2cpp/main.go @@ -0,0 +1,50 @@ +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) + } +} diff --git a/backend/go/trellis2cpp/package.sh b/backend/go/trellis2cpp/package.sh new file mode 100755 index 000000000..20d7dedc2 --- /dev/null +++ b/backend/go/trellis2cpp/package.sh @@ -0,0 +1,63 @@ +#!/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/ diff --git a/backend/go/trellis2cpp/run.sh b/backend/go/trellis2cpp/run.sh new file mode 100755 index 000000000..78f5ec362 --- /dev/null +++ b/backend/go/trellis2cpp/run.sh @@ -0,0 +1,61 @@ +#!/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 "$@" diff --git a/backend/go/trellis2cpp/trellis2.go b/backend/go/trellis2cpp/trellis2.go new file mode 100644 index 000000000..addbf1bb0 --- /dev/null +++ b/backend/go/trellis2cpp/trellis2.go @@ -0,0 +1,511 @@ +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 `_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) +} diff --git a/backend/go/trellis2cpp/trellis2_test.go b/backend/go/trellis2cpp/trellis2_test.go new file mode 100644 index 000000000..d492756e4 --- /dev/null +++ b/backend/go/trellis2cpp/trellis2_test.go @@ -0,0 +1,247 @@ +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()) + }) +}) diff --git a/backend/go/valkey-store/Makefile b/backend/go/valkey-store/Makefile new file mode 100644 index 000000000..0d7ddb327 --- /dev/null +++ b/backend/go/valkey-store/Makefile @@ -0,0 +1,17 @@ +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 diff --git a/backend/go/valkey-store/config.go b/backend/go/valkey-store/config.go new file mode 100644 index 000000000..02c5711b4 --- /dev/null +++ b/backend/go/valkey-store/config.go @@ -0,0 +1,261 @@ +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 "" +} diff --git a/backend/go/valkey-store/encoding.go b/backend/go/valkey-store/encoding.go new file mode 100644 index 000000000..cec086070 --- /dev/null +++ b/backend/go/valkey-store/encoding.go @@ -0,0 +1,80 @@ +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) +} diff --git a/backend/go/valkey-store/encoding_test.go b/backend/go/valkey-store/encoding_test.go new file mode 100644 index 000000000..1c89f1ab1 --- /dev/null +++ b/backend/go/valkey-store/encoding_test.go @@ -0,0 +1,77 @@ +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()) + }) +}) diff --git a/backend/go/valkey-store/main.go b/backend/go/valkey-store/main.go new file mode 100644 index 000000000..ba030899a --- /dev/null +++ b/backend/go/valkey-store/main.go @@ -0,0 +1,25 @@ +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) + } +} diff --git a/backend/go/valkey-store/package.sh b/backend/go/valkey-store/package.sh new file mode 100644 index 000000000..c6b487836 --- /dev/null +++ b/backend/go/valkey-store/package.sh @@ -0,0 +1,12 @@ +#!/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/ diff --git a/backend/go/valkey-store/run.sh b/backend/go/valkey-store/run.sh new file mode 100644 index 000000000..caf370e2e --- /dev/null +++ b/backend/go/valkey-store/run.sh @@ -0,0 +1,6 @@ +#!/bin/bash +set -ex + +CURDIR=$(dirname "$(realpath "$0")") + +exec "$CURDIR"/valkey-store "$@" diff --git a/backend/go/valkey-store/store.go b/backend/go/valkey-store/store.go new file mode 100644 index 000000000..b2b8814a5 --- /dev/null +++ b/backend/go/valkey-store/store.go @@ -0,0 +1,692 @@ +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() +} diff --git a/backend/go/valkey-store/store_suite_test.go b/backend/go/valkey-store/store_suite_test.go new file mode 100644 index 000000000..d15f3fe2a --- /dev/null +++ b/backend/go/valkey-store/store_suite_test.go @@ -0,0 +1,13 @@ +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") +} diff --git a/backend/go/valkey-store/store_test.go b/backend/go/valkey-store/store_test.go new file mode 100644 index 000000000..c6d563d2b --- /dev/null +++ b/backend/go/valkey-store/store_test.go @@ -0,0 +1,598 @@ +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")) + }) +}) diff --git a/backend/go/vibevoice-cpp/package.sh b/backend/go/vibevoice-cpp/package.sh index 62860b8d6..36feaee9e 100755 --- a/backend/go/vibevoice-cpp/package.sh +++ b/backend/go/vibevoice-cpp/package.sh @@ -17,34 +17,7 @@ cp -fv $CURDIR/libgovibevoicecpp-*.dylib $CURDIR/package/ 2>/dev/null || true cp -fv $CURDIR/run.sh $CURDIR/package/ # Detect architecture and copy appropriate libraries -if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then - 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 - 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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries based on BUILD_TYPE GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh" diff --git a/backend/go/vllm-cpp/Makefile b/backend/go/vllm-cpp/Makefile index 0f009340c..11173432f 100644 --- a/backend/go/vllm-cpp/Makefile +++ b/backend/go/vllm-cpp/Makefile @@ -11,7 +11,30 @@ JOBS?=$(shell nproc --ignore=1 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || e # vllm.cpp version VLLM_CPP_REPO?=https://github.com/mudler/vllm.cpp -VLLM_CPP_VERSION?=9e1c9025ae61167a3335454d7cc0de6093c21845 +VLLM_CPP_VERSION?=0757cac231ecd571a83c4fd2f50805c9251fc225 + +# MLX GEMM provider (darwin/metal only; see the metal branch below for why). +# Consumed as the prebuilt pip wheel: building MLX from source needs `xcrun +# metal`, i.e. a full Xcode the macOS runners do not have, while the wheel ships +# include/, lib/libmlx.dylib and the compiled mlx.metallib ready to link. +# +# DEFAULT ON, but ONLY because VLLM_CPP_VERSION above is pinned at or past +# vllm.cpp 89c46aeb, which SHAPE-GATES the provider to prefill. The ordering is +# load-bearing, not incidental: +# +# pin >= 89c46aeb, MLX on -> 99.1% of MLX-LM (gated: prefill only) +# pin < 89c46aeb, MLX on -> ~51% (ungated: it also takes decode) +# +# MLX's steel GEMM wins prefill (537 ms TTFT against 602) and loses decode badly, +# because the provider pays an mx::eval sync plus an output memcpy per call and +# decode makes ~112 calls per TOKEN. Ungated it does both; gated it does only the +# good half. So if this pin is ever moved BACKWARDS, this default must go with it. +VLLM_CPP_MLX?=on +MLX_VERSION?=0.29.4 +MLX_VENV?=$(abspath ./mlx-venv) +# Resolved lazily (recursive `=`, not `:=`): the glob only matches once the venv +# target has run, and the interpreter version in the path varies per runner. +MLX_ROOT=$(shell echo $(MLX_VENV)/lib/python*/site-packages/mlx) # The backend consumes only the stable C ABI (libvllm + include/vllm.h), so the # server, examples and tests of the engine are never built here. @@ -49,6 +72,23 @@ else ifeq ($(BUILD_TYPE),vulkan) CMAKE_ARGS+=-DVLLM_CPP_VULKAN=ON -DVLLM_CPP_CUDA=OFF else ifeq ($(BUILD_TYPE),metal) CMAKE_ARGS+=-DVLLM_CPP_METAL=ON + # The optional MLX GEMM provider. vllm.cpp keeps it OFF by default because it + # is a ~19 MB libmlx.dylib plus a ~105 MB mlx.metallib, and upstream's + # position is that it must earn that cost by measurement. It does, on the + # only hardware this build targets: measured on an Apple M4 against the + # native MSL GEMM in the SAME binary (arms toggled by + # VT_OP_PROVIDER_DISABLE=mlx), Qwen3-1.7B-bf16 p=512 g=128, it is 1.5x to + # 2.2x aggregate throughput and 2x to 3x faster TTFT, at equal peak memory + # and bit-identical output on every parity shape. See vllm.cpp + # docs/BENCHMARKS.md "MLX GEMM provider A/B on Apple M4". + # + # MLX delegates the dense GEMM ONLY: kPagedAttention stays vllm.cpp's own + # kernel, because MLX has no paged-KV primitive at all. + # + # Set VLLM_CPP_MLX=off for a Metal build without it (smaller image, slower). + ifeq ($(VLLM_CPP_MLX),on) + MLX_ENABLED=1 + endif else CMAKE_ARGS+=-DVLLM_CPP_CUDA=OFF endif @@ -68,10 +108,54 @@ sources/vllm.cpp: git fetch --depth 1 origin $(VLLM_CPP_VERSION) && \ git checkout FETCH_HEAD -$(LIB): sources/vllm.cpp +ifeq ($(MLX_ENABLED),1) +# A stamp FILE, not a phony target: a phony prerequisite is always "newer" than +# $(LIB) and would re-link libvllm on every invocation. Keyed on the version so +# a MLX_VERSION bump reinstalls instead of silently reusing the old wheel. +MLX_STAMP=$(MLX_VENV)/.mlx-$(MLX_VERSION).stamp +MLX_CMAKE_ARGS=-DVLLM_CPP_MLX=ON -DMLX_ROOT=$(MLX_ROOT) + +$(MLX_STAMP): + @if [ ! -x "$(MLX_VENV)/bin/pip" ]; then \ + python3 -m venv "$(MLX_VENV)" || { echo "vllm-cpp: python3 with venv is required to build the MLX provider; pass VLLM_CPP_MLX=off to build Metal without it" >&2; exit 1; }; \ + fi + "$(MLX_VENV)"/bin/pip install --quiet --disable-pip-version-check "mlx==$(MLX_VERSION)" + @# Resolved in the SHELL, not by $(MLX_ROOT): make expands a whole recipe + @# before running its first line, so the glob would still be unmatched here. + @# Every later use (the cmake args, package.sh) expands after this target has + @# completed, where $(MLX_ROOT) does resolve. + @root=$$(echo "$(MLX_VENV)"/lib/python*/site-packages/mlx); \ + test -f "$$root/lib/libmlx.dylib" -a -f "$$root/include/mlx/array.h" || \ + { echo "vllm-cpp: mlx==$(MLX_VERSION) did not provide lib/libmlx.dylib + include/mlx/array.h under $$root" >&2; exit 1; } + touch $@ +else +MLX_STAMP= +MLX_CMAKE_ARGS= +endif + +# govllmcpp.go mirrors vllm.h by hand, and the only guard against the two +# drifting apart is the vllm_abi_version check inside registerLib - which fires +# at runtime, on the user's machine, taking down every model load (issue +# #11379). Compare the two here instead, so moving VLLM_CPP_VERSION past the +# mirrors turns the build red while the header is still around to diff. +abi-check: sources/vllm.cpp + @engine=$$(sed -n 's/^#define VLLM_ABI_VERSION \([0-9][0-9]*\).*/\1/p' sources/vllm.cpp/include/vllm.h); \ + backend=$$(sed -n 's/^const abiVersion = \([0-9][0-9]*\).*/\1/p' govllmcpp.go); \ + if [ -z "$$engine" ] || [ -z "$$backend" ]; then \ + echo "vllm-cpp: cannot read the ABI version (engine='$$engine' backend='$$backend')" >&2; exit 1; \ + fi; \ + if [ "$$engine" != "$$backend" ]; then \ + echo "vllm-cpp: ABI mismatch: vllm.cpp $(VLLM_CPP_VERSION) is v$$engine, govllmcpp.go mirrors v$$backend." >&2; \ + echo " Update the struct mirrors and abiVersion in govllmcpp.go (and the offsets in vllmcpp_test.go) to v$$engine." >&2; \ + exit 1; \ + fi; \ + echo "vllm-cpp: ABI v$$engine matches the pinned engine" + +$(LIB): sources/vllm.cpp $(MLX_STAMP) + $(MAKE) abi-check mkdir -p build && \ cd build && \ - cmake ../sources/vllm.cpp $(CMAKE_ARGS) && \ + cmake ../sources/vllm.cpp $(CMAKE_ARGS) $(MLX_CMAKE_ARGS) && \ cmake --build . --config Release -j$(JOBS) --target vllm_shared cp -fL build/$(LIB) ./$(LIB) @@ -79,16 +163,18 @@ vllm-cpp: main.go govllmcpp.go backend.go options.go $(LIB) CGO_ENABLED=0 $(GOCMD) build -tags "$(GO_TAGS)" -o vllm-cpp ./ package: vllm-cpp - bash package.sh + MLX_ROOT="$(MLX_ROOT)" bash package.sh build: package clean: purge - rm -rf libvllm.so libvllm.dylib package sources/vllm.cpp vllm-cpp + rm -rf libvllm.so libvllm.dylib package sources/vllm.cpp vllm-cpp "$(MLX_VENV)" purge: rm -rf build +.PHONY: abi-check + .NOTPARALLEL: # The unit specs are pure Go (struct mirrors, option mapping, load diff --git a/backend/go/vllm-cpp/README.md b/backend/go/vllm-cpp/README.md index a2a10070c..302119bfe 100644 --- a/backend/go/vllm-cpp/README.md +++ b/backend/go/vllm-cpp/README.md @@ -6,7 +6,7 @@ safetensors + GGUF loading, CUDA / CPU / Metal / Vulkan) with no Python at inference time. The backend dlopens the engine's stable C ABI (`libvllm`, `include/vllm.h`, -ABI v2) through purego: +ABI v10) through purego: - `Load` -> `vllm_engine_load`: accepts a `.gguf` file or a HF-style model directory (`config.json` + safetensors). `context_size` maps to @@ -29,6 +29,12 @@ ABI v2) through purego: LocalAI's Go-side grammar-constrained tool calling; JSON-schema / regex / choice constraints are also exposed by the ABI. +The struct mirrors in `govllmcpp.go` are hand-written against one ABI version, +and the engine refuses to load against any other. Moving `VLLM_CPP_VERSION` in +the Makefile therefore means updating `abiVersion` plus the mirrors (and their +offsets in `vllmcpp_test.go`) in the same change; `make abi-check` compares the +pinned header against the bindings and the library build runs it first. + ## Hardware coverage The CUDA builds require the CUDA 13 toolchain and target Blackwell only: @@ -64,5 +70,50 @@ options: - max_num_seqs:16 ``` +## Apple Silicon: the MLX GEMM provider (ON by default, gated to prefill) + +`BUILD_TYPE=metal` builds vllm.cpp's MLX provider for the dense GEMM +(`VLLM_CPP_MLX=on`, the default here). It is on because upstream now SHAPE-GATES +it to prefill; it was briefly off in this branch's history, and that was correct +at the time for an ungated provider. + +The gate matters more than the flag. MLX's steel GEMM wins prefill but loses +decode, because the provider pays an `mx::eval` synchronisation plus an output +memcpy on every call and decode makes ~112 calls *per token*. Measured on an +Apple M4, Qwen3-1.7B-bf16 warm at p=512 g=128: + +| configuration | prefill TTFT | warm throughput | +|---|--:|--:| +| MLX **gated to prefill** (pin >= 89c46aeb) | **524.5 ms** | **24.37 tok/s, 97.6% of MLX-LM** | +| MLX ungated (older pins) | 537 ms | 12.7 tok/s | +| MLX off | 602 ms | 23.9 tok/s, 95.9% | + +Ratios are against an MLX-LM baseline measured INTERLEAVED with ours over four +ABBA blocks (its spread 0.34%, ours 0.12%). An earlier revision of this file +claimed 99.1%; that used a two-run MLX-LM baseline containing an outlier and +overstated us by about 1.5 points. + +**`VLLM_CPP_VERSION` and this flag are coupled.** Moving the pin back before +`89c46aeb` while leaving `VLLM_CPP_MLX=on` would take the middle row — roughly +half throughput. If you roll the pin back, roll the default back with it. + +One caveat: MLX's GEMM is not bit-identical to the native kernel, so an MLX build +produces a different greedy sequence than a non-MLX one. That is a property of the +provider, not of the gate, and it predates this packaging. Full disposition in +vllm.cpp `docs/BENCHMARKS.md`. + +Build knobs: + +- `VLLM_CPP_MLX=off` builds Metal without the provider: ~124 MB smaller, and + 96.4% of MLX-LM instead of 99.1%. +- `MLX_VERSION` pins the wheel (default `0.29.4`). MLX is consumed as the + prebuilt pip wheel because building it from source needs `xcrun metal`, i.e. a + full Xcode the macOS runners do not have. + +Packaging vendors `libmlx.dylib`, `mlx.metallib` and MLX's MIT license into +`package/lib/`, and rewrites `libvllm.dylib`'s rpath to `@loader_path/lib` +(re-signing it, since `install_name_tool` invalidates the signature). The +metallib must stay beside `libmlx.dylib`: MLX looks for it there. + Testing: `make test` runs the unit specs; export `VLLM_CPP_MODEL=` (and optionally `VLLM_CPP_LIBRARY=`) to enable the e2e specs. diff --git a/backend/go/vllm-cpp/govllmcpp.go b/backend/go/vllm-cpp/govllmcpp.go index bd39e5428..5f525f873 100644 --- a/backend/go/vllm-cpp/govllmcpp.go +++ b/backend/go/vllm-cpp/govllmcpp.go @@ -1,6 +1,6 @@ package main -// purego bindings for the vllm.cpp stable C ABI (include/vllm.h, ABI v2). +// purego bindings for the vllm.cpp stable C ABI (include/vllm.h, ABI v10). // // The structs below are hand-mirrored PODs of the C declarations, with // explicit padding so the Go layout matches the C layout on linux/darwin @@ -17,15 +17,21 @@ import ( "github.com/ebitengine/purego" ) -// abiVersion is the VLLM_ABI_VERSION this file mirrors (vllm.h). -const abiVersion = 5 +// abiVersion is the VLLM_ABI_VERSION this file mirrors (vllm.h). It must track +// the header of the VLLM_CPP_VERSION pinned in the Makefile: the build checks +// the two against each other, because a mismatch is only caught at runtime by +// registerLib, where it takes the backend down on every load (issue #11379). +const abiVersion = 10 // vllm_status (vllm.h). const ( vllmOK = 0 ) -// cModelParams mirrors vllm_model_params. +// cModelParams mirrors vllm_model_params. The fields the backend does not set +// are still mirrored: the engine reads the whole struct, so the Go value must +// be the same size as the C one. Every one of them is inert when zeroed, which +// is what keeps the engine byte-identical to the pre-v6 behavior. type cModelParams struct { ModelPath uintptr // const char* TokenizerConfigPath uintptr // const char* @@ -35,11 +41,18 @@ type cModelParams struct { MaxNumSeqs int32 ToolParser uintptr // const char*; NULL = auto-detect (ABI v4) ReasoningParser uintptr // const char*; NULL = auto-detect (ABI v5) + SpeculativeConfig uintptr // const char*; NULL = no speculation (ABI v6) + EnablePrefixCaching int32 // 0 = model default, 1 = on, 2 = off (ABI v7) + MaxNumBatchedTokens int32 // <= 0 = per-arch default (ABI v9) + SchedulingPolicy uintptr // const char*; NULL = "fcfs" (ABI v9) + KVTransferConfig uintptr // const char*; NULL = no connector (ABI v9) + EnableJumpForward int32 // 0 = env-resolved (off), 1 = on, 2 = off (ABI v10) + _ [4]byte } -// cSamplingParams mirrors vllm_sampling_params (ABI v2, structured fields -// included). Padding matches the C compiler's: the uint64 seed is 8-aligned, -// and each pointer following an int32 is 8-aligned. +// cSamplingParams mirrors vllm_sampling_params (structured fields included). +// Padding matches the C compiler's: the uint64 seed is 8-aligned, and each +// pointer following an int32 is 8-aligned. type cSamplingParams struct { Temperature float32 TopP float32 @@ -65,6 +78,10 @@ type cSamplingParams struct { StructuredGrammar uintptr // const char* StructuredJSONObject int32 _ [4]byte + // Per-request custom logits processor (ABI v8). Left NULL: a Go callback + // would have to run inside the sampler's decode step for every token. + LogitsProcessor uintptr // vllm_logits_processor; NULL = none + LogitsProcessorUserData uintptr // void*, passed back to the callback } // cCompletion mirrors vllm_completion. diff --git a/backend/go/vllm-cpp/package.sh b/backend/go/vllm-cpp/package.sh index 78dc49178..30a21d219 100644 --- a/backend/go/vllm-cpp/package.sh +++ b/backend/go/vllm-cpp/package.sh @@ -43,6 +43,50 @@ elif [ -f "/lib/ld-linux-aarch64.so.1" ]; then cp -arfLv /lib/aarch64-linux-gnu/libpthread.so.0 $CURDIR/package/lib/libpthread.so.0 elif [ $(uname -s) = "Darwin" ]; then echo "Detected Darwin" + # Vendor the optional MLX GEMM provider, when libvllm was built against it. + # Three facts drive every line below, each verified on an Apple M4 before it + # was written: + # 1. libvllm.dylib carries an LC_LOAD_DYLIB on @rpath/libmlx.dylib, and its + # build-time LC_RPATH points inside the build venv. That path does not + # exist on a user's machine, so it must become @loader_path/lib. + # 2. MLX finds its ~100 MB mlx.metallib beside its OWN dylib, so the two + # files have to land in the same directory or every Metal op dies with + # "Failed to load the default metallib". + # 3. install_name_tool invalidates the code signature, and macOS refuses to + # load an arm64 image whose signature does not match, so the patched + # library must be re-signed ad-hoc afterwards. + if otool -L "$CURDIR/package/libvllm.dylib" 2>/dev/null | grep -q "libmlx.dylib"; then + MLX_LIB_DIR="${MLX_ROOT}/lib" + if [ ! -f "$MLX_LIB_DIR/libmlx.dylib" ] || [ ! -f "$MLX_LIB_DIR/mlx.metallib" ]; then + echo "Error: libvllm.dylib links libmlx.dylib but $MLX_LIB_DIR is missing libmlx.dylib/mlx.metallib" >&2 + exit 1 + fi + echo "Vendoring the MLX GEMM provider from $MLX_LIB_DIR" + cp -fLv "$MLX_LIB_DIR/libmlx.dylib" "$CURDIR/package/lib/" + cp -fLv "$MLX_LIB_DIR/mlx.metallib" "$CURDIR/package/lib/" + # MLX is MIT and we redistribute its binaries, so its license ships with + # them. mlx-metal is the wheel carrying the dylib and the metallib. + MLX_LICENSE=$(ls "${MLX_ROOT}"/../mlx_metal-*.dist-info/licenses/LICENSE 2>/dev/null | head -1) + if [ -z "$MLX_LICENSE" ]; then + MLX_LICENSE=$(ls "${MLX_ROOT}"/../mlx-*.dist-info/licenses/LICENSE 2>/dev/null | head -1) + fi + if [ -z "$MLX_LICENSE" ]; then + echo "Error: could not find the MLX LICENSE to redistribute alongside libmlx.dylib" >&2 + exit 1 + fi + cp -fLv "$MLX_LICENSE" "$CURDIR/package/lib/LICENSE.mlx" + # Drop every build-tree rpath, then point at the packaged copy. + otool -l "$CURDIR/package/libvllm.dylib" | awk '/LC_RPATH/{f=1;next} f&&/ path /{print $2;f=0}' | while read -r rp; do + install_name_tool -delete_rpath "$rp" "$CURDIR/package/libvllm.dylib" 2>/dev/null || true + done + install_name_tool -add_rpath "@loader_path/lib" "$CURDIR/package/libvllm.dylib" + codesign -f -s - "$CURDIR/package/libvllm.dylib" + # A broken rpath must fail the BUILD, not the user's first inference. + if ! otool -l "$CURDIR/package/libvllm.dylib" | grep -q "@loader_path/lib"; then + echo "Error: libvllm.dylib did not get the @loader_path/lib rpath" >&2 + exit 1 + fi + fi else echo "Error: Could not detect architecture" exit 1 diff --git a/backend/go/vllm-cpp/vllmcpp_test.go b/backend/go/vllm-cpp/vllmcpp_test.go index 297d15944..7ab6deedd 100644 --- a/backend/go/vllm-cpp/vllmcpp_test.go +++ b/backend/go/vllm-cpp/vllmcpp_test.go @@ -16,10 +16,17 @@ func TestVllmCpp(t *testing.T) { RunSpecs(t, "vllm-cpp suite") } -// The Go POD mirrors must match the C struct layout of vllm.h (ABI v2) +// The Go POD mirrors must match the C struct layout of vllm.h (ABI v10) // byte-for-byte: these offsets are the C offsets on LP64 (linux/darwin // amd64+arm64). A failure here means govllmcpp.go drifted from vllm.h. var _ = Describe("C ABI struct mirrors", func() { + It("declares the ABI version the pinned engine reports", func() { + // VLLM_ABI_VERSION in the vllm.h of VLLM_CPP_VERSION (Makefile). + // Moving the pin past this without growing the mirrors below ships a + // backend that refuses every load at startup (issue #11379). + Expect(abiVersion).To(Equal(10)) + }) + It("cModelParams matches vllm_model_params", func() { var p cModelParams Expect(unsafe.Offsetof(p.ModelPath)).To(Equal(uintptr(0))) @@ -30,10 +37,16 @@ var _ = Describe("C ABI struct mirrors", func() { Expect(unsafe.Offsetof(p.MaxNumSeqs)).To(Equal(uintptr(28))) Expect(unsafe.Offsetof(p.ToolParser)).To(Equal(uintptr(32))) Expect(unsafe.Offsetof(p.ReasoningParser)).To(Equal(uintptr(40))) - Expect(unsafe.Sizeof(p)).To(Equal(uintptr(48))) + Expect(unsafe.Offsetof(p.SpeculativeConfig)).To(Equal(uintptr(48))) + Expect(unsafe.Offsetof(p.EnablePrefixCaching)).To(Equal(uintptr(56))) + Expect(unsafe.Offsetof(p.MaxNumBatchedTokens)).To(Equal(uintptr(60))) + Expect(unsafe.Offsetof(p.SchedulingPolicy)).To(Equal(uintptr(64))) + Expect(unsafe.Offsetof(p.KVTransferConfig)).To(Equal(uintptr(72))) + Expect(unsafe.Offsetof(p.EnableJumpForward)).To(Equal(uintptr(80))) + Expect(unsafe.Sizeof(p)).To(Equal(uintptr(88))) }) - It("cSamplingParams matches vllm_sampling_params (ABI v2)", func() { + It("cSamplingParams matches vllm_sampling_params", func() { var p cSamplingParams Expect(unsafe.Offsetof(p.Temperature)).To(Equal(uintptr(0))) Expect(unsafe.Offsetof(p.TopP)).To(Equal(uintptr(4))) @@ -55,7 +68,9 @@ var _ = Describe("C ABI struct mirrors", func() { Expect(unsafe.Offsetof(p.NStructuredChoice)).To(Equal(uintptr(96))) Expect(unsafe.Offsetof(p.StructuredGrammar)).To(Equal(uintptr(104))) Expect(unsafe.Offsetof(p.StructuredJSONObject)).To(Equal(uintptr(112))) - Expect(unsafe.Sizeof(p)).To(Equal(uintptr(120))) + Expect(unsafe.Offsetof(p.LogitsProcessor)).To(Equal(uintptr(120))) + Expect(unsafe.Offsetof(p.LogitsProcessorUserData)).To(Equal(uintptr(128))) + Expect(unsafe.Sizeof(p)).To(Equal(uintptr(136))) }) It("cCompletion matches vllm_completion", func() { diff --git a/backend/go/voice-detect/package.sh b/backend/go/voice-detect/package.sh index de95c8ce2..ff8776239 100755 --- a/backend/go/voice-detect/package.sh +++ b/backend/go/voice-detect/package.sh @@ -25,34 +25,7 @@ cp -avf "$CURDIR"/libvoicedetect.so* "$CURDIR/package/lib/" 2>/dev/null || { # Detect architecture and copy the core runtime libs libvoicedetect.so links # against, plus the matching dynamic loader as lib/ld.so. -if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then - 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 - 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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries (CUDA/ROCm/Intel/Vulkan loader + ICDs + drivers) based on # BUILD_TYPE so the backend can reach the GPU without the runtime base image diff --git a/backend/go/voxtral/package.sh b/backend/go/voxtral/package.sh index 8465a36da..6cc94fc40 100644 --- a/backend/go/voxtral/package.sh +++ b/backend/go/voxtral/package.sh @@ -16,43 +16,13 @@ cp -fv $CURDIR/libgovoxtral-*.dylib $CURDIR/package/ 2>/dev/null || true 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 - # OpenBLAS if available - if [ -f /usr/lib/x86_64-linux-gnu/libopenblas.so.0 ]; then - cp -arfLv /usr/lib/x86_64-linux-gnu/libopenblas.so.0 $CURDIR/package/lib/ - fi -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 - # OpenBLAS if available - if [ -f /usr/lib/aarch64-linux-gnu/libopenblas.so.0 ]; then - cp -arfLv /usr/lib/aarch64-linux-gnu/libopenblas.so.0 $CURDIR/package/lib/ - fi -elif [ $(uname -s) = "Darwin" ]; then - echo "Detected Darwin — system frameworks linked dynamically, no bundled libs needed" -else - echo "Error: Could not detect architecture" - exit 1 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" + +# OpenBLAS if available +if [ -f /usr/lib/x86_64-linux-gnu/libopenblas.so.0 ]; then + cp -arfLv /usr/lib/x86_64-linux-gnu/libopenblas.so.0 $CURDIR/package/lib/ +elif [ -f /usr/lib/aarch64-linux-gnu/libopenblas.so.0 ]; then + cp -arfLv /usr/lib/aarch64-linux-gnu/libopenblas.so.0 $CURDIR/package/lib/ fi # Package GPU libraries based on BUILD_TYPE diff --git a/backend/go/whisper/Makefile b/backend/go/whisper/Makefile index a042c7050..e760bea02 100644 --- a/backend/go/whisper/Makefile +++ b/backend/go/whisper/Makefile @@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1) # whisper.cpp version WHISPER_REPO?=https://github.com/ggml-org/whisper.cpp -WHISPER_CPP_VERSION?=080bbbe85230f624f0b52127f1ae1218247989f9 +WHISPER_CPP_VERSION?=306c88f4d1286aec1bf96e544632897886af5501 SO_TARGET?=libgowhisper.so CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF diff --git a/backend/go/whisper/package.sh b/backend/go/whisper/package.sh index efeaa7009..a32d577ae 100755 --- a/backend/go/whisper/package.sh +++ b/backend/go/whisper/package.sh @@ -17,40 +17,7 @@ cp -fv $CURDIR/libgowhisper-*.dylib $CURDIR/package/ 2>/dev/null || true 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/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/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/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/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 +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" # Package GPU libraries based on BUILD_TYPE # The GPU library packaging script will detect BUILD_TYPE and copy appropriate GPU libraries diff --git a/backend/index.yaml b/backend/index.yaml index b6026a6e6..93ff6ebfe 100644 --- a/backend/index.yaml +++ b/backend/index.yaml @@ -129,6 +129,38 @@ nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-ds4" metal: "metal-ds4" metal-darwin-arm64: "metal-ds4" +- &audiocpp + name: "audio-cpp" + alias: "audio-cpp" + license: apache-2.0 + description: | + 0xShug0/audio.cpp - a ggml audio inference framework covering text to + speech, voice cloning, transcription, forced alignment, voice activity + detection, speaker diarization, source separation, voice conversion and + music generation, across 30+ model families. Consumes audio.cpp-native + GGUF packages from huggingface.co/audio-cpp/audio.cpp-gguf. + urls: + - https://github.com/0xShug0/audio.cpp + - https://huggingface.co/audio-cpp/audio.cpp-gguf + tags: + - TTS + - text-to-speech + - audio-transcription + - CPU + - CUDA + - Metal + # No vulkan key: the vulkan image would carry a Vulkan loader with no Mesa ICD + # (see the audio-cpp block in .github/backend-matrix.yml). Pointing a + # capability at a tag CI never builds hands the user a pull failure; leaving it + # out lets SystemState.Capability() fall through to `default` and the CPU + # image. + capabilities: + default: "cpu-audio-cpp" + nvidia: "cuda12-audio-cpp" + nvidia-cuda-12: "cuda12-audio-cpp" + nvidia-cuda-13: "cuda13-audio-cpp" + metal: "metal-audio-cpp" + metal-darwin-arm64: "metal-audio-cpp" - &whispercpp name: "whisper" alias: "whisper" @@ -161,16 +193,27 @@ alias: "vllm-cpp" license: apache-2.0 description: | - vllm.cpp is a from-scratch C++20 port of vLLM created and maintained by the LocalAI team. - It mirrors vLLM's V1 architecture (paged KV cache, continuous batching, prefix caching, - scheduler, sampler) on a portable tensor runtime with no Python, PyTorch or ggml at - 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. + ALPHA development builds. Try it, but llama-cpp stays the recommendation for + production use. + + vllm.cpp is an Apache-2.0 C++20 inference engine maintained by the LocalAI team, + developed in its own repository and usable without LocalAI. It began as a port of + vLLM and keeps vLLM as its reference implementation, checking output against it and + benchmarking against it, while growing a featureset of its own. It implements vLLM's + V1 architecture (paged KV cache, continuous batching, prefix caching, scheduler, + sampler) on a portable tensor runtime with no Python, PyTorch or ggml at inference + time. It loads GGUF as well as Hugging Face safetensors, supports structured output + (JSON schema / regex / choice / GBNF grammar) enforced in-engine, ships speculative + decoding and KV offload, 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. + + The project is expected to be renamed as it diverges further from vLLM; the new + name is still to be decided. urls: - https://github.com/mudler/vllm.cpp tags: @@ -437,6 +480,32 @@ 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" @@ -1467,6 +1536,7 @@ alias: "kokoro" name: "kokoro" capabilities: + default: "cpu-kokoro" nvidia: "cuda12-kokoro" intel: "intel-kokoro" amd: "rocm-kokoro" @@ -1848,6 +1918,23 @@ 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: @@ -1969,6 +2056,15 @@ nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-ds4-development" metal: "metal-ds4-development" metal-darwin-arm64: "metal-ds4-development" +- !!merge <<: *audiocpp + name: "audio-cpp-development" + capabilities: + default: "cpu-audio-cpp-development" + nvidia: "cuda12-audio-cpp-development" + nvidia-cuda-12: "cuda12-audio-cpp-development" + nvidia-cuda-13: "cuda13-audio-cpp-development" + metal: "metal-audio-cpp-development" + metal-darwin-arm64: "metal-audio-cpp-development" - !!merge <<: *stablediffusionggml name: "stablediffusion-ggml-development" capabilities: @@ -1983,6 +2079,18 @@ 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" @@ -2380,6 +2488,35 @@ 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" @@ -3699,6 +3836,77 @@ 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" @@ -5198,11 +5406,22 @@ - !!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" @@ -6673,3 +6892,44 @@ uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-vllm-cpp" mirrors: - localai/localai-backends:master-metal-darwin-arm64-vllm-cpp +## audio-cpp +- !!merge <<: *audiocpp + name: "cpu-audio-cpp" + uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-audio-cpp" + mirrors: + - localai/localai-backends:latest-cpu-audio-cpp +- !!merge <<: *audiocpp + name: "cpu-audio-cpp-development" + uri: "quay.io/go-skynet/local-ai-backends:master-cpu-audio-cpp" + mirrors: + - localai/localai-backends:master-cpu-audio-cpp +- !!merge <<: *audiocpp + name: "cuda12-audio-cpp" + uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-12-audio-cpp" + mirrors: + - localai/localai-backends:latest-gpu-nvidia-cuda-12-audio-cpp +- !!merge <<: *audiocpp + name: "cuda12-audio-cpp-development" + uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-12-audio-cpp" + mirrors: + - localai/localai-backends:master-gpu-nvidia-cuda-12-audio-cpp +- !!merge <<: *audiocpp + name: "cuda13-audio-cpp" + uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-13-audio-cpp" + mirrors: + - localai/localai-backends:latest-gpu-nvidia-cuda-13-audio-cpp +- !!merge <<: *audiocpp + name: "cuda13-audio-cpp-development" + uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-13-audio-cpp" + mirrors: + - localai/localai-backends:master-gpu-nvidia-cuda-13-audio-cpp +- !!merge <<: *audiocpp + name: "metal-audio-cpp" + uri: "quay.io/go-skynet/local-ai-backends:latest-metal-darwin-arm64-audio-cpp" + mirrors: + - localai/localai-backends:latest-metal-darwin-arm64-audio-cpp +- !!merge <<: *audiocpp + name: "metal-audio-cpp-development" + uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-audio-cpp" + mirrors: + - localai/localai-backends:master-metal-darwin-arm64-audio-cpp diff --git a/backend/python/chatterbox/requirements-cublas12.txt b/backend/python/chatterbox/requirements-cublas12.txt index 2467c9162..1ca1ad8d6 100644 --- a/backend/python/chatterbox/requirements-cublas12.txt +++ b/backend/python/chatterbox/requirements-cublas12.txt @@ -1,6 +1,7 @@ -torch -torchaudio -transformers +--extra-index-url https://download.pytorch.org/whl/cu124 +torch==2.6.0+cu124 +torchaudio==2.6.0+cu124 +transformers<5 numpy>=1.24.0,<1.26.0 # chatterbox-tts itself is installed with --no-deps in install.sh. # These are its real runtime deps, mirroring upstream's pyproject.toml @@ -15,4 +16,4 @@ conformer safetensors spacy-pkuseg pykakasi==2.3.0 -accelerate +accelerate \ No newline at end of file diff --git a/backend/python/chatterbox/requirements.txt b/backend/python/chatterbox/requirements.txt index 55a0867f0..07d7d0706 100644 --- a/backend/python/chatterbox/requirements.txt +++ b/backend/python/chatterbox/requirements.txt @@ -2,5 +2,5 @@ grpcio==1.71.0 protobuf certifi packaging -setuptools +setuptools<81 poetry \ No newline at end of file diff --git a/backend/python/common/vllm_utils.py b/backend/python/common/vllm_utils.py index 9124645ac..60c38f5bd 100644 --- a/backend/python/common/vllm_utils.py +++ b/backend/python/common/vllm_utils.py @@ -4,11 +4,208 @@ Generic helpers (``parse_options``, ``messages_to_dicts``) live in ``python_utils`` and are re-exported here for backwards compatibility with existing imports in both backends. """ +import dataclasses +import difflib +import json import sys from python_utils import messages_to_dicts, parse_options -__all__ = ["parse_options", "messages_to_dicts", "setup_parsers"] +__all__ = [ + "parse_options", + "messages_to_dicts", + "setup_parsers", + "apply_options_to_engine_args", + "normalize_option_key", +] + +# Options[] entries LocalAI itself acts on, in their normalized spelling. +_BACKEND_LEVEL_OPTIONS = {"tool_parser", "reasoning_parser"} + +_TRUTHY = {"true", "1", "yes", "on", "t", "y"} +_FALSY = {"false", "0", "no", "off", "f", "n"} + + +def normalize_option_key(key): + """Normalize an Options[] key to its engine/field spelling. + + Users copy flags straight out of vLLM's CLI docs (``--reasoning-parser``), + so accept that spelling everywhere LocalAI looks options up by name. + """ + return key.strip().lstrip("-").replace("-", "_") + + +_BASE_HINTS = { + "bool": "bool", + "int": "int", + "float": "float", + "str": "str", + "dict": "dict", + "mapping": "dict", +} + + +def _hint_from_annotation(annotation): + """Target type name for a field annotation, or None when it is not a scalar. + + ``dataclasses.fields()`` hands back either a real type, a ``typing`` + construct or - under PEP 563, which vLLM's arg_utils uses - the annotation + as a plain string, so work on the textual form. Match the *base* of the + annotation rather than searching it: ``Literal["auto", "float16"]`` (vLLM's + dtype) contains "float" but is not a float. + """ + if annotation is None: + return None + if isinstance(annotation, type): + text = annotation.__name__ + elif isinstance(annotation, str): + text = annotation + else: + text = str(annotation) + text = text.replace("typing.", "").strip() + while text.lower().startswith("optional[") and text.endswith("]"): + text = text[len("optional["):-1].strip() + if "|" in text: + parts = [p.strip() for p in text.split("|") if p.strip().lower() not in ("none", "nonetype")] + if len(parts) != 1: + return None + text = parts[0] + base = text.split("[", 1)[0].strip().lower() + return _BASE_HINTS.get(base) + + +def _hint_from_value(current): + """Target type name inferred from a field's current value.""" + if isinstance(current, bool): + return "bool" + if isinstance(current, dict): + return "dict" + if isinstance(current, float): + return "float" + if isinstance(current, int): + return "int" + if isinstance(current, str): + return "str" + return None + + +def _type_hint(annotation, current): + """Best-effort target type name for a dataclass field.""" + hint = _hint_from_annotation(annotation) + if hint is None: + hint = _hint_from_value(current) + return hint + + +def _coerce_option(raw, hint): + """Coerce a CLI-supplied string to the field's type. Raises ValueError.""" + if hint == "bool": + low = raw.lower() + if low in _TRUTHY: + return True + if low in _FALSY: + return False + raise ValueError(f"{raw!r} is not a boolean") + if hint == "int": + return int(raw) + if hint == "float": + return float(raw) + if hint == "dict": + return json.loads(raw) + if hint == "str": + return raw + + # Untyped (or union-typed) field: infer from the literal itself. + low = raw.lower() + if low in _TRUTHY: + return True + if low in _FALSY: + return False + for cast in (int, float): + try: + return cast(raw) + except ValueError: + pass + if raw.startswith("{") or raw.startswith("["): + return json.loads(raw) + return raw + + +def apply_options_to_engine_args(engine_args, options): + """Apply CLI-style ``--flag[:value]`` entries from Options[] to engine args. + + ``options:`` is a loose bag shared with backend-level settings + (``tool_parser:``, ``reasoning_parser:``, ``vad_only``, …), so only + ``--`` prefixed entries are treated as engine flags. Names are normalized + the way vLLM's own CLI does (``--enable-prefix-caching`` → + ``enable_prefix_caching``) and values are coerced to the target field's + type. Both ``--flag:value`` (LocalAI convention) and ``--flag=value`` + (vLLM CLI convention) are accepted; a bare ``--flag`` sets a boolean field. + + Unknown or uncoercible flags warn on stderr and are skipped instead of + failing the load - unlike ``engine_args:``, which is engine-only and + therefore strict, Options[] carries entries this function knows nothing + about. + + Returns a new dataclass instance via ``dataclasses.replace`` so the + engine's ``__post_init__`` re-runs, or the original when nothing applied. + """ + if not options: + return engine_args + + fields = {f.name: f for f in dataclasses.fields(type(engine_args))} + updates = {} + for opt in options: + opt = opt.strip() + if not opt.startswith("--"): + continue + body = opt[2:] + seps = [i for i in (body.find(":"), body.find("=")) if i != -1] + if seps: + cut = min(seps) + name, raw = body[:cut], body[cut + 1:].strip() + else: + name, raw = body, None + name = normalize_option_key(name) + + if name not in fields: + # LocalAI reads these itself; whether the engine dataclass carries + # them too varies by vLLM version, so never flag them as unknown. + if name in _BACKEND_LEVEL_OPTIONS: + continue + suggestion = difflib.get_close_matches(name, fields, n=1) + hint = f" did you mean {suggestion[0]!r}?" if suggestion else "" + print( + f"[vllm_utils] unknown engine option {opt!r} (field {name!r}), skipping.{hint}", + file=sys.stderr, + ) + continue + + target = _type_hint(fields[name].type, getattr(engine_args, name, None)) + if raw is None: + if target != "bool": + print( + f"[vllm_utils] engine option {opt!r} needs a value " + f"(field {name!r} is not a flag), skipping", + file=sys.stderr, + ) + continue + updates[name] = True + continue + + try: + updates[name] = _coerce_option(raw, target) + except (ValueError, TypeError) as err: + print( + f"[vllm_utils] cannot apply engine option {opt!r} to field " + f"{name!r}: {err}, skipping", + file=sys.stderr, + ) + + if not updates: + return engine_args + print(f"[vllm_utils] engine options from Options[]: {updates}", file=sys.stderr) + return dataclasses.replace(engine_args, **updates) def setup_parsers(opts): diff --git a/backend/python/common/vllm_utils_test.py b/backend/python/common/vllm_utils_test.py new file mode 100644 index 000000000..b3bbcbb5b --- /dev/null +++ b/backend/python/common/vllm_utils_test.py @@ -0,0 +1,160 @@ +"""Unit tests for the shared vLLM backend helpers (vllm_utils.py). + +Run standalone (Python standard library only, no backend venv needed): + cd backend/python/common && python3 -m unittest vllm_utils_test + +``vllm_utils`` imports vLLM lazily (inside functions), so the module is +importable without the vLLM wheel. ``AsyncEngineArgs`` is stood in for by a +local dataclass that mirrors the field shapes that matter: plain scalars, +``Optional[...]`` scalars (vLLM's tri-state flags) and dict-valued fields. +""" + +import contextlib +import dataclasses +import io +import unittest +from typing import Dict, Literal, Optional, Union + +from vllm_utils import apply_options_to_engine_args, normalize_option_key + + +@dataclasses.dataclass +class FakeEngineArgs: + model: str = "" + quantization: Optional[str] = None + kv_cache_dtype: str = "auto" + enable_prefix_caching: Optional[bool] = None + enforce_eager: bool = False + max_model_len: Optional[int] = None + tensor_parallel_size: int = 1 + gpu_memory_utilization: float = 0.9 + limit_mm_per_prompt: Optional[Dict[str, int]] = None + # vLLM types dtype as a Literal of strings, several of which contain the + # word "float" - a naive match reads that as a float field. + dtype: Literal["auto", "float16", "bfloat16"] = "auto" + seed: Union[int, str, None] = None + + +# vLLM's arg_utils is annotated under PEP 563, where dataclasses.fields() hands +# back annotations as plain strings rather than types. +StringAnnotatedEngineArgs = dataclasses.make_dataclass( + "StringAnnotatedEngineArgs", + [ + ("max_model_len", "int | None", dataclasses.field(default=None)), + ("enable_prefix_caching", "Optional[bool]", dataclasses.field(default=None)), + ("kv_cache_dtype", "str", dataclasses.field(default="auto")), + ], +) + + +def _apply(options, **overrides): + """Apply options to a fresh FakeEngineArgs, returning (result, stderr).""" + err = io.StringIO() + with contextlib.redirect_stderr(err): + out = apply_options_to_engine_args(FakeEngineArgs(**overrides), options) + return out, err.getvalue() + + +class TestApplyOptionsToEngineArgs(unittest.TestCase): + def test_string_option_is_applied(self): + out, _ = _apply(["--quantization:gptq_marlin"]) + self.assertEqual(out.quantization, "gptq_marlin") + + def test_valueless_flag_enables_boolean_field(self): + out, _ = _apply(["--enable-prefix-caching"]) + self.assertIs(out.enable_prefix_caching, True) + + def test_boolean_field_accepts_explicit_false(self): + out, _ = _apply(["--enable-prefix-caching:false"], enable_prefix_caching=True) + self.assertIs(out.enable_prefix_caching, False) + + def test_integer_field_is_coerced(self): + out, _ = _apply(["--max-model-len:4096", "--tensor-parallel-size:2"]) + self.assertEqual(out.max_model_len, 4096) + self.assertEqual(out.tensor_parallel_size, 2) + + def test_float_field_is_coerced(self): + out, _ = _apply(["--gpu-memory-utilization:0.85"]) + self.assertAlmostEqual(out.gpu_memory_utilization, 0.85) + + def test_equals_separator_is_accepted(self): + out, _ = _apply(["--kv-cache-dtype=fp8_e5m2"]) + self.assertEqual(out.kv_cache_dtype, "fp8_e5m2") + + def test_dict_field_is_parsed_as_json(self): + out, _ = _apply(['--limit-mm-per-prompt:{"image": 4}']) + self.assertEqual(out.limit_mm_per_prompt, {"image": 4}) + + def test_non_flag_options_are_left_alone(self): + out, err = _apply(["tool_parser:hermes", "reasoning_parser:qwen3", "vad_only"]) + self.assertEqual(out, FakeEngineArgs()) + self.assertEqual(err, "") + + def test_unknown_flag_warns_and_is_skipped(self): + out, err = _apply(["--not-a-real-flag:1", "--quantization:awq"]) + self.assertEqual(out.quantization, "awq") + self.assertIn("not_a_real_flag", err) + self.assertIn("unknown", err.lower()) + + def test_parser_flags_are_not_reported_as_unknown(self): + # LocalAI consumes these itself; whether the engine dataclass also has + # the field depends on the vLLM version, and warning about them would + # send users chasing a non-problem. + out, err = _apply(["--tool-parser:hermes", "--reasoning-parser:qwen3"]) + self.assertEqual(out, FakeEngineArgs()) + self.assertEqual(err, "") + + def test_unknown_flag_hints_at_the_closest_field(self): + _, err = _apply(["--max-model-length:4096"]) + self.assertIn("max_model_len", err) + + def test_uncoercible_value_warns_and_is_skipped(self): + out, err = _apply(["--max-model-len:lots"]) + self.assertIsNone(out.max_model_len) + self.assertIn("max_model_len", err) + + def test_valueless_flag_on_non_boolean_field_warns_and_is_skipped(self): + out, err = _apply(["--quantization"]) + self.assertIsNone(out.quantization) + self.assertIn("quantization", err) + + def test_empty_options_returns_the_same_engine_args(self): + original = FakeEngineArgs(model="m") + self.assertIs(apply_options_to_engine_args(original, []), original) + + +class TestFieldTypeInference(unittest.TestCase): + def test_literal_typed_field_keeps_its_string_value(self): + out, err = _apply(["--dtype:bfloat16"]) + self.assertEqual(out.dtype, "bfloat16") + self.assertEqual(err.count("skipping"), 0) + + def test_ambiguous_union_falls_back_to_literal_inference(self): + out, _ = _apply(["--seed:42"]) + self.assertEqual(out.seed, 42) + + def test_string_annotations_are_understood(self): + err = io.StringIO() + with contextlib.redirect_stderr(err): + out = apply_options_to_engine_args( + StringAnnotatedEngineArgs(), + ["--max-model-len:4096", "--enable-prefix-caching", "--kv-cache-dtype:fp8_e5m2"], + ) + self.assertEqual(out.max_model_len, 4096) + self.assertIs(out.enable_prefix_caching, True) + self.assertEqual(out.kv_cache_dtype, "fp8_e5m2") + + +class TestNormalizeOptionKey(unittest.TestCase): + def test_cli_flag_becomes_a_field_name(self): + self.assertEqual(normalize_option_key("--reasoning-parser"), "reasoning_parser") + + def test_plain_key_is_untouched(self): + self.assertEqual(normalize_option_key("tool_parser"), "tool_parser") + + def test_surrounding_whitespace_is_stripped(self): + self.assertEqual(normalize_option_key(" --tool-parser "), "tool_parser") + + +if __name__ == "__main__": + unittest.main() diff --git a/backend/python/diffusers/backend.py b/backend/python/diffusers/backend.py index cb43c88c0..539ce5444 100755 --- a/backend/python/diffusers/backend.py +++ b/backend/python/diffusers/backend.py @@ -883,6 +883,34 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): return backend_pb2.Result(message="Media generated", success=True) + def UpscaleImage(self, request, context): + try: + if not request.src: + return backend_pb2.Result(success=False, message="No source image provided") + if not request.dst: + return backend_pb2.Result(success=False, message="No destination path provided") + + scale = request.scale if request.scale > 0 else 2 + image = Image.open(request.src).convert("RGB") + + # If the loaded pipeline supports upscaling (e.g. StableDiffusionUpscalePipeline), + # use it; otherwise fall back to high-quality Lanczos resize. + if self.pipe is not None and self.PipelineType in ("StableDiffusionUpscalePipeline", "StableDiffusionLatentUpscalePipeline"): + print(f"UpscaleImage: using diffusers upscale pipeline ({self.PipelineType})", file=sys.stderr) + upscaled = self.pipe(prompt="", image=image).images[0] + else: + # Fallback: high-quality Lanczos resize + print(f"UpscaleImage: no upscale pipeline loaded, using Lanczos resize (scale={scale})", file=sys.stderr) + new_w = image.width * scale + new_h = image.height * scale + upscaled = image.resize((new_w, new_h), Image.LANCZOS) + + upscaled.save(request.dst) + return backend_pb2.Result(message="Image upscaled", success=True) + except Exception as e: + print(f"UpscaleImage error: {e}", file=sys.stderr) + return backend_pb2.Result(success=False, message=str(e)) + def GenerateVideo(self, request, context): try: prompt = request.prompt diff --git a/backend/python/mlx-vlm/requirements-mps.txt b/backend/python/mlx-vlm/requirements-mps.txt index 882c2af5a..ee5351beb 100644 --- a/backend/python/mlx-vlm/requirements-mps.txt +++ b/backend/python/mlx-vlm/requirements-mps.txt @@ -1 +1,3 @@ -git+https://github.com/Blaizzy/mlx-vlm@v0.4.4 \ No newline at end of file +git+https://github.com/Blaizzy/mlx-vlm@v0.4.4 +torch +torchvision diff --git a/backend/python/nemo/install.sh b/backend/python/nemo/install.sh index bb7faeef4..a30fd27c8 100755 --- a/backend/python/nemo/install.sh +++ b/backend/python/nemo/install.sh @@ -14,4 +14,28 @@ 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 diff --git a/backend/python/sglang/requirements-cublas12-after.txt b/backend/python/sglang/requirements-cublas12-after.txt index 435075ecb..4c201fd7d 100644 --- a/backend/python/sglang/requirements-cublas12-after.txt +++ b/backend/python/sglang/requirements-cublas12-after.txt @@ -2,3 +2,25 @@ # (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 + +# Same failure mode as the nvidia-modelopt bound above, via a different +# package. sglang -> flashinfer-python -> cuda-tile, unbounded, and the +# global --prerelease=allow resolves it to 1.6.0rc3, whose build backend +# imports wheel_stub without declaring it in build-system.requires. With +# --no-build-isolation nothing installs it and the build dies with +# "No module named 'wheel_stub'". 1.5.0 is the newest stable release. +# Raise the bound once 1.6.0 final ships. +cuda-tile<1.6 diff --git a/backend/python/sglang/requirements-cublas13-after.txt b/backend/python/sglang/requirements-cublas13-after.txt index 435075ecb..4c201fd7d 100644 --- a/backend/python/sglang/requirements-cublas13-after.txt +++ b/backend/python/sglang/requirements-cublas13-after.txt @@ -2,3 +2,25 @@ # (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 + +# Same failure mode as the nvidia-modelopt bound above, via a different +# package. sglang -> flashinfer-python -> cuda-tile, unbounded, and the +# global --prerelease=allow resolves it to 1.6.0rc3, whose build backend +# imports wheel_stub without declaring it in build-system.requires. With +# --no-build-isolation nothing installs it and the build dies with +# "No module named 'wheel_stub'". 1.5.0 is the newest stable release. +# Raise the bound once 1.6.0 final ships. +cuda-tile<1.6 diff --git a/backend/python/sglang/requirements-l4t13-after.txt b/backend/python/sglang/requirements-l4t13-after.txt index fc2ca2030..d405b4ab0 100644 --- a/backend/python/sglang/requirements-l4t13-after.txt +++ b/backend/python/sglang/requirements-l4t13-after.txt @@ -13,3 +13,12 @@ # FunctionCallParser, ReasoningParser); the [all] extras are optional # accelerators not required at import time. sglang>=0.5.11 + +# Same failure mode the cublas profiles carry an nvidia-modelopt bound for, +# reached through a different package. sglang -> flashinfer-python -> +# cuda-tile, unbounded, and the global --prerelease=allow resolves it to +# 1.6.0rc3, whose build backend imports wheel_stub without declaring it in +# build-system.requires. With --no-build-isolation nothing installs it and +# the build dies with "No module named 'wheel_stub'". 1.5.0 is the newest +# stable release. Raise the bound once 1.6.0 final ships. +cuda-tile<1.6 diff --git a/backend/python/vllm/backend.py b/backend/python/vllm/backend.py index 8a42e7bcb..a5a85041b 100644 --- a/backend/python/vllm/backend.py +++ b/backend/python/vllm/backend.py @@ -22,6 +22,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common')) sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common')) from grpc_auth import get_auth_interceptors from model_utils import resolve_model_reference +from vllm_utils import apply_options_to_engine_args, normalize_option_key from vllm.engine.arg_utils import AsyncEngineArgs from vllm.engine.async_llm_engine import AsyncLLMEngine @@ -101,13 +102,18 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): return decoded_text def _parse_options(self, options_list): - """Parse Options[] key:value string list into a dict.""" + """Parse Options[] key:value string list into a dict. + + Keys are normalized to their field spelling so the CLI form users copy + from vLLM's docs (``--reasoning-parser:qwen3``) selects the same parser + as LocalAI's own (``reasoning_parser:qwen3``). + """ opts = {} for opt in options_list: if ":" not in opt: continue key, value = opt.split(":", 1) - opts[key.strip()] = value.strip() + opts[normalize_option_key(key)] = value.strip() return opts def _apply_engine_args(self, engine_args, engine_args_json): @@ -229,6 +235,12 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): "audio": max(request.LimitAudioPerPrompt, 1) } + # CLI-style flags in options: (--quantization:gptq_marlin, + # --enable-prefix-caching, ...) land on the engine args too - they must + # be applied *before* the engine is created or they do nothing. + # engine_args: is applied after this, so it stays the last word. + engine_args = apply_options_to_engine_args(engine_args, request.Options) + # engine_args from YAML overrides typed fields above so operators can # tune anything the AsyncEngineArgs dataclass exposes without waiting # on protobuf changes. diff --git a/backend/rust/kokoros/package.sh b/backend/rust/kokoros/package.sh index 80e06740d..86668796e 100644 --- a/backend/rust/kokoros/package.sh +++ b/backend/rust/kokoros/package.sh @@ -31,11 +31,7 @@ if [ -d "/etc/ssl/certs" ]; then fi # Copy the dynamic linker -if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then - cp -arfLv /lib64/ld-linux-x86-64.so.2 $CURDIR/package/lib/ld.so -elif [ -f "/lib/ld-linux-aarch64.so.1" ]; then - cp -arfLv /lib/ld-linux-aarch64.so.1 $CURDIR/package/lib/ld.so -fi +source "$CURDIR/../../../scripts/build/package-system-libs.sh" "$CURDIR/package/lib" "" echo "Packaging completed successfully" ls -liah $CURDIR/package/ diff --git a/backend/rust/kokoros/src/service.rs b/backend/rust/kokoros/src/service.rs index 495b0a423..f2fbe6cc4 100644 --- a/backend/rust/kokoros/src/service.rs +++ b/backend/rust/kokoros/src/service.rs @@ -334,6 +334,13 @@ impl Backend for KokorosService { Err(Status::unimplemented("Not supported")) } + async fn generate3_d( + &self, + _: Request, + ) -> Result, Status> { + Err(Status::unimplemented("Not supported")) + } + async fn audio_transcription( &self, _: Request, @@ -408,6 +415,13 @@ impl Backend for KokorosService { Err(Status::unimplemented("Not supported")) } + async fn detokenize( + &self, + _: Request, + ) -> Result, Status> { + Err(Status::unimplemented("Not supported")) + } + async fn detect( &self, _: Request, diff --git a/cmd/local-ai/main.go b/cmd/local-ai/main.go index fca6b2638..b1799dbb8 100644 --- a/cmd/local-ai/main.go +++ b/cmd/local-ai/main.go @@ -1,6 +1,7 @@ package main import ( + "errors" "os" "path/filepath" @@ -107,6 +108,13 @@ For documentation and support: // Run the thing! err = ctx.Run(&cli.CLI.Context) if err != nil { + // A command that has already told the user what went wrong returns + // only a status. Logging it as well would print a bare "exit status 1" + // underneath the explanation they just read. + var reported cli.ExitCodeError + if errors.As(err, &reported) { + os.Exit(reported.Code) + } xlog.Fatal("Error running the application", "error", err) } } diff --git a/core/application/application.go b/core/application/application.go index df904c02d..7633c8f81 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -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, storeName, "") + return corebackend.StoreBackend(ml, appConfig, app.backendLoader, 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, storeName, "") + return corebackend.StoreBackend(ml, appConfig, app.backendLoader, storeName, "") } app.voiceRegistry = voicerecognition.NewStoreRegistry(voiceStoreResolver, voiceStoreName, voiceEmbeddingDim) @@ -553,12 +553,17 @@ func (a *Application) start() error { // once at startup and reused across chat sessions that opt in via metadata. if !a.applicationConfig.DisableLocalAIAssistant { holder := mcpTools.NewLocalAIAssistantHolder() + var nodeRegistry *nodes.NodeRegistry + if a.distributed != nil { + nodeRegistry = a.distributed.Registry + } assistantClient := localaiInproc.New( a.applicationConfig, a.applicationConfig.SystemState, a.backendLoader, a.modelLoader, a.galleryService, + nodeRegistry, ) // Wire usage tracking so the assistant's get_usage_stats tool // returns real data; nil values keep the tool returning a clear diff --git a/core/application/router_factories.go b/core/application/router_factories.go index 879c43a83..523c4465e 100644 --- a/core/application/router_factories.go +++ b/core/application/router_factories.go @@ -46,12 +46,12 @@ type lazyScorer struct { modelName string } -func (l *lazyScorer) Score(ctx context.Context, prompt string, candidates []string) ([]backend.CandidateScore, error) { +func (l *lazyScorer) Score(ctx context.Context, prompt string, stablePrefixLen int, 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, candidates) + return backend.NewScorer(l.app.modelLoader, *cfg, l.app.applicationConfig).Score(ctx, prompt, stablePrefixLen, 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, storeName) + return backend.NewVectorStore(a.modelLoader, a.applicationConfig, a.backendLoader, storeName) } diff --git a/core/application/router_factories_test.go b/core/application/router_factories_test.go index 5a6988a88..91b26d491 100644 --- a/core/application/router_factories_test.go +++ b/core/application/router_factories_test.go @@ -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", []string{"a"}) + _, err := sc.Score(context.Background(), "prompt", 0, []string{"a"}) Expect(err).To(HaveOccurred()) Expect(err.Error()).To(ContainSubstring("no longer available")) }) diff --git a/core/application/startup.go b/core/application/startup.go index a28f0f847..b46af7046 100644 --- a/core/application/startup.go +++ b/core/application/startup.go @@ -24,6 +24,7 @@ import ( "github.com/mudler/LocalAI/core/services/routing/router" "github.com/mudler/LocalAI/core/services/storage" coreStartup "github.com/mudler/LocalAI/core/startup" + "github.com/mudler/LocalAI/core/trace" "github.com/mudler/LocalAI/internal" "github.com/mudler/LocalAI/pkg/downloader" "github.com/mudler/LocalAI/pkg/modelartifacts" @@ -60,6 +61,7 @@ func New(opts ...config.AppOption) (*Application, error) { options.Threads = xsysinfo.CPUPhysicalCores() } + trace.ConfigureBackendTracePersistence(options.DataPath) application := newApplication(options) application.startupConfig = &startupConfigCopy @@ -133,7 +135,6 @@ func New(opts ...config.AppOption) (*Application, error) { migrateDataFiles(options.DynamicConfigsDir, options.DataPath) } } - // Initialize auth database if auth is enabled if options.Auth.Enabled { // Auto-generate HMAC secret if not provided @@ -443,6 +444,13 @@ func New(opts ...config.AppOption) (*Application, error) { // when gallery data refreshes instead of using a fixed TTL. vram.SetGalleryGenerationFunc(gallery.GalleryGeneration) + // Fill those caches ahead of the first visitor. An estimate for an entry + // nobody has asked about yet costs a remote probe of its weight files, and + // the model gallery asks for one per row, so without this the first page + // spends seconds filling in its own sizes while somebody watches it. + // Non-blocking, and bounded: see DefaultEstimateWarmConfig. + gallery.WarmEstimateCache(options.Context, options.Galleries, options.SystemState, gallery.EstimateWarmConfigFromEnv()) + if options.ConfigFile != "" { if err := application.ModelConfigLoader().LoadMultipleModelConfigsSingleFile(options.ConfigFile, configLoaderOpts...); err != nil { xlog.Error("error loading config file", "error", err) diff --git a/core/backend/audio_transform.go b/core/backend/audio_transform.go index bb006fc10..99c46a4ff 100644 --- a/core/backend/audio_transform.go +++ b/core/backend/audio_transform.go @@ -33,6 +33,22 @@ type AudioTransformOutputs struct { Dst string AudioPath string ReferencePath string + // Stems are the other named outputs the same run produced, in the model's + // own order and including the one whose content Dst carries. Empty for a + // single-output transform. + // + // A separation backend writes every stem beside Dst from ONE inference. + // Dropping them here would mean a caller who wants drums as well as vocals + // has to run the whole separation again per stem, which is precisely what + // the single run exists to avoid. + Stems []AudioTransformStem +} + +// AudioTransformStem is one named output of a multi-output transform, e.g. the +// "vocals" track of a source separation. +type AudioTransformStem struct { + Name string + Dst string } // ModelAudioTransform runs the unary AudioTransform RPC and returns the @@ -128,9 +144,40 @@ func ModelAudioTransform( Dst: dst, AudioPath: persistedAudio, ReferencePath: persistedRef, + Stems: collectStems(res, audioDir), }, res, nil } +// collectStems turns the backend's reported stems into the caller-facing list. +// +// Every path is checked to be a direct child of audioDir, the generated-content +// directory this request handed the backend. A backend is a separate process +// and its response is not this process's data: a stem path pointing at /etc or +// at another user's file would otherwise be served straight back through the +// HTTP layer, which resolves these into URLs. A stem that fails the check is +// dropped rather than fatal, so a well-behaved majority still reaches the +// caller. +func collectStems(res *proto.AudioTransformResult, audioDir string) []AudioTransformStem { + if res == nil || len(res.GetStems()) == 0 { + return nil + } + stems := make([]AudioTransformStem, 0, len(res.GetStems())) + for _, stem := range res.GetStems() { + name, path := stem.GetName(), stem.GetDst() + if name == "" || path == "" { + continue + } + if filepath.Dir(filepath.Clean(path)) != filepath.Clean(audioDir) { + continue + } + stems = append(stems, AudioTransformStem{Name: name, Dst: path}) + } + if len(stems) == 0 { + return nil + } + return stems +} + // ModelAudioTransformStream opens the bidirectional AudioTransformStream RPC // and returns the underlying stream client. The caller is responsible for // sending the initial Config message, subsequent Frame messages, and for diff --git a/core/backend/detokenize.go b/core/backend/detokenize.go new file mode 100644 index 000000000..05c16c4a5 --- /dev/null +++ b/core/backend/detokenize.go @@ -0,0 +1,67 @@ +package backend + +import ( + "time" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/schema" + "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" +) + +func ModelDetokenize(tokens []int32, loader *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig) (schema.DetokenizeResponse, error) { + + var inferenceModel grpc.Backend + var err error + + opts := ModelOptions(modelConfig, appConfig) + inferenceModel, err = loader.Load(opts...) + if err != nil { + recordModelLoadFailure(appConfig, modelConfig.Name, modelConfig.Backend, err, nil) + return schema.DetokenizeResponse{}, err + } + + var startTime time.Time + if appConfig.EnableTracing { + trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) + startTime = time.Now() + } + + resp, err := inferenceModel.Detokenize(appConfig.Context, &pb.DetokenizeRequest{Tokens: tokens}) + + if appConfig.EnableTracing { + errStr := "" + if err != nil { + errStr = err.Error() + } + + content := "" + if resp != nil { + content = resp.Content + } + + trace.RecordBackendTrace(trace.BackendTrace{ + Timestamp: startTime, + Duration: time.Since(startTime), + Type: trace.BackendTraceTokenize, + ModelName: modelConfig.Name, + Backend: modelConfig.Backend, + Summary: trace.TruncateString(content, 200), + Error: errStr, + Data: map[string]any{ + "token_count": len(tokens), + "output_text": trace.TruncateString(content, 1000), + }, + }) + } + + if err != nil { + return schema.DetokenizeResponse{}, err + } + + return schema.DetokenizeResponse{ + Content: resp.Content, + }, nil +} diff --git a/core/backend/model3d.go b/core/backend/model3d.go new file mode 100644 index 000000000..1952501d9 --- /dev/null +++ b/core/backend/model3d.go @@ -0,0 +1,104 @@ +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 +} diff --git a/core/backend/options.go b/core/backend/options.go index 72f49f1ed..c7525fee8 100644 --- a/core/backend/options.go +++ b/core/backend/options.go @@ -166,6 +166,21 @@ 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), @@ -178,16 +193,7 @@ func ModelOptions(c config.ModelConfig, so *config.ApplicationConfig, opts ...mo defOpts = append(defOpts, model.WithModelFile(c.ModelFileName())) } - threads := 1 - - if c.Threads != nil { - threads = *c.Threads - } - - if so.Threads != 0 { - threads = so.Threads - } - + threads := effectiveThreads(c, so.Threads) c.Threads = &threads grpcOpts := grpcModelOpts(c, so.SystemState.Model.ModelsPath) @@ -416,6 +422,7 @@ 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), @@ -475,6 +482,7 @@ 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, } } diff --git a/core/backend/options_internal_test.go b/core/backend/options_internal_test.go index bf4258bbf..19e62504c 100644 --- a/core/backend/options_internal_test.go +++ b/core/backend/options_internal_test.go @@ -120,6 +120,7 @@ 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() { @@ -128,6 +129,14 @@ 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() { @@ -355,3 +364,23 @@ 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)) + }) +}) diff --git a/core/backend/preload.go b/core/backend/preload.go index 103d36efc..3525da299 100644 --- a/core/backend/preload.go +++ b/core/backend/preload.go @@ -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) + stages, err := pipelineStages(cl, &cfg.Pipeline, ml.ModelPath, appConfig.ToConfigLoaderOptions()...) 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) ([]PreloadStage, error) { +func pipelineStages(cl *config.ModelConfigLoader, p *config.Pipeline, modelPath string, opts ...config.ConfigLoaderOption) ([]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) + cfg, err := cl.LoadResolvedModelConfig(s.name, modelPath, opts...) if err != nil { return nil, fmt.Errorf("%s (%s): %w", s.role, s.name, err) } @@ -87,9 +87,11 @@ 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 -// (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. +// 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. func PreloadStages(ctx context.Context, ml *model.ModelLoader, appConfig *config.ApplicationConfig, stages []PreloadStage) ([]string, error) { var ( wg sync.WaitGroup @@ -98,7 +100,7 @@ func PreloadStages(ctx context.Context, ml *model.ModelLoader, appConfig *config errs []error ) for _, s := range stages { - if s.Cfg == nil { + if s.Cfg == nil || s.Cfg.Name == "" { continue } wg.Add(1) diff --git a/core/backend/preload_internal_test.go b/core/backend/preload_internal_test.go index f92d2b015..5a0fc0f2c 100644 --- a/core/backend/preload_internal_test.go +++ b/core/backend/preload_internal_test.go @@ -103,12 +103,13 @@ var _ = Describe("PreloadStages", func() { return PreloadStage{Role: role, Cfg: &config.ModelConfig{Name: name}} } - It("loads every present stage, skips absent (nil-config) ones, and returns the loaded names", func() { + It("loads every present stage, skips absent stages, and returns the loaded names", func() { stubLoader(nil) loaded, err := PreloadStages(context.Background(), nil, nil, []PreloadStage{ mkStage("vad", "vad-m"), - {Role: "transcription"}, // absent stage + {Role: "transcription"}, + mkStage("tts", ""), mkStage("llm", "llm-m"), }) diff --git a/core/backend/score.go b/core/backend/score.go index 7ce795f60..174e86500 100644 --- a/core/backend/score.go +++ b/core/backend/score.go @@ -23,6 +23,10 @@ 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 @@ -42,9 +46,13 @@ 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. +// 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). type Scorer interface { - Score(ctx context.Context, prompt string, candidates []string) ([]CandidateScore, error) + Score(ctx context.Context, prompt string, stablePrefixLen int, candidates []string) ([]CandidateScore, error) } // NewScorer binds (loader, modelConfig, appConfig) into a Scorer. The @@ -61,8 +69,8 @@ type modelScorer struct { appConfig *config.ApplicationConfig } -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) +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) if err != nil { return nil, err } @@ -103,6 +111,7 @@ 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 { diff --git a/core/backend/stores.go b/core/backend/stores.go index 480400f42..1c324ab1a 100644 --- a/core/backend/stores.go +++ b/core/backend/stores.go @@ -9,6 +9,7 @@ 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" ) @@ -23,21 +24,25 @@ type VectorStore interface { // NewVectorStore returns a VectorStore backed by the local-store // gRPC backend, namespaced by storeName so two routers don't collide. -func NewVectorStore(loader *model.ModelLoader, appConfig *config.ApplicationConfig, storeName string) VectorStore { +// 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 { if storeName == "" { return nil } - return &localVectorStore{loader: loader, appConfig: appConfig, storeName: storeName} + return &localVectorStore{loader: loader, appConfig: appConfig, cl: cl, 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.storeName, "") + return StoreBackend(s.loader, s.appConfig, s.cl, s.storeName, "") } func (s *localVectorStore) Search(ctx context.Context, vec []float32) (sim float64, payload []byte, ok bool, err error) { @@ -121,7 +126,24 @@ func (s *localVectorStore) recordTrace(start time.Time, op string, vecDim int, s }) } -func StoreBackend(sl *model.ModelLoader, appConfig *config.ApplicationConfig, storeName string, backend string) (grpc.Backend, error) { +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 + } + } + if backend == "" { backend = model.LocalStoreBackend } @@ -145,5 +167,12 @@ func StoreBackend(sl *model.ModelLoader, appConfig *config.ApplicationConfig, st 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...) } diff --git a/core/backend/upscale.go b/core/backend/upscale.go new file mode 100644 index 000000000..c821f56cd --- /dev/null +++ b/core/backend/upscale.go @@ -0,0 +1,37 @@ +package backend + +import ( + "context" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/grpc/proto" + model "github.com/mudler/LocalAI/pkg/model" +) + +// ImageUpscale loads the model specified in modelConfig and calls UpscaleImage +// on the backend, writing the result to dst. +func ImageUpscale(ctx context.Context, src, dst string, scale int, loader *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig) (func() error, error) { + opts := ModelOptions(modelConfig, appConfig, model.WithContext(ctx)) + inferenceModel, err := loader.Load(opts...) + if err != nil { + recordModelLoadFailure(appConfig, modelConfig.Name, modelConfig.Backend, err, nil) + return nil, err + } + + fn := func() error { + _, err := inferenceModel.UpscaleImage( + ctx, + &proto.UpscaleImageRequest{ + Src: src, + Dst: dst, + Scale: int32(scale), + }, + ) + return err + } + + return fn, nil +} + +// ImageUpscaleFunc is a test-friendly indirection. +var ImageUpscaleFunc = ImageUpscale diff --git a/core/cli/chat/chat.go b/core/cli/chat/chat.go deleted file mode 100644 index 071d3a785..000000000 --- a/core/cli/chat/chat.go +++ /dev/null @@ -1,30 +0,0 @@ -package chat - -import ( - "context" - "io" - "strings" -) - -type Options struct { - Model string - BaseURL string - APIKey string - In io.Reader - Out io.Writer -} - -func Run(ctx context.Context, opts Options) error { - if opts.In == nil { - opts.In = strings.NewReader("") - } - if opts.Out == nil { - opts.Out = io.Discard - } - - session, err := newChatSession(ctx, newLocalAIChatClient(opts.BaseURL, opts.APIKey), opts.Model) - if err != nil { - return err - } - return runTerminalChat(ctx, session, opts.In, opts.Out) -} diff --git a/core/cli/chat/chat_test.go b/core/cli/chat/chat_test.go deleted file mode 100644 index 7399c3802..000000000 --- a/core/cli/chat/chat_test.go +++ /dev/null @@ -1,172 +0,0 @@ -package chat - -import ( - "bytes" - "encoding/json" - "fmt" - "io" - "net/http" - "net/http/httptest" - "strings" - - . "github.com/onsi/ginkgo/v2" - . "github.com/onsi/gomega" -) - -var _ = Describe("Run chat", func() { - It("streams a single chat response", func() { - var capturedModel string - var capturedAuth string - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == "/v1/models" { - w.Header().Set("Content-Type", "application/json") - writeResponse(w, `{"object":"list","data":[{"id":"test-model","object":"model"}]}`) - return - } - - Expect(r.URL.Path).To(Equal("/v1/chat/completions")) - capturedAuth = r.Header.Get("Authorization") - - var body struct { - Model string `json:"model"` - Messages []struct { - Role string `json:"role"` - Content string `json:"content"` - } `json:"messages"` - } - Expect(json.NewDecoder(r.Body).Decode(&body)).To(Succeed()) - capturedModel = body.Model - Expect(body.Messages).To(HaveLen(1)) - Expect(body.Messages[0].Role).To(Equal("user")) - Expect(body.Messages[0].Content).To(Equal("hello")) - - w.Header().Set("Content-Type", "text/event-stream") - writeResponse(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}]}\n\n") - writeResponse(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"!\"}}]}\n\n") - writeResponse(w, "data: [DONE]\n\n") - })) - defer server.Close() - - var out bytes.Buffer - err := Run(GinkgoT().Context(), Options{ - Model: "test-model", - BaseURL: server.URL + "/v1", - APIKey: "secret", - In: strings.NewReader("hello\n/exit\n"), - Out: &out, - }) - - Expect(err).ToNot(HaveOccurred()) - Expect(capturedModel).To(Equal("test-model")) - Expect(capturedAuth).To(Equal("Bearer secret")) - Expect(out.String()).To(ContainSubstring("assistant: hi!")) - Expect(out.String()).To(ContainSubstring("bye")) - }) - - It("auto-selects the only available model", func() { - server := chatTestServer([]string{"solo"}, nil) - defer server.Close() - - var out bytes.Buffer - err := Run(GinkgoT().Context(), Options{ - BaseURL: server.URL + "/v1", - In: strings.NewReader("/exit\n"), - Out: &out, - }) - - Expect(err).ToNot(HaveOccurred()) - Expect(out.String()).To(ContainSubstring("LocalAI chat (solo)")) - }) - - It("returns an actionable error when no models are installed", func() { - server := chatTestServer(nil, nil) - defer server.Close() - - err := Run(GinkgoT().Context(), Options{ - BaseURL: server.URL + "/v1", - In: strings.NewReader(""), - }) - - Expect(err).To(HaveOccurred()) - Expect(err.Error()).To(ContainSubstring("no chat models are installed")) - Expect(err.Error()).To(ContainSubstring("local-ai models install ")) - }) - - It("returns an actionable error when multiple models are available without a selection", func() { - server := chatTestServer([]string{"alpha", "beta"}, nil) - defer server.Close() - - err := Run(GinkgoT().Context(), Options{ - BaseURL: server.URL + "/v1", - In: strings.NewReader(""), - }) - - Expect(err).To(HaveOccurred()) - Expect(err.Error()).To(ContainSubstring("multiple models are available")) - Expect(err.Error()).To(ContainSubstring("--model")) - Expect(err.Error()).To(ContainSubstring("alpha")) - Expect(err.Error()).To(ContainSubstring("beta")) - }) - - It("lists and switches models inside the chat", func() { - requestedModels := []string{} - server := chatTestServer([]string{"alpha", "beta"}, func(model string) { - requestedModels = append(requestedModels, model) - }) - defer server.Close() - - var out bytes.Buffer - err := Run(GinkgoT().Context(), Options{ - Model: "alpha", - BaseURL: server.URL + "/v1", - In: strings.NewReader("/models\n/model beta\nhello\n/exit\n"), - Out: &out, - }) - - Expect(err).ToNot(HaveOccurred()) - Expect(out.String()).To(ContainSubstring("* alpha")) - Expect(out.String()).To(ContainSubstring(" beta")) - Expect(out.String()).To(ContainSubstring("switched to beta; conversation cleared")) - Expect(requestedModels).To(Equal([]string{"beta"})) - }) -}) - -func chatTestServer(models []string, onChat func(model string)) *httptest.Server { - return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case "/v1/models": - w.Header().Set("Content-Type", "application/json") - writeResponse(w, `{"object":"list","data":[`) - for i, model := range models { - if i > 0 { - writeResponse(w, ",") - } - writeResponsef(w, `{"id":%q,"object":"model"}`, model) - } - writeResponse(w, `]}`) - case "/v1/chat/completions": - var body struct { - Model string `json:"model"` - } - Expect(json.NewDecoder(r.Body).Decode(&body)).To(Succeed()) - if onChat != nil { - onChat(body.Model) - } - w.Header().Set("Content-Type", "text/event-stream") - writeResponse(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"}}]}\n\n") - writeResponse(w, "data: [DONE]\n\n") - default: - w.WriteHeader(http.StatusNotFound) - } - })) -} - -func writeResponse(w io.Writer, text string) { - _, err := fmt.Fprint(w, text) - Expect(err).ToNot(HaveOccurred()) -} - -func writeResponsef(w io.Writer, format string, args ...any) { - _, err := fmt.Fprintf(w, format, args...) - Expect(err).ToNot(HaveOccurred()) -} diff --git a/core/cli/chat/client.go b/core/cli/chat/client.go deleted file mode 100644 index 407845d0b..000000000 --- a/core/cli/chat/client.go +++ /dev/null @@ -1,114 +0,0 @@ -package chat - -import ( - "context" - "errors" - "fmt" - "io" - "sort" - "strings" - - openai "github.com/sashabaranov/go-openai" -) - -type chatClient interface { - ListModels(ctx context.Context) ([]string, error) - StreamChat(ctx context.Context, model string, messages []chatMessage, out io.Writer) (string, error) -} - -type localAIChatClient struct { - client *openai.Client -} - -func newLocalAIChatClient(baseURL string, apiKey string) *localAIChatClient { - cfg := openai.DefaultConfig(apiKey) - cfg.BaseURL = baseURL - return &localAIChatClient{client: openai.NewClientWithConfig(cfg)} -} - -func (c *localAIChatClient) ListModels(ctx context.Context) ([]string, error) { - resp, err := c.client.ListModels(ctx) - if err != nil { - return nil, err - } - - models := make([]string, 0, len(resp.Models)) - for _, model := range resp.Models { - if model.ID != "" { - models = append(models, model.ID) - } - } - sort.Strings(models) - return models, nil -} - -func (c *localAIChatClient) StreamChat(ctx context.Context, model string, messages []chatMessage, out io.Writer) (string, error) { - stream, err := c.client.CreateChatCompletionStream(ctx, openai.ChatCompletionRequest{ - Model: model, - Messages: openAIChatMessages(messages), - }) - if err != nil { - return "", friendlyChatError(err, model) - } - defer func() { - _ = stream.Close() - }() - - var answer strings.Builder - for { - resp, err := stream.Recv() - if errors.Is(err, io.EOF) { - break - } - if err != nil { - return answer.String(), friendlyChatError(err, model) - } - if len(resp.Choices) == 0 { - continue - } - - token := resp.Choices[0].Delta.Content - if token == "" { - continue - } - answer.WriteString(token) - if _, err := fmt.Fprint(out, token); err != nil { - return answer.String(), err - } - } - - return answer.String(), nil -} - -func openAIChatMessages(messages []chatMessage) []openai.ChatCompletionMessage { - converted := make([]openai.ChatCompletionMessage, len(messages)) - for i, message := range messages { - converted[i] = openai.ChatCompletionMessage{ - Role: message.Role, - Content: message.Content, - } - } - return converted -} - -func friendlyChatError(err error, model string) error { - var apiErr *openai.APIError - if errors.As(err, &apiErr) { - switch apiErr.HTTPStatusCode { - case 404: - return fmt.Errorf("model %q is not available. Run `local-ai models list`, install a model with `local-ai models install `, or switch with `/model `", model) - case 403: - return fmt.Errorf("model %q is disabled. Enable it from LocalAI settings or choose another model with `/model `", model) - } - if apiErr.Message != "" { - return errors.New(apiErr.Message) - } - } - - msg := err.Error() - if strings.Contains(msg, "model") && strings.Contains(msg, "not found") { - return fmt.Errorf("model %q is not available. Run `local-ai models list`, install a model with `local-ai models install `, or switch with `/model `", model) - } - - return err -} diff --git a/core/cli/chat/models.go b/core/cli/chat/models.go deleted file mode 100644 index 291ec15aa..000000000 --- a/core/cli/chat/models.go +++ /dev/null @@ -1,17 +0,0 @@ -package chat - -import "strings" - -func formatChatModelList(models []string, current string) string { - var b strings.Builder - for _, model := range models { - prefix := " " - if model == current { - prefix = "* " - } - b.WriteString(prefix) - b.WriteString(model) - b.WriteByte('\n') - } - return b.String() -} diff --git a/core/cli/chat/paths.go b/core/cli/chat/paths.go new file mode 100644 index 000000000..1cfa34821 --- /dev/null +++ b/core/cli/chat/paths.go @@ -0,0 +1,153 @@ +package chat + +import ( + "fmt" + "os" + "path/filepath" + + "gopkg.in/yaml.v3" +) + +// stateDirMode matches the mode nib uses for the same directory. The directory +// holds an API key, so it stays owner-only. +const stateDirMode = 0o700 + +// configFileMode keeps the config owner-only: nib stores the user's API key in +// it alongside the keys written here. +const configFileMode = 0o600 + +// StateDir resolves where the chat agent keeps its config, plugins, and +// skills. This is user-scoped rather than server-scoped: chat is a client that +// may target a remote LocalAI, so it does not belong under LOCALAI_CONFIG_DIR. +func StateDir(override string) (string, error) { + if override != "" { + return override, nil + } + if xdg := os.Getenv("XDG_CONFIG_HOME"); xdg != "" { + return filepath.Join(xdg, "localai", "chat"), nil + } + home, err := os.UserHomeDir() + if err != nil { + return "", fmt.Errorf("resolving home directory for the agent state dir: %w", err) + } + return filepath.Join(home, ".config", "localai", "chat"), nil +} + +// ConfigPath is the agent's config file inside dir. +func ConfigPath(dir string) string { return filepath.Join(dir, "config.yaml") } + +// EnsureStateDir creates dir and, on first run only, seeds a config file +// pointing at baseURL. It deliberately does not seed a model: a baked-in model +// name goes stale as soon as the user installs a different one. +// +// The config file is machine-managed from here on: nib rewrites it whenever it +// self-configures, so hand-written comments in it do not survive. +func EnsureStateDir(dir, baseURL string) error { + if err := os.MkdirAll(dir, stateDirMode); err != nil { + return fmt.Errorf("creating agent state dir %s: %w", dir, err) + } + path := ConfigPath(dir) + if _, err := os.Stat(path); err == nil { + return nil // already configured; never overwrite the user's file + } else if !os.IsNotExist(err) { + return fmt.Errorf("checking agent config %s: %w", path, err) + } + + seed := map[string]string{"base_url": baseURL} + data, err := yaml.Marshal(seed) + if err != nil { + return fmt.Errorf("encoding seed agent config: %w", err) + } + if err := writeConfigFile(path, data); err != nil { + return fmt.Errorf("writing seed agent config: %w", err) + } + return nil +} + +// PersistModel records the chosen model in the agent config, preserving every +// other key the user may have set, including the api_key nib writes there. +// +// The file is machine-managed: this overlays the model onto the parsed keys and +// re-marshals, which drops comments. That is deliberate rather than an +// oversight, because nib's own save path does the same thing and would erase +// them on its next write regardless. +func PersistModel(dir, model string) error { + // PersistModel is callable before EnsureStateDir, so it cannot assume the + // directory exists. + if err := os.MkdirAll(dir, stateDirMode); err != nil { + return fmt.Errorf("creating agent state dir %s: %w", dir, err) + } + path := ConfigPath(dir) + + values := map[string]any{} + // #nosec G304 -- path is the fixed config.yaml name under the user-selected + // chat state directory; selecting that directory is the documented override. + data, err := os.ReadFile(path) + if err != nil && !os.IsNotExist(err) { + return fmt.Errorf("reading agent config %s: %w", path, err) + } + if err == nil { + if err := yaml.Unmarshal(data, &values); err != nil { + return fmt.Errorf("parsing agent config %s: %w", path, err) + } + } + values["model"] = model + + out, err := yaml.Marshal(values) + if err != nil { + return fmt.Errorf("encoding agent config: %w", err) + } + if err := writeConfigFile(path, out); err != nil { + return fmt.Errorf("writing agent config: %w", err) + } + return nil +} + +// writeConfigFile replaces path with data atomically: it writes a temporary +// file next to the target and renames it over the target. Writing the target in +// place would truncate it first, so an interrupted or out-of-disk write would +// leave a half-written config and destroy the api_key nib keeps in the same +// file. The temporary file must share the directory because rename is only +// atomic within one filesystem. +func writeConfigFile(path string, data []byte) error { + dir := filepath.Dir(path) + + // A randomized name rather than a fixed config.yaml.tmp, so two concurrent + // writers cannot corrupt each other's temporary file. + tmp, err := os.CreateTemp(dir, "config.yaml.*.tmp") + if err != nil { + return fmt.Errorf("creating temp file in %s: %w", dir, err) + } + tmpPath := tmp.Name() + renamed := false + defer func() { + if !renamed { + // Leave no litter behind on any failure path. + _ = os.Remove(tmpPath) + } + }() + + if _, err := tmp.Write(data); err != nil { + _ = tmp.Close() + return fmt.Errorf("writing %s: %w", tmpPath, err) + } + // Flush before the rename: renaming a file whose contents are still only in + // the page cache can still lose them across a crash. + if err := tmp.Sync(); err != nil { + _ = tmp.Close() + return fmt.Errorf("syncing %s: %w", tmpPath, err) + } + if err := tmp.Close(); err != nil { + return fmt.Errorf("closing %s: %w", tmpPath, err) + } + // CreateTemp already asks for 0600, but the umask can only ever clear bits, + // so set the mode explicitly rather than inheriting whatever survived. + if err := os.Chmod(tmpPath, configFileMode); err != nil { + return fmt.Errorf("setting mode on %s: %w", tmpPath, err) + } + if err := os.Rename(tmpPath, path); err != nil { + return fmt.Errorf("replacing %s: %w", path, err) + } + renamed = true + return nil +} diff --git a/core/cli/chat/paths_test.go b/core/cli/chat/paths_test.go new file mode 100644 index 000000000..7a4b9aa29 --- /dev/null +++ b/core/cli/chat/paths_test.go @@ -0,0 +1,186 @@ +package chat + +import ( + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gopkg.in/yaml.v3" +) + +// richConfig stands in for a config nib has already taken ownership of: a +// comment, a secret, and a nested block. A flat scalar alone would not catch a +// writer that mangles structure or drops a key it does not know about. +const richConfig = `# hand written note +base_url: http://x.invalid/v1 +api_key: secret-token +mcp_servers: + files: + command: mcp-files + args: + - --root + - /tmp +` + +var _ = Describe("Agent state directory", func() { + Describe("StateDir", func() { + It("prefers an explicit override", func() { + Expect(StateDir("/custom/dir")).To(Equal("/custom/dir")) + }) + + It("uses XDG_CONFIG_HOME when set", func() { + tmp := GinkgoT().TempDir() + GinkgoT().Setenv("XDG_CONFIG_HOME", tmp) + Expect(StateDir("")).To(Equal(filepath.Join(tmp, "localai", "chat"))) + }) + + It("falls back to ~/.config/localai/chat", func() { + tmp := GinkgoT().TempDir() + GinkgoT().Setenv("XDG_CONFIG_HOME", "") + GinkgoT().Setenv("HOME", tmp) + Expect(StateDir("")).To(Equal(filepath.Join(tmp, ".config", "localai", "chat"))) + }) + + It("fails when neither XDG_CONFIG_HOME nor a home directory is resolvable", func() { + GinkgoT().Setenv("XDG_CONFIG_HOME", "") + GinkgoT().Setenv("HOME", "") + + dir, err := StateDir("") + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("agent state dir")) + // No silent fallback to a relative path: writing an API key into the + // working directory would be worse than refusing. + Expect(dir).To(BeEmpty()) + }) + }) + + Describe("EnsureStateDir", func() { + It("creates the directory and seeds base_url on first run", func() { + dir := filepath.Join(GinkgoT().TempDir(), "chat") + Expect(EnsureStateDir(dir, "http://127.0.0.1:8080/v1")).To(Succeed()) + + data, err := os.ReadFile(ConfigPath(dir)) + Expect(err).ToNot(HaveOccurred()) + Expect(string(data)).To(ContainSubstring("base_url: http://127.0.0.1:8080/v1")) + // A model must NOT be seeded: it goes stale as soon as the user + // installs a different one. + Expect(string(data)).ToNot(ContainSubstring("model:")) + }) + + It("keeps the seeded config and its directory owner-only", func() { + dir := filepath.Join(GinkgoT().TempDir(), "chat") + Expect(EnsureStateDir(dir, "http://127.0.0.1:8080/v1")).To(Succeed()) + + // nib writes the user's api_key into this same file, so the modes are + // load-bearing, not cosmetic. + config, err := os.Stat(ConfigPath(dir)) + Expect(err).ToNot(HaveOccurred()) + Expect(config.Mode().Perm()).To(Equal(os.FileMode(0o600))) + + state, err := os.Stat(dir) + Expect(err).ToNot(HaveOccurred()) + Expect(state.Mode().Perm()).To(Equal(os.FileMode(0o700))) + }) + + It("leaves an existing config byte-for-byte untouched", func() { + dir := GinkgoT().TempDir() + Expect(os.WriteFile(ConfigPath(dir), []byte(richConfig), 0o600)).To(Succeed()) + + Expect(EnsureStateDir(dir, "http://127.0.0.1:8080/v1")).To(Succeed()) + + data, err := os.ReadFile(ConfigPath(dir)) + Expect(err).ToNot(HaveOccurred()) + // Byte-exact against a fixture carrying a comment and a nested block: + // an implementation that "preserves" by re-marshaling through a map + // fails here rather than passing on a flat scalar. + Expect(string(data)).To(Equal(richConfig)) + }) + }) + + Describe("PersistModel", func() { + It("adds a model to an existing config, preserving other keys", func() { + dir := GinkgoT().TempDir() + Expect(os.WriteFile(ConfigPath(dir), []byte("base_url: http://x.invalid/v1\n"), 0o600)).To(Succeed()) + + Expect(PersistModel(dir, "chosen-model")).To(Succeed()) + + data, err := os.ReadFile(ConfigPath(dir)) + Expect(err).ToNot(HaveOccurred()) + Expect(string(data)).To(ContainSubstring("base_url: http://x.invalid/v1")) + Expect(string(data)).To(ContainSubstring("model: chosen-model")) + }) + + It("replaces an existing model rather than duplicating the key", func() { + dir := GinkgoT().TempDir() + Expect(os.WriteFile(ConfigPath(dir), []byte("model: old\nbase_url: http://x.invalid/v1\n"), 0o600)).To(Succeed()) + + Expect(PersistModel(dir, "new")).To(Succeed()) + + data, err := os.ReadFile(ConfigPath(dir)) + Expect(err).ToNot(HaveOccurred()) + Expect(string(data)).To(ContainSubstring("model: new")) + Expect(string(data)).ToNot(ContainSubstring("model: old")) + }) + + It("preserves secrets and nested blocks it does not understand", func() { + dir := GinkgoT().TempDir() + Expect(os.WriteFile(ConfigPath(dir), []byte(richConfig), 0o600)).To(Succeed()) + + Expect(PersistModel(dir, "chosen-model")).To(Succeed()) + + data, err := os.ReadFile(ConfigPath(dir)) + Expect(err).ToNot(HaveOccurred()) + + var got map[string]any + Expect(yaml.Unmarshal(data, &got)).To(Succeed()) + Expect(got).To(HaveKeyWithValue("model", "chosen-model")) + Expect(got).To(HaveKeyWithValue("base_url", "http://x.invalid/v1")) + // Losing this key logs the user out of their own server. + Expect(got).To(HaveKeyWithValue("api_key", "secret-token")) + Expect(got).To(HaveKeyWithValue("mcp_servers", + HaveKeyWithValue("files", And( + HaveKeyWithValue("command", "mcp-files"), + HaveKeyWithValue("args", ConsistOf("--root", "/tmp")), + )), + )) + + // Documented, accepted behavior rather than an aspiration: the overlay + // re-marshals, so comments do not survive. nib's own save path erases + // them too, so preserving them here would buy nothing. + Expect(string(data)).ToNot(ContainSubstring("# hand written note")) + }) + + It("keeps the rewritten config owner-only and leaves no temp file behind", func() { + dir := GinkgoT().TempDir() + Expect(os.WriteFile(ConfigPath(dir), []byte(richConfig), 0o600)).To(Succeed()) + + Expect(PersistModel(dir, "chosen-model")).To(Succeed()) + + info, err := os.Stat(ConfigPath(dir)) + Expect(err).ToNot(HaveOccurred()) + Expect(info.Mode().Perm()).To(Equal(os.FileMode(0o600))) + + // The atomic write stages through a sibling temp file; it must not + // survive a successful write. + entries, err := os.ReadDir(dir) + Expect(err).ToNot(HaveOccurred()) + names := []string{} + for _, entry := range entries { + names = append(names, entry.Name()) + } + Expect(names).To(ConsistOf("config.yaml")) + }) + + It("creates the state directory when it does not exist yet", func() { + // Task 4 may persist a picked model before anything else has run. + dir := filepath.Join(GinkgoT().TempDir(), "chat") + + Expect(PersistModel(dir, "chosen-model")).To(Succeed()) + + data, err := os.ReadFile(ConfigPath(dir)) + Expect(err).ToNot(HaveOccurred()) + Expect(string(data)).To(ContainSubstring("model: chosen-model")) + }) + }) +}) diff --git a/core/cli/chat/probe.go b/core/cli/chat/probe.go new file mode 100644 index 000000000..6b5032fbd --- /dev/null +++ b/core/cli/chat/probe.go @@ -0,0 +1,86 @@ +package chat + +import ( + "context" + "errors" + "fmt" + "net/http" + "net/url" + + openai "github.com/sashabaranov/go-openai" +) + +var ( + // ErrUnreachable means nothing answered at the endpoint. Callers use this + // to decide whether offering to start a server makes sense. + ErrUnreachable = errors.New("no LocalAI server reachable") + // ErrUnauthorized means the server answered but rejected the credentials. + ErrUnauthorized = errors.New("LocalAI server rejected the API key") +) + +// Probe lists the models the endpoint advertises. It classifies the two +// failures that need different advice: nothing listening, and bad credentials. +// +// The returned list is what the server advertises, verbatim and in server +// order. LocalAI happily lists non-model entries it finds in the models +// directory (stray archives, dotfiles), and guessing which advertised IDs are +// real belongs to whoever presents them, not here. +func Probe(ctx context.Context, baseURL, apiKey string) ([]string, error) { + cfg := openai.DefaultConfig(apiKey) + cfg.BaseURL = baseURL + + resp, err := openai.NewClientWithConfig(cfg).ListModels(ctx) + if err != nil { + if status, answered := responseStatus(err); answered { + if status == http.StatusUnauthorized || status == http.StatusForbidden { + return nil, fmt.Errorf("%w: %w", ErrUnauthorized, err) + } + // The server answered, so it is up; surface its error as-is. + return nil, fmt.Errorf("listing models at %s: %w", baseURL, err) + } + // A caller who cancelled the probe learned nothing about the endpoint, + // so claiming it is unreachable would send them to fix a server that + // may be fine. A deadline is left alone: an endpoint that cannot answer + // within the probe's budget is unreachable for our purposes. + var urlErr *url.Error + if errors.As(err, &urlErr) && !errors.Is(err, context.Canceled) { + // Only a failure to complete the round trip means nothing is + // listening. A reply we could not parse is a different problem, + // so it falls through to the generic error below. + return nil, fmt.Errorf("%w at %s: %w", ErrUnreachable, baseURL, err) + } + return nil, fmt.Errorf("listing models at %s: %w", baseURL, err) + } + + models := make([]string, 0, len(resp.Models)) + for _, m := range resp.Models { + if m.ID != "" { + models = append(models, m.ID) + } + } + return models, nil +} + +// responseStatus reports the HTTP status a failed call came back with, and +// whether there was one at all. +// +// go-openai splits this across two types depending on the error body, and both +// occur against a real LocalAI: it returns *openai.APIError when the body +// parses as an OpenAI error envelope, which is what LocalAI's normal error +// handler sends, and *openai.RequestError when it does not, which is what +// LocalAI sends when started with opaque errors, since that handler replies +// with a bare status and no body. +func responseStatus(err error) (int, bool) { + // *RequestError is checked first because it is the outer type when + // go-openai nests one error inside the other; the inner value in that case + // carries no status. + var reqErr *openai.RequestError + if errors.As(err, &reqErr) { + return reqErr.HTTPStatusCode, true + } + var apiErr *openai.APIError + if errors.As(err, &apiErr) { + return apiErr.HTTPStatusCode, true + } + return 0, false +} diff --git a/core/cli/chat/probe_test.go b/core/cli/chat/probe_test.go new file mode 100644 index 000000000..ec9d4705a --- /dev/null +++ b/core/cli/chat/probe_test.go @@ -0,0 +1,169 @@ +package chat + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Probe", func() { + It("returns the advertised models", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + Expect(json.NewEncoder(w).Encode(map[string]any{ + "object": "list", + "data": []map[string]string{ + {"id": "model-a", "object": "model"}, + {"id": "model-b", "object": "model"}, + }, + })).To(Succeed()) + })) + defer srv.Close() + + models, err := Probe(context.Background(), srv.URL+"/v1", "") + Expect(err).ToNot(HaveOccurred()) + Expect(models).To(Equal([]string{"model-a", "model-b"})) + }) + + It("reports an unreachable server distinguishably", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) + url := srv.URL + srv.Close() // nothing is listening now + + _, err := Probe(context.Background(), url+"/v1", "") + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrUnreachable)).To(BeTrue(), "want ErrUnreachable, got %v", err) + }) + + It("reports an auth failure distinguishably", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + })) + defer srv.Close() + + _, err := Probe(context.Background(), srv.URL+"/v1", "bad-key") + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrUnauthorized)).To(BeTrue(), "want ErrUnauthorized, got %v", err) + }) + + // LocalAI's normal error handler replies with an OpenAI error envelope, and + // its opaque-errors handler replies with a bare status and no body. Those + // reach the client as two different go-openai types, so both have to be + // classified the same way. + It("reports an auth failure carrying an error envelope distinguishably", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + Expect(json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{"message": "invalid api key", "code": http.StatusUnauthorized}, + })).To(Succeed()) + })) + defer srv.Close() + + _, err := Probe(context.Background(), srv.URL+"/v1", "bad-key") + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrUnauthorized)).To(BeTrue(), "want ErrUnauthorized, got %v", err) + }) + + It("does not call a server that answered with an error unreachable", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + defer srv.Close() + + _, err := Probe(context.Background(), srv.URL+"/v1", "") + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrUnreachable)).To(BeFalse(), "a server that replied is not unreachable, got %v", err) + Expect(errors.Is(err, ErrUnauthorized)).To(BeFalse(), "500 is not an auth failure, got %v", err) + }) + + // Pointing chat at some other service that happens to be listening is a + // different problem from nothing listening, and needs different advice. + It("does not call a reply it could not parse unreachable", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html") + _, err := w.Write([]byte("not LocalAI")) + Expect(err).ToNot(HaveOccurred()) + })) + defer srv.Close() + + _, err := Probe(context.Background(), srv.URL+"/v1", "") + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrUnreachable)).To(BeFalse(), "something answered, got %v", err) + }) + + It("returns every advertised id, including ones that are not models", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + Expect(json.NewEncoder(w).Encode(map[string]any{ + "object": "list", + "data": []map[string]string{ + {"id": "zeta", "object": "model"}, + {"id": ".gitignore", "object": "model"}, + {"id": "alpha", "object": "model"}, + {"id": "voice.tar.bz2", "object": "model"}, + }, + })).To(Succeed()) + })) + defer srv.Close() + + // Verbatim and in server order: deciding which of these are real, and + // what order to show them in, belongs to the caller. + models, err := Probe(context.Background(), srv.URL+"/v1", "") + Expect(err).ToNot(HaveOccurred()) + Expect(models).To(Equal([]string{"zeta", ".gitignore", "alpha", "voice.tar.bz2"})) + }) + + It("stops early when the context is already cancelled", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + Expect(json.NewEncoder(w).Encode(map[string]any{"object": "list", "data": []any{}})).To(Succeed()) + })) + defer srv.Close() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := Probe(ctx, srv.URL+"/v1", "") + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, context.Canceled)).To(BeTrue(), "want the cancellation preserved, got %v", err) + // A cancelled probe learned nothing about the endpoint, so it must not + // send the caller off to start a server that may already be running. + Expect(errors.Is(err, ErrUnreachable)).To(BeFalse(), "cancelling is not a verdict on the server, got %v", err) + }) + + It("reports a server that never answers as unreachable", func() { + release := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + <-release + })) + defer srv.Close() + defer close(release) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + + _, err := Probe(ctx, srv.URL+"/v1", "") + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrUnreachable)).To(BeTrue(), "want ErrUnreachable, got %v", err) + Expect(errors.Is(err, context.DeadlineExceeded)).To(BeTrue(), "want the deadline preserved, got %v", err) + }) + + It("returns an empty list when the server has no models", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + Expect(json.NewEncoder(w).Encode(map[string]any{"object": "list", "data": []any{}})).To(Succeed()) + })) + defer srv.Close() + + models, err := Probe(context.Background(), srv.URL+"/v1", "") + Expect(err).ToNot(HaveOccurred()) + Expect(models).To(BeEmpty()) + }) +}) diff --git a/core/cli/chat/resolve.go b/core/cli/chat/resolve.go new file mode 100644 index 000000000..8b39f8138 --- /dev/null +++ b/core/cli/chat/resolve.go @@ -0,0 +1,95 @@ +package chat + +import ( + "errors" + "fmt" + "slices" + "sort" + "strings" + + "github.com/mudler/xlog" +) + +// ModelChooser asks the user to pick one of models. It is nil when the session +// is not interactive. +type ModelChooser func(models []string) (string, error) + +// ModelRequest is everything model resolution needs. +type ModelRequest struct { + Flag string // --model + Configured string // model recorded in the agent config + Available []string // models the server advertises + StateDir string // where an interactive choice is persisted + Choose ModelChooser // nil means non-interactive + // Notify reports a problem that is worth telling the user about but not + // worth failing over. Nil discards it. It exists because the one such + // problem here, a choice that could not be saved, changes what the user + // should expect next: they will be asked again. A log line does not reach + // them, since the agent runs at log level error by default. + Notify func(message string) +} + +// ResolveModel picks the model for this invocation. A flag or a configured +// value wins outright and is not persisted; only an interactive choice is +// written back, so the prompt appears at most once. +// +// Available is used exactly as the server gave it. LocalAI advertises stray +// files it finds in the models directory alongside real models, but real model +// IDs contain dots too (lfm2.5-8b-a1b), so any client-side "looks like a +// filename" heuristic would eventually hide a model the user has. Deciding +// which advertised IDs are real belongs to the endpoint, not to a guess here. +func ResolveModel(req ModelRequest) (string, error) { + if req.Flag != "" { + return req.Flag, nil + } + if req.Configured != "" { + return req.Configured, nil + } + + // The server's /v1/models ordering is not stable between calls, so sort + // before showing or listing: the same number must mean the same model on + // the next run. Sort a copy; the caller's slice is not ours to reorder. + available := append([]string(nil), req.Available...) + sort.Strings(available) + + switch len(available) { + case 0: + return "", errors.New("the LocalAI server has no models installed. Install one with 'local-ai models install ', then run 'local-ai chat' again") + case 1: + return available[0], nil + } + + if req.Choose == nil { + return "", fmt.Errorf( + "several models are available; pick one with --model. Available: %s", + strings.Join(available, ", "), + ) + } + + chosen, err := req.Choose(available) + if err != nil { + return "", err + } + // Choose is an interface, so its answer is checked rather than trusted. + // What comes back is persisted and every later run starts against it, so a + // chooser that returns an empty string or a name of its own would record a + // model the server never offered and there would be nothing left to catch + // it. + if !slices.Contains(available, chosen) { + return "", fmt.Errorf( + "the model chooser answered %q, which is not one of the available models: %s", + chosen, strings.Join(available, ", "), + ) + } + if req.StateDir != "" { + if err := PersistModel(req.StateDir, chosen); err != nil { + // A failure to remember the choice must not block the session: the + // user picked a model, so honour it and say what will happen. + xlog.Warn("could not save the model choice", "error", err, "model", chosen) + if req.Notify != nil { + req.Notify(fmt.Sprintf("Your choice of %s could not be saved, so this question comes back next time: %v", chosen, err)) + } + } + } + return chosen, nil +} diff --git a/core/cli/chat/resolve_test.go b/core/cli/chat/resolve_test.go new file mode 100644 index 000000000..bbb49d06e --- /dev/null +++ b/core/cli/chat/resolve_test.go @@ -0,0 +1,156 @@ +package chat + +import ( + "errors" + "os" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("ResolveModel", func() { + It("prefers the flag over everything", func() { + got, err := ResolveModel(ModelRequest{ + Flag: "from-flag", + Configured: "from-config", + Available: []string{"a", "b"}, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(got).To(Equal("from-flag")) + }) + + It("uses the configured model when no flag is given", func() { + got, err := ResolveModel(ModelRequest{ + Configured: "from-config", + Available: []string{"a", "b"}, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(got).To(Equal("from-config")) + }) + + It("auto-selects when the server offers exactly one model", func() { + got, err := ResolveModel(ModelRequest{Available: []string{"only-one"}}) + Expect(err).ToNot(HaveOccurred()) + Expect(got).To(Equal("only-one")) + }) + + It("errors and lists the options when several models exist and there is no chooser", func() { + _, err := ResolveModel(ModelRequest{Available: []string{"a", "b"}}) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("a")) + Expect(err.Error()).To(ContainSubstring("b")) + Expect(err.Error()).To(ContainSubstring("--model")) + }) + + It("sorts before offering, so the same number means the same model next run", func() { + var offered []string + available := []string{"zeta", "alpha", "mid"} + _, err := ResolveModel(ModelRequest{ + Available: available, + StateDir: GinkgoT().TempDir(), + Choose: func(models []string) (string, error) { + offered = models + return models[0], nil + }, + }) + Expect(err).ToNot(HaveOccurred()) + // The server's /v1/models ordering is unstable between calls. + Expect(offered).To(Equal([]string{"alpha", "mid", "zeta"})) + // Sorting must happen on a copy: the caller still owns this slice, and + // reordering it under them would move whatever they index into it. + Expect(available).To(Equal([]string{"zeta", "alpha", "mid"})) + }) + + It("lists models in sorted order in the several-models error", func() { + _, err := ResolveModel(ModelRequest{Available: []string{"zeta", "alpha"}}) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("alpha, zeta")) + }) + + It("asks the chooser when several models exist, and persists the answer", func() { + dir := GinkgoT().TempDir() + got, err := ResolveModel(ModelRequest{ + Available: []string{"a", "b"}, + StateDir: dir, + Choose: func(models []string) (string, error) { return models[1], nil }, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(got).To(Equal("b")) + + data, err := os.ReadFile(ConfigPath(dir)) + Expect(err).ToNot(HaveOccurred()) + Expect(string(data)).To(ContainSubstring("model: b")) + }) + + // The answer is persisted and every later run starts against it, and + // ModelChooser is exported, so the invariant has to hold for choosers this + // package did not write. + DescribeTable("refuses an answer the chooser was not offered", + func(answer string) { + dir := GinkgoT().TempDir() + got, err := ResolveModel(ModelRequest{ + Available: []string{"alpha", "zeta"}, + StateDir: dir, + Choose: func([]string) (string, error) { return answer, nil }, + }) + Expect(err).To(HaveOccurred()) + Expect(got).To(BeEmpty()) + Expect(err.Error()).To(ContainSubstring("alpha, zeta")) + + _, statErr := os.Stat(ConfigPath(dir)) + Expect(os.IsNotExist(statErr)).To(BeTrue(), "nothing may be recorded for an answer that was refused") + }, + Entry("nothing at all", ""), + Entry("a model the server never offered", "gamma"), + Entry("an offered model with stray whitespace", " alpha"), + Entry("an offered model in the wrong case", "Alpha"), + ) + + It("notifies, and still honours the choice, when it cannot be persisted", func() { + dir := GinkgoT().TempDir() + // A directory where the config file belongs: the write fails for any + // user, including root. + Expect(os.MkdirAll(ConfigPath(dir), 0o700)).To(Succeed()) + + var notices []string + got, err := ResolveModel(ModelRequest{ + Available: []string{"a", "b"}, + StateDir: dir, + Choose: func(models []string) (string, error) { return models[0], nil }, + Notify: func(message string) { notices = append(notices, message) }, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(got).To(Equal("a")) + Expect(notices).To(HaveLen(1)) + Expect(notices[0]).To(ContainSubstring("a")) + Expect(notices[0]).To(ContainSubstring("could not be saved")) + }) + + It("says nothing when the choice was saved", func() { + var notices []string + _, err := ResolveModel(ModelRequest{ + Available: []string{"a", "b"}, + StateDir: GinkgoT().TempDir(), + Choose: func(models []string) (string, error) { return models[0], nil }, + Notify: func(message string) { notices = append(notices, message) }, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(notices).To(BeEmpty()) + }) + + It("propagates a chooser cancellation", func() { + cancelled := errors.New("cancelled") + _, err := ResolveModel(ModelRequest{ + Available: []string{"a", "b"}, + StateDir: GinkgoT().TempDir(), + Choose: func([]string) (string, error) { return "", cancelled }, + }) + Expect(errors.Is(err, cancelled)).To(BeTrue()) + }) + + It("errors with an install hint when the server has no models", func() { + _, err := ResolveModel(ModelRequest{Available: nil}) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("local-ai models install")) + }) +}) diff --git a/core/cli/chat/run.go b/core/cli/chat/run.go new file mode 100644 index 000000000..3da8e08d2 --- /dev/null +++ b/core/cli/chat/run.go @@ -0,0 +1,475 @@ +package chat + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "os" + "os/signal" + "strconv" + "strings" + "syscall" + "time" + + "github.com/mudler/nib/app" + nibcmd "github.com/mudler/nib/cmd" + nibconfig "github.com/mudler/nib/config" + nibtypes "github.com/mudler/nib/types" + "golang.org/x/term" +) + +// Options is everything the chat command passes down from its flags. +type Options struct { + Args []string // forwarded to the agent verbatim + Endpoint string // the server root, e.g. http://127.0.0.1:8080 + BaseURL string // the API base, e.g. http://127.0.0.1:8080/v1 + APIKey string + Model string + StateDir string + TraceDir string + Yolo bool + // ProbeTimeout bounds each check of the server. Zero means + // defaultProbeTimeout. + ProbeTimeout time.Duration + + In io.Reader + Out io.Writer + ErrOut io.Writer +} + +// ExitStatus reports the status the process should exit with for an agent run +// that failed, and whether err is such a failure. +// +// nib writes what went wrong to the error stream itself and hands back nothing +// but a code, so an error that satisfies this has already been explained to the +// user and must not be reported a second time. The refusal to open a +// full-screen session on a stdin that cannot be read arrives this way, and it +// is the one a user is most likely to meet: 'echo q | local-ai chat' names +// --cli, and burying that under a second message would hide the fix. +func ExitStatus(err error) (int, bool) { + var exit app.ExitError + if errors.As(err, &exit) { + return exit.Code, true + } + return 0, false +} + +// shutdownSignals end the session. SIGHUP is one of them because this is a +// terminal program: once the terminal is gone there is nobody left to talk to, +// and a server started for the session has to go with it. +var shutdownSignals = []os.Signal{os.Interrupt, syscall.SIGTERM, syscall.SIGHUP} + +// shutdownContext derives a context that is cancelled when the process is +// asked to stop. +// +// Without it a signal kills this process where it stands, skipping every +// deferred call, and a 'local-ai run' started for the session is reparented to +// init with nothing left that knows to shut it down. An interactive Ctrl+C is +// safe on its own, because the child shares this process' foreground process +// group and the terminal signals all of it, but a SIGTERM from a supervisor or +// a script reaches only this process. +// +// Since nib v0.5.1 cancelling this context does end the session: RunTUI passes +// it to bubbletea, which unwinds the program and reports the context's own +// error. The server is still stopped on cancellation rather than on the way +// out (see runSession), because registering here removes SIGHUP's default +// terminate disposition, and a guarantee about a server this process owns is +// not worth resting on how promptly a third party unwinds its interface. +// +// A handler rather than SysProcAttr.Pdeathsig on the child: Pdeathsig is +// Linux-only, and in Go it is delivered when the OS thread that forked exits +// rather than when the process does, so it can fire on a perfectly healthy +// parent. Setpgid is not an alternative either, since taking the child out of +// the foreground process group is what would break the Ctrl+C that works +// today. SIGKILL stays uncovered, as it must: nothing in the process can +// observe it. +func shutdownContext(parent context.Context) (context.Context, context.CancelFunc) { + return signal.NotifyContext(parent, shutdownSignals...) +} + +// Run starts the agent: resolve where state lives, make sure a server is +// reachable, pick a model, then hand off to nib. +func Run(ctx context.Context, opts Options) error { + ctx, stop := shutdownContext(ctx) + defer stop() + + p, err := prepare(ctx, opts, isTerminal(opts.In)) + if err != nil { + return err + } + // A server this process started belongs to this session, and Stop is + // nil-safe and idempotent, so one defer covers both cases and costs nothing + // when runSession has already stopped it. + defer p.server.Stop() + + return runSession(ctx, p.server, func(ctx context.Context) error { + return runAgent(ctx, p.dir, p.model, opts) + }) +} + +// runSession hands the terminal to agent, and stops a server started for this +// session as soon as the context is cancelled rather than when agent returns. +// +// The difference matters because the deferred Stop in Run is only reached once +// agent returns, and how long that takes is nib's business rather than ours. +// nib v0.5.1 does unwind the TUI on a cancelled context, so it does return; a +// SIGHUP no longer leaves the interface on screen with the server behind it, +// which it did before, when bubbletea's own SIGINT and SIGTERM handler was the +// only thing that ever quit the program and registering for SIGHUP had removed +// the default disposition that used to end the process. Watching the context +// keeps the guarantee independent of what the agent does with it. +func runSession(ctx context.Context, server *StartedServer, agent func(context.Context) error) error { + returned := make(chan struct{}) + defer close(returned) + + go func() { + select { + case <-ctx.Done(): + server.Stop() + case <-returned: + } + }() + + return agent(ctx) +} + +// preparation is what the agent needs once the environment is ready: where its +// state lives, which model to talk to, and the server this process started on +// the user's behalf, if any. +type preparation struct { + dir string + model string + server *StartedServer +} + +// prepare does everything that has to happen before the agent takes over the +// terminal. It is split out of Run because all of it is testable and none of +// what follows is: once app.Run has the terminal there is no seam left. +// +// interactive says whether there is a user to prompt. It is a parameter rather +// than a second read of opts.In so the prompts can be driven over a pipe. +func prepare(ctx context.Context, opts Options, interactive bool) (_ *preparation, err error) { + dir, dirErr := StateDir(opts.StateDir) + if dirErr != nil { + return nil, dirErr + } + if err := EnsureStateDir(dir, opts.BaseURL); err != nil { + return nil, err + } + + if isLocalOnlyArgs(opts.Args) { + return &preparation{dir: dir}, nil + } + + // One prompter for every question this run asks; see its doc comment for + // why the reader cannot be rebuilt per question. + var prompts *prompter + if interactive { + prompts = newPrompter(opts.In, opts.ErrOut) + } + + var started *StartedServer + defer func() { + // Nothing after the spawn may leave a server behind: the caller only + // learns about it through a successful return. + if err != nil { + started.Stop() + } + }() + + models, err := probeModels(ctx, opts) + if err != nil { + if errors.Is(err, ErrUnauthorized) { + return nil, fmt.Errorf("the LocalAI server at %s rejected the API key. Pass --api-key or set LOCALAI_API_KEY", opts.Endpoint) + } + if !errors.Is(err, ErrUnreachable) { + return nil, err + } + + var confirm Confirmer + if interactive { + confirm = prompts.yesNo + } + var startErr error + started, startErr = OfferToStart(ctx, StartOptions{ + Endpoint: opts.Endpoint, + Confirm: confirm, + Stderr: opts.ErrOut, + }) + if startErr != nil { + err = startErr + if errors.Is(startErr, ErrDeclined) { + err = fmt.Errorf("no LocalAI server at %s. Start one with 'local-ai run', or point elsewhere with --endpoint", opts.Endpoint) + } + return nil, err + } + say(opts.ErrOut, "Started a temporary LocalAI server; it stops when you exit. Use 'local-ai run' for a persistent one.\n") + + if models, err = probeModels(ctx, opts); err != nil { + return nil, err + } + } + + var chooser ModelChooser + if interactive { + chooser = prompts.choose + } + model, err := ResolveModel(ModelRequest{ + Flag: opts.Model, + Configured: configuredModel(dir), + Available: models, + StateDir: dir, + Choose: chooser, + Notify: func(message string) { say(opts.ErrOut, "%s\n", message) }, + }) + if err != nil { + return nil, err + } + + return &preparation{dir: dir, model: model, server: started}, nil +} + +func runAgent(ctx context.Context, dir, model string, opts Options) error { + return app.Run(ctx, agentOptions(dir, model, opts)) +} + +// agentOptions builds the request handed to nib. It is split out of runAgent +// because app.Run takes the terminal and cannot be called from a test, while +// what is asked of it is exactly the part worth pinning. +// +// The stream fields are the interesting ones, and they are not symmetric. +// +// nib reads a non-nil stream as "the embedder wants this used", and refuses +// every mode but --cli when such a stream is not a terminal, because the +// full-screen interface renders on /dev/tty and would otherwise ignore it in +// silence. Nil means "not injected": nib falls back to the process stream and +// behaves as standalone nib does. +// +// Stdin is passed through as it comes. A piped or redirected stdin really is +// ignored by the interface, so the refusal is the honest answer there, and it +// is the one users meet: 'echo q | local-ai chat' says to re-run with --cli +// rather than opening a full-screen session that will never read the question. +// +// Stdout is different, and the process stream is deliberately sent as nil. The +// interface does write to stdout even when it is a pipe: that is the whole of +// nib's shell-capture idiom, out=$(local-ai chat --height 50%), which is what +// the Ctrl+Space widget emitted by --init is built on. Injecting os.Stdout +// there would refuse the widget for a stream nib was going to use anyway. +// +// The test is identity with os.Stdout rather than whether it happens to be a +// terminal, which means a shell redirect goes the same way as the widget: +// 'local-ai chat > out.txt' no longer refuses either, and renders on /dev/tty +// with the capture line landing in the file. That is not a second decision, it +// is the same one. Both are the process stdout as the shell handed it over, +// differing only in being a pipe rather than a regular file, which nib's gate +// does not look at and should not. Refusing one would refuse the other. +// +// What stays injected, and so stays subject to the refusal, is a writer some +// in-process caller chose for itself rather than inherited: a bytes.Buffer, or +// an *os.File it opened. The specs rely on that. +// +// Stderr is never gated by nib, so it is passed through unchanged. +// +// The config values go through Overrides rather than Defaults, and that is not +// a detail. Defaults are seeds: they sit BENEATH the config file, so the file +// silently undoes them. Everything here is a decision this invocation already +// made on the user's behalf, and a flag that the file can undo is not a flag. +// It was not a rare case either, since EnsureStateDir writes base_url on the +// first run and an interactive choice writes model, so from the second run on +// the file carried a value for both and --endpoint and --model did nothing. +// +// The one asymmetry to plan around is that nib cannot tell "set to the zero +// value" from "not set", so an override only ever raises a field. --yolo can +// turn approval off, but nothing on the command line can turn it back on over +// an approval_mode: auto in the file; that needs a config edit. Same shape for +// the strings, which is what makes an unset --api-key or --trace-dir leave the +// file's value standing, as it should. +// +// nib's own --trace-dir and --yolo, and their NIB_TRACE_DIR and NIB_YOLO twins, +// are resolved after the config load and so still outrank these. That is +// deliberate upstream: they are instructions to nib rather than ambient +// environment. +func agentOptions(dir, model string, opts Options) app.Options { + // Model is the model this run resolved, which already prefers --model and + // falls back to the file's own model, so the override restates the file's + // value rather than fighting it whenever no flag was given. + // + // BaseURL is the endpoint this run probed, offered to start a server for, + // and seeded the config with. Handing nib a different one is precisely the + // split that made --endpoint a no-op, so the agent talks to the server + // LocalAI checked. Pointing somewhere else for good is LOCALAI_CHAT_ENDPOINT + // or --endpoint, not a hand-edited base_url the probe never reads. + // + // APIKey and TraceDir are the flags as given, empty when they were not, and + // an empty override leaves the file alone. TraceDir is runtime-only in nib + // (yaml:"-"), so no file value exists for it to beat today; it belongs here + // with the other flags rather than one rung down for a reason that could + // quietly stop being true. + overrides := nibtypes.Config{ + Model: model, + APIKey: opts.APIKey, + BaseURL: opts.BaseURL, + TraceDir: opts.TraceDir, + } + if opts.Yolo { + overrides.ApprovalMode = "auto" + } + + return app.Options{ + Args: opts.Args, + ProgramName: "local-ai chat", + BaseDir: dir, + Overrides: overrides, + SkipSetup: true, + SkipBareEnv: true, + Stdin: opts.In, + Stdout: ownStdout(opts.Out), + Stderr: opts.ErrOut, + } +} + +// ownStdout reports the writer as nib's own rather than as an injected one when +// it is the process stdout, by answering nil for it. See agentOptions for why +// that distinction is the difference between a working Ctrl+Space widget and a +// refused one. +func ownStdout(w io.Writer) io.Writer { + if f, ok := w.(*os.File); ok && f == os.Stdout { + return nil + } + return w +} + +// defaultProbeTimeout bounds a check of the server. Listing models is cheap, +// so this is long enough that a loaded server is never given up on and short +// enough that a hung one does not leave the user staring at nothing. +const defaultProbeTimeout = 30 * time.Second + +// probeModels lists what the endpoint offers, under a budget. +func probeModels(ctx context.Context, opts Options) ([]string, error) { + timeout := opts.ProbeTimeout + if timeout <= 0 { + timeout = defaultProbeTimeout + } + // A real deadline rather than a cancel plus a timer. Probe reads + // context.Canceled as "the caller gave up", which is a statement about the + // caller and not about the endpoint, and only a deadline as "nothing + // answered in time". Expiring the budget as a cancellation would stop + // ErrUnreachable firing for precisely the hung servers that the offer to + // start one exists for. + probeCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + return Probe(probeCtx, opts.BaseURL, opts.APIKey) +} + +// isLocalOnlyArgs reports whether the forwarded arguments do their work +// without ever reaching a model, in which case demanding a running server (and +// offering to start one) would be an obstacle rather than a service. +// +// Two groups qualify. The management subcommands edit nib's own state: plugin, +// skill, and the mcp verbs that add or remove configured servers, which is +// asked of nib rather than restated, because bare 'mcp' and its transport +// flags do serve the agent and do need a model. The other group is the flags +// that only print something, above all --init: its shell snippet goes into an +// rc file, typically long before any server exists. +func isLocalOnlyArgs(args []string) bool { + if len(args) == 0 { + return false + } + // A scan rather than a look at args[0]: the mode flags this command + // translates are prepended, so --init is not necessarily first. Positional + // text cannot be mistaken for a flag here, since nib ignores what is left + // after flag parsing. + for _, a := range args { + switch { + case a == "--init", a == "-init", strings.HasPrefix(a, "--init="), strings.HasPrefix(a, "-init="): + return true + case a == "--version", a == "-version": + return true + } + } + switch args[0] { + case "plugin", "skill": + return true + case "mcp": + return len(args) >= 2 && nibcmd.IsMCPManageSubcommand(args[1]) + } + return false +} + +// configuredModel reads the model already recorded in the agent config, if any. +func configuredModel(dir string) string { + cfg := nibconfig.LoadWith(nibconfig.LoadOptions{BaseDir: dir, SkipBareEnv: true}) + return cfg.Model +} + +func isTerminal(in io.Reader) bool { + f, ok := in.(*os.File) + return ok && term.IsTerminal(int(f.Fd())) +} + +// say writes a line of interactive chatter: a question, or a notice about +// something that did not stop the session. A write that fails is not worth +// failing over, and when the terminal really is gone the read that follows the +// question says so. +func say(w io.Writer, format string, args ...any) { + _, _ = fmt.Fprintf(w, format, args...) +} + +// prompter asks this run's questions on the user's terminal. +// +// It owns the buffered reader rather than wrapping opts.In per question, +// because bufio reads ahead: a throwaway reader for the "start a server?" +// question swallows the model choice that was typed behind it, and the next +// question then sees EOF. A real run asks both, one after the other. +type prompter struct { + in *bufio.Reader + out io.Writer +} + +func newPrompter(in io.Reader, out io.Writer) *prompter { + return &prompter{in: bufio.NewReader(in), out: out} +} + +// yesNo satisfies Confirmer. Anything that is not an explicit yes is a no, so +// a closed stream declines rather than proceeding on the user's behalf. +func (p *prompter) yesNo(question string) (bool, error) { + say(p.out, "%s [y/N]: ", question) + line, err := p.in.ReadString('\n') + if err != nil && !errors.Is(err, io.EOF) { + return false, fmt.Errorf("reading the answer: %w", err) + } + switch strings.ToLower(strings.TrimSpace(line)) { + case "y", "yes": + return true, nil + } + return false, nil +} + +// choose satisfies ModelChooser. It answers with a list index rather than with +// what the user typed, so the result can only ever be one of the models it was +// offered: a model name is not something to accept unvalidated here, since +// ResolveModel persists whatever comes back and every later run then starts +// against it. +func (p *prompter) choose(models []string) (string, error) { + if len(models) == 0 { + return "", errors.New("there is nothing to choose from") + } + say(p.out, "Several models are available:\n") + for i, m := range models { + say(p.out, " %d) %s\n", i+1, m) + } + say(p.out, "Pick one [1-%d]: ", len(models)) + + line, err := p.in.ReadString('\n') + if err != nil && !errors.Is(err, io.EOF) { + return "", fmt.Errorf("reading the choice: %w", err) + } + answer := strings.TrimSpace(line) + n, err := strconv.Atoi(answer) + if err != nil || n < 1 || n > len(models) { + return "", fmt.Errorf("not a valid choice: %q. Pick a number between 1 and %d, or pass --model", answer, len(models)) + } + return models[n-1], nil +} diff --git a/core/cli/chat/run_test.go b/core/cli/chat/run_test.go new file mode 100644 index 000000000..e334e0a0b --- /dev/null +++ b/core/cli/chat/run_test.go @@ -0,0 +1,629 @@ +package chat + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "syscall" + "time" + + "github.com/mudler/nib/app" + nibconfig "github.com/mudler/nib/config" + nibtypes "github.com/mudler/nib/types" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// modelServer answers /v1/models with the given ids, as LocalAI does. +func modelServer(ids ...string) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + data := make([]map[string]string, 0, len(ids)) + for _, id := range ids { + data = append(data, map[string]string{"id": id, "object": "model"}) + } + w.Header().Set("Content-Type", "application/json") + Expect(json.NewEncoder(w).Encode(map[string]any{"object": "list", "data": data})).To(Succeed()) + })) +} + +var _ = Describe("prepare", func() { + var ( + dir string + errOut *bytes.Buffer + ) + + BeforeEach(func() { + dir = GinkgoT().TempDir() + errOut = &bytes.Buffer{} + }) + + // optionsFor points a run at srv, with no input to read: the default is a + // session nobody can be asked anything in. + optionsFor := func(srv *httptest.Server) Options { + endpoint := "http://127.0.0.1:0" + base := endpoint + "/v1" + if srv != nil { + endpoint, base = srv.URL, srv.URL+"/v1" + } + return Options{ + Endpoint: endpoint, + BaseURL: base, + StateDir: dir, + In: strings.NewReader(""), + Out: &bytes.Buffer{}, + ErrOut: errOut, + } + } + + It("uses the only model the server offers", func() { + srv := modelServer("the-only-model") + defer srv.Close() + + p, err := prepare(context.Background(), optionsFor(srv), false) + Expect(err).ToNot(HaveOccurred()) + Expect(p.model).To(Equal("the-only-model")) + Expect(p.dir).To(Equal(dir)) + Expect(p.server).To(BeNil(), "nothing was started, so nothing is owned") + }) + + It("seeds the agent config with the endpoint on first run", func() { + srv := modelServer("m") + defer srv.Close() + + _, err := prepare(context.Background(), optionsFor(srv), false) + Expect(err).ToNot(HaveOccurred()) + + data, err := os.ReadFile(ConfigPath(dir)) + Expect(err).ToNot(HaveOccurred()) + Expect(string(data)).To(ContainSubstring(srv.URL + "/v1")) + }) + + It("lets --model win over what the server offers", func() { + srv := modelServer("a", "b") + defer srv.Close() + + opts := optionsFor(srv) + opts.Model = "not-listed-yet" + p, err := prepare(context.Background(), opts, false) + Expect(err).ToNot(HaveOccurred()) + Expect(p.model).To(Equal("not-listed-yet")) + }) + + It("advises about the API key when the server rejects it", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + })) + defer srv.Close() + + _, err := prepare(context.Background(), optionsFor(srv), false) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("--api-key")) + Expect(err.Error()).To(ContainSubstring(srv.URL)) + }) + + // Not interactive means nobody can answer the offer, so the advice has to + // stand on its own. + It("advises how to start a server when none is reachable", func() { + srv := modelServer() + url := srv.URL + srv.Close() // nothing is listening now + + opts := optionsFor(nil) + opts.Endpoint, opts.BaseURL = url, url+"/v1" + _, err := prepare(context.Background(), opts, false) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("local-ai run")) + Expect(err.Error()).To(ContainSubstring(url)) + }) + + // A server that accepts the connection and then never replies is the case + // the offer to start one exists for, so the budget has to expire as a + // deadline: Probe reads a cancellation as "the caller gave up" and refuses + // to call the endpoint unreachable on the strength of it. + It("treats a server that never answers as one that is not there", func(ctx SpecContext) { + release := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case <-release: + case <-r.Context().Done(): + } + })) + defer srv.Close() + defer close(release) + + opts := optionsFor(srv) + opts.ProbeTimeout = 100 * time.Millisecond + _, err := prepare(context.Background(), opts, false) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("local-ai run"), "want the offer-a-server advice, got %v", err) + }, SpecTimeout(30*time.Second)) + + It("asks which model to use and remembers the answer", func() { + srv := modelServer("zeta", "alpha") + defer srv.Close() + + opts := optionsFor(srv) + opts.In = strings.NewReader("2\n") + p, err := prepare(context.Background(), opts, true) + Expect(err).ToNot(HaveOccurred()) + // The list is sorted before it is shown, so 2 is zeta, not the second + // thing the server happened to name. + Expect(p.model).To(Equal("zeta")) + Expect(errOut.String()).To(ContainSubstring("1) alpha")) + Expect(errOut.String()).To(ContainSubstring("2) zeta")) + + data, err := os.ReadFile(ConfigPath(dir)) + Expect(err).ToNot(HaveOccurred()) + Expect(string(data)).To(ContainSubstring("zeta")) + }) + + // The choice is prompted for once and remembered. When remembering it fails + // the user is about to be asked again on every future run, so they have to + // be told here: a log line is invisible at the default log level. + It("says so on the prompt when the choice cannot be remembered", func() { + srv := modelServer("zeta", "alpha") + defer srv.Close() + + // A directory where the config file belongs: writable state dir, + // unwritable config, on any platform and as any user. + Expect(os.MkdirAll(ConfigPath(dir), 0o700)).To(Succeed()) + + opts := optionsFor(srv) + opts.In = strings.NewReader("1\n") + p, err := prepare(context.Background(), opts, true) + + // Failing to remember the choice must not cost the user their session. + Expect(err).ToNot(HaveOccurred()) + Expect(p.model).To(Equal("alpha")) + Expect(errOut.String()).To(ContainSubstring("could not be saved"), "the user has to learn they will be asked again") + }) + + It("does not ask again once a model is recorded", func() { + srv := modelServer("zeta", "alpha") + defer srv.Close() + + Expect(PersistModel(dir, "alpha")).To(Succeed()) + + opts := optionsFor(srv) + opts.In = strings.NewReader("") // an answer would have nothing to read + p, err := prepare(context.Background(), opts, true) + Expect(err).ToNot(HaveOccurred()) + Expect(p.model).To(Equal("alpha")) + Expect(errOut.String()).To(BeEmpty()) + }) + + It("says what to install when the server has no models", func() { + srv := modelServer() + defer srv.Close() + + _, err := prepare(context.Background(), optionsFor(srv), false) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("models install")) + }) + + Describe("arguments that only touch local state", func() { + unreachable := func(args ...string) Options { + opts := optionsFor(nil) // port 0: nothing can ever answer here + opts.Args = args + return opts + } + + DescribeTable("skips the server entirely", + func(args ...string) { + p, err := prepare(context.Background(), unreachable(args...), false) + Expect(err).ToNot(HaveOccurred()) + Expect(p.model).To(BeEmpty()) + Expect(p.server).To(BeNil()) + }, + Entry("plugin", "plugin", "list"), + Entry("skill", "skill", "list"), + Entry("mcp add", "mcp", "add", "srv"), + Entry("mcp list", "mcp", "list"), + // The shell snippet is what a user puts in their rc file, long + // before any server exists. + Entry("the shell integration script", "--init", "zsh"), + Entry("the version", "--version"), + ) + + // Bare 'mcp' and its transport flags serve the agent over MCP, so they + // need a model like any other session. Only the verbs that edit the + // configured servers are local. + DescribeTable("still needs a server", + func(args ...string) { + _, err := prepare(context.Background(), unreachable(args...), false) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("local-ai run")) + }, + Entry("mcp over stdio", "mcp", "--stdio"), + Entry("bare mcp", "mcp"), + ) + }) + + // A reader per question would read ahead into a buffer it then discards, so + // the second question would see EOF whenever both answers were typed ahead. + // That is the shape of a real run: the offer to start a server is followed + // by the model prompt. + It("keeps reading answers from the same stream across questions", func() { + out := &bytes.Buffer{} + p := newPrompter(strings.NewReader("y\n2\n"), out) + + yes, err := p.yesNo("Start one now?") + Expect(err).ToNot(HaveOccurred()) + Expect(yes).To(BeTrue()) + + chosen, err := p.choose([]string{"alpha", "zeta"}) + Expect(err).ToNot(HaveOccurred()) + Expect(chosen).To(Equal("zeta")) + }) + + // Whatever the chooser returns is persisted and used for every later run, + // so an answer that is not one of the offered models must never come back + // as one. + Describe("the model prompt", func() { + offered := []string{"alpha", "zeta"} + + DescribeTable("refuses an answer that is not one of the numbers shown", + func(answer string) { + chosen, err := newPrompter(strings.NewReader(answer), &bytes.Buffer{}).choose(offered) + Expect(err).To(HaveOccurred()) + Expect(chosen).To(BeEmpty()) + }, + Entry("nothing at all", ""), + Entry("a blank line", "\n"), + Entry("only spaces", " \n"), + Entry("zero", "0\n"), + Entry("past the end", "3\n"), + Entry("negative", "-1\n"), + Entry("a model name", "zeta\n"), + Entry("a number with a suffix", "1x\n"), + ) + + It("says how to answer when the answer was not a number", func() { + _, err := newPrompter(strings.NewReader("banana\n"), &bytes.Buffer{}).choose(offered) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("between 1 and 2")) + Expect(err.Error()).To(ContainSubstring("--model")) + }) + + It("returns the model shown against the number", func() { + chosen, err := newPrompter(strings.NewReader("1\n"), &bytes.Buffer{}).choose(offered) + Expect(err).ToNot(HaveOccurred()) + Expect(chosen).To(Equal("alpha")) + }) + + It("refuses to ask when there is nothing to offer", func() { + chosen, err := newPrompter(strings.NewReader("1\n"), &bytes.Buffer{}).choose(nil) + Expect(err).To(HaveOccurred()) + Expect(chosen).To(BeEmpty()) + }) + }) + + // A server started for this session is stopped by a deferred call, which a + // signal skips: the process dies where it stands and leaves 'local-ai run' + // reparented to init. + Describe("shutdown signals", func() { + It("ends the session when the terminal goes away", func() { + ctx, stop := shutdownContext(context.Background()) + defer stop() + + self, err := os.FindProcess(os.Getpid()) + Expect(err).ToNot(HaveOccurred()) + Expect(self.Signal(syscall.SIGHUP)).To(Succeed()) + + Eventually(ctx.Done()).WithTimeout(5 * time.Second).Should(BeClosed()) + Expect(ctx.Err()).To(MatchError(context.Canceled)) + }) + + // SIGINT and SIGTERM cannot be delivered here to prove the same thing: + // Ginkgo registers for both to abort the suite, and a signal goes to + // every registered listener. + It("also listens for an interrupt and a terminate", func() { + Expect(shutdownSignals).To(ContainElements(os.Signal(os.Interrupt), os.Signal(syscall.SIGTERM))) + }) + }) + + // Cancelling the context does unwind nib's TUI since v0.5.1, but how long + // that takes is nib's business, and the deferred Stop in Run is only reached + // once the agent returns. A server this process started is ours to end, so + // the guarantee is made here instead, where it does not depend on the agent + // at all. Before v0.5.1 there was no guarantee to be had on the SIGHUP path: + // bubbletea's own SIGINT and SIGTERM handler was the only thing that ever + // quit the program, and registering for SIGHUP took away the default + // disposition that used to end the process. + Describe("runSession", func() { + It("stops the session's server on cancellation, without waiting for the agent", func() { + server, proc := stoppableServer() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + err := runSession(ctx, server, func(ctx context.Context) error { + cancel() + Eventually(func() int32 { return proc.interrupts.Load() }). + WithTimeout(5 * time.Second). + Should(BeNumerically(">", 0), "the server has to be stopped while the agent is still running") + return nil + }) + Expect(err).ToNot(HaveOccurred()) + Expect(proc.lastSignal.Load()).To(Equal(os.Interrupt)) + }) + + It("leaves the server alone for as long as the session lasts", func() { + server, proc := stoppableServer() + + Expect(runSession(context.Background(), server, func(context.Context) error { + return nil + })).To(Succeed()) + Expect(proc.interrupts.Load()).To(BeZero()) + Expect(proc.kills.Load()).To(BeZero()) + }) + + It("returns what the agent returned", func() { + failed := errors.New("the agent gave up") + server, _ := stoppableServer() + + Expect(runSession(context.Background(), server, func(context.Context) error { + return failed + })).To(MatchError(failed)) + }) + + // Most sessions run against a server the user already had, and there is + // nothing to stop then. + It("copes with a session that started no server", func() { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + Expect(runSession(ctx, nil, func(context.Context) error { + return nil + })).To(Succeed()) + }) + }) + + // Which streams reach nib decides two user-visible behaviours at once, and + // they pull in opposite directions, so both are pinned here rather than left + // to whoever next edits the literal. + // + // nib refuses every mode but --cli when a stream it was handed is not a + // terminal. That refusal is wanted for stdin, where it is what tells someone + // piping a question to re-run with --cli. It is not wanted for the process + // stdout, where it would refuse the Ctrl+Space widget that --init emits: + // out=$(local-ai chat --height 50%) puts a pipe on stdout by construction, + // and writing the chosen command into that pipe is the entire point. + Describe("agentOptions", func() { + // optionsWithStreams is a request that differs from the next only in + // what it was told to read and write. + optionsWithStreams := func(in io.Reader, out, errOut io.Writer) Options { + return Options{ + BaseURL: "http://127.0.0.1:8080/v1", + In: in, + Out: out, + ErrOut: errOut, + } + } + + Describe("stdout", func() { + // The regression this exists to catch: reinstating + // 'Stdout: opts.Out' breaks Ctrl+Space and nothing else notices. + It("hands nib nothing for the process stdout, so the capture widget is not refused", func() { + o := agentOptions(dir, "a-model", optionsWithStreams(os.Stdin, os.Stdout, os.Stderr)) + Expect(o.Stdout).To(BeNil(), "injecting os.Stdout is what refuses out=$(local-ai chat)") + }) + + It("keeps a stdout the caller chose, which the refusal still guards", func() { + out := &bytes.Buffer{} + o := agentOptions(dir, "a-model", optionsWithStreams(os.Stdin, out, os.Stderr)) + Expect(o.Stdout).To(BeIdenticalTo(out)) + }) + + // Being an *os.File is not what makes a stream nib's own; being the + // process stdout is. This is a file an in-process caller opened for + // itself, not one a shell redirect handed over as stdout, which + // still arrives as os.Stdout and is still nil-ed. It was never going + // to receive the interface, so it stays injected and stays refused. + It("keeps a file that is not the process stdout", func() { + f, err := os.CreateTemp(GinkgoT().TempDir(), "captured") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(f.Close) + + o := agentOptions(dir, "a-model", optionsWithStreams(os.Stdin, f, os.Stderr)) + Expect(o.Stdout).To(BeIdenticalTo(f)) + }) + }) + + Describe("stdin", func() { + // The opposite regression: nilling stdin the way stdout is nilled + // would silently drop the refusal that names --cli. + It("hands the process stdin over, so a piped session is still refused", func() { + o := agentOptions(dir, "a-model", optionsWithStreams(os.Stdin, os.Stdout, os.Stderr)) + Expect(o.Stdin).To(BeIdenticalTo(os.Stdin)) + }) + + It("hands over a stdin the caller chose", func() { + in := strings.NewReader("a question") + o := agentOptions(dir, "a-model", optionsWithStreams(in, os.Stdout, os.Stderr)) + Expect(o.Stdin).To(BeIdenticalTo(in)) + }) + }) + + // nib gates stdin and stdout and nothing else, so there is no reason to + // hide the error stream from it. + It("hands the error stream over whatever it is", func() { + errOut := &bytes.Buffer{} + o := agentOptions(dir, "a-model", optionsWithStreams(os.Stdin, os.Stdout, errOut)) + Expect(o.Stderr).To(BeIdenticalTo(errOut)) + + o = agentOptions(dir, "a-model", optionsWithStreams(os.Stdin, os.Stdout, os.Stderr)) + Expect(o.Stderr).To(BeIdenticalTo(os.Stderr)) + }) + + It("names the command a user would type, not the binary nib ships as", func() { + o := agentOptions(dir, "a-model", optionsWithStreams(os.Stdin, os.Stdout, os.Stderr)) + Expect(o.ProgramName).To(Equal("local-ai chat"), + "the --init widget invokes this name, so a user has to be able to run it") + }) + + It("carries the resolved session through to nib", func() { + opts := optionsWithStreams(os.Stdin, os.Stdout, os.Stderr) + opts.Args = []string{"--cli"} + opts.APIKey = "a-key" + opts.TraceDir = "/traces" + + o := agentOptions(dir, "the-model", opts) + Expect(o.Args).To(Equal([]string{"--cli"})) + Expect(o.BaseDir).To(Equal(dir)) + Expect(o.Overrides.Model).To(Equal("the-model")) + Expect(o.Overrides.APIKey).To(Equal("a-key")) + Expect(o.Overrides.BaseURL).To(Equal("http://127.0.0.1:8080/v1")) + Expect(o.Overrides.TraceDir).To(Equal("/traces")) + // The model and the server are settled before nib starts, and the + // bare MODEL and API_KEY variables belong to some other tool. + Expect(o.SkipSetup).To(BeTrue()) + Expect(o.SkipBareEnv).To(BeTrue()) + }) + + // Defaults sit beneath the config file. Anything routed through them is + // accepted from the command line and then thrown away the moment the + // file carries the same key, which is the normal state rather than an + // edge case. Nothing this command resolves belongs there, so the channel + // stays empty and this says so: it is what fails if the block is moved + // back a rung. + It("seeds nothing, because a seed is not a flag", func() { + opts := optionsWithStreams(os.Stdin, os.Stdout, os.Stderr) + opts.APIKey = "a-key" + opts.TraceDir = "/traces" + opts.Yolo = true + + Expect(agentOptions(dir, "the-model", opts).Defaults).To(Equal(nibtypes.Config{}), + "Defaults lose to the config file, so a value placed there is a flag that does nothing") + }) + + It("asks for automatic approval only when --yolo was given", func() { + opts := optionsWithStreams(os.Stdin, os.Stdout, os.Stderr) + Expect(agentOptions(dir, "a-model", opts).Overrides.ApprovalMode).To(BeEmpty()) + + opts.Yolo = true + Expect(agentOptions(dir, "a-model", opts).Overrides.ApprovalMode).To(Equal("auto")) + }) + + // The specs above pin what is handed over. These pin what nib does with + // it, which is the part that was wrong: every value below reached + // app.Options intact and was then discarded by the config load, so a + // spec that stops at the struct cannot see the bug. Resolving the config + // the way app.Run resolves it can. + Describe("the config nib actually resolves", func() { + // writeConfig puts a config file where nib will read it, with values + // that disagree with every flag under test. + writeConfig := func(body string) { + Expect(os.WriteFile(ConfigPath(dir), []byte(body), 0o600)).To(Succeed()) + } + + // resolve loads the config exactly as app.Run does, so the precedence + // under test is nib's own rather than a restatement of it here. + resolve := func(o app.Options) nibtypes.Config { + return nibconfig.LoadWith(nibconfig.LoadOptions{ + BaseDir: o.BaseDir, + Defaults: o.Defaults, + Overrides: o.Overrides, + SkipBareEnv: o.SkipBareEnv, + }) + } + + It("sends the requests to the endpoint the flag named, not the one on disk", func() { + writeConfig("base_url: http://127.0.0.1:9999/v1\n") + + opts := optionsWithStreams(os.Stdin, os.Stdout, os.Stderr) + opts.BaseURL = "http://127.0.0.1:8080/v1" + + cfg := resolve(agentOptions(dir, "a-model", opts)) + Expect(cfg.BaseURL).To(Equal("http://127.0.0.1:8080/v1"), + "--endpoint probed 8080; every turn has to go there too") + }) + + It("uses the model the flag named, not the one the picker recorded", func() { + writeConfig("model: recorded-model\n") + + cfg := resolve(agentOptions(dir, "flag-model", optionsWithStreams(os.Stdin, os.Stdout, os.Stderr))) + Expect(cfg.Model).To(Equal("flag-model")) + }) + + It("uses the key the flag named, not the one nib saved", func() { + writeConfig("api_key: saved-key\n") + + opts := optionsWithStreams(os.Stdin, os.Stdout, os.Stderr) + opts.APIKey = "flag-key" + + cfg := resolve(agentOptions(dir, "a-model", opts)) + Expect(cfg.APIKey).To(Equal("flag-key")) + }) + + It("turns approval off for --yolo even when the file demands it", func() { + writeConfig("approval_mode: prompt\n") + + opts := optionsWithStreams(os.Stdin, os.Stdout, os.Stderr) + opts.Yolo = true + + cfg := resolve(agentOptions(dir, "a-model", opts)) + Expect(cfg.ApprovalMode).To(Equal("auto")) + }) + + // The other half of the same rule, and the reason an unset flag is + // not a demand for the empty string: an override only ever raises a + // field, so what the user configured survives a run that said + // nothing about it. + It("leaves what the file configured alone when no flag was given", func() { + writeConfig("api_key: saved-key\napproval_mode: prompt\n") + + cfg := resolve(agentOptions(dir, "a-model", optionsWithStreams(os.Stdin, os.Stdout, os.Stderr))) + Expect(cfg.APIKey).To(Equal("saved-key")) + Expect(cfg.ApprovalMode).To(Equal("prompt")) + }) + }) + }) + + // nib reports its own failures on the error stream and returns nothing but + // a status, so anything that reaches here as one has already been explained + // once. The refusal to open a full-screen session on a stdin that cannot be + // read is the one users meet: 'echo q | local-ai chat' names --cli, and a + // second message on top would bury the fix. + Describe("ExitStatus", func() { + It("recognises a status the agent already explained", func() { + code, reported := ExitStatus(app.ExitError{Code: 2}) + Expect(reported).To(BeTrue()) + Expect(code).To(Equal(2)) + }) + + It("finds one that has been wrapped", func() { + code, reported := ExitStatus(fmt.Errorf("running the agent: %w", app.ExitError{Code: 1})) + Expect(reported).To(BeTrue()) + Expect(code).To(Equal(1)) + }) + + It("leaves an ordinary failure to be reported", func() { + _, reported := ExitStatus(errors.New("no LocalAI server at http://127.0.0.1:8080")) + Expect(reported).To(BeFalse()) + }) + + It("says nothing about a run that succeeded", func() { + _, reported := ExitStatus(nil) + Expect(reported).To(BeFalse()) + }) + }) + + It("reports a state dir it cannot create", func() { + blocked := filepath.Join(dir, "a-file") + Expect(os.WriteFile(blocked, []byte("not a dir"), 0o600)).To(Succeed()) + + opts := optionsFor(nil) + opts.StateDir = filepath.Join(blocked, "chat") + _, err := prepare(context.Background(), opts, false) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("agent state dir")) + }) +}) diff --git a/core/cli/chat/server.go b/core/cli/chat/server.go new file mode 100644 index 000000000..3d23d09c7 --- /dev/null +++ b/core/cli/chat/server.go @@ -0,0 +1,276 @@ +package chat + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "os" + "os/exec" + "strings" + "sync" + "time" + + "github.com/mudler/LocalAI/pkg/httpclient" +) + +// ErrDeclined means no server was started, either because the session is not +// interactive or because the user said no. +var ErrDeclined = errors.New("no server started") + +// errServerExited means the process we spawned died before it ever reported +// ready, so there is no point in polling out the rest of the budget. +var errServerExited = errors.New("the LocalAI server exited before it became ready") + +const ( + // defaultReadyTimeout bounds the wait for a freshly spawned server. A cold + // start probes hardware and may pull a backend, so the budget is generous. + defaultReadyTimeout = 2 * time.Minute + // readyPollInterval is how long to wait between readiness polls. + readyPollInterval = 500 * time.Millisecond + // readyProbeTimeout bounds a single readiness request, so one connection + // that hangs cannot swallow the whole budget. + readyProbeTimeout = 5 * time.Second + // shutdownGrace is how long a server we started gets to unload models and + // stop its backends after SIGINT before it is killed outright. + shutdownGrace = 10 * time.Second + // childOutputDrainDelay bounds how long cmd.Wait keeps copying the child's + // output after the child itself has exited. + // + // This is not a theoretical guard for LocalAI. 'local-ai run' spawns backend + // subprocesses, and they inherit the write end of the pipe exec created for + // the child's stderr. A backend that outlives its parent holds that pipe + // open, so an unbounded cmd.Wait would block on the copy goroutine long + // after the server itself is gone: exited would never close, Stop would burn + // its whole grace period even on a clean shutdown, and the waiter goroutine + // would leak. + // + // The value is long enough that a legitimate final burst of logs is never + // truncated even on a loaded machine, where the copy itself takes + // microseconds. It must stay strictly below shutdownGrace: at or above it, + // every wedged-pipe shutdown would exhaust the grace period and then SIGKILL + // a process that had already exited cleanly. + childOutputDrainDelay = 5 * time.Second +) + +// Confirmer asks a yes/no question. Nil means the session is not interactive. +type Confirmer func(question string) (bool, error) + +// StartOptions configures OfferToStart. +type StartOptions struct { + // Endpoint is the address the user expected a server on, used in the + // question and polled for readiness. This is the endpoint root, not the + // /v1 API base URL: readiness is served at the root. + Endpoint string + // Confirm asks whether to start a server. Nil means never start. + Confirm Confirmer + // Stderr receives the child's output. + Stderr io.Writer + // Executable overrides the binary to run. Empty means os.Executable(). + Executable string + // ReadyTimeout bounds the wait for readiness. Zero means defaultReadyTimeout. + ReadyTimeout time.Duration +} + +// StartedServer is a server this process started and is responsible for. +type StartedServer struct { + // exited is closed once the child has been reaped. One background waiter + // owns cmd.Wait: it may only be called once, and it is what closes the + // pipes exec created for Stdout/Stderr and joins the goroutines copying + // them, so calling os.Process.Wait directly instead would leak both. + exited chan struct{} + // waitErr is the child's exit status. It is written before exited is + // closed and must only be read after that channel is observed closed. + waitErr error + + // proc is the child. It is an interface rather than *os.Process so that + // Stop's contract, in particular that the child is asked to stop exactly + // once however often Stop is called, can be pinned without a live process + // to signal. Nil means nothing was ever started. + proc processControl + + stopOnce sync.Once +} + +// processControl is the part of *os.Process that Stop needs. +// +// One interface rather than a pair of independent function fields: two fields +// can be wired to each other's operation, or one left nil, and no test can tell, +// because a fake satisfies any combination. There is nothing to swap or forget +// here, since the sole implementation is the real process and the method names +// carry the meaning. +type processControl interface { + Signal(os.Signal) error + Kill() error +} + +// *os.Process satisfies processControl unmodified, so production needs no +// adapter and no nil branch: the wiring is a single assignment. +var _ processControl = (*os.Process)(nil) + +// newServerCommand builds the child process. Split out from OfferToStart so the +// process' configuration can be asserted on without spawning anything. +func newServerCommand(bin string, stderr io.Writer) *exec.Cmd { + cmd := exec.Command(bin, "run") + // Stdin is left nil, so the child gets /dev/null: it is a background + // server, and sharing the terminal would have it stealing keystrokes from + // the agent. + cmd.Stdout = stderr // the child's logs are diagnostics, not chat output + cmd.Stderr = stderr + // Bound the wait for the child's output pipes; see childOutputDrainDelay. + cmd.WaitDelay = childOutputDrainDelay + return cmd +} + +// OfferToStart asks whether to start a LocalAI server and, if allowed, spawns +// one and waits for it to report ready. +// +// A child process rather than an in-process boot: RunCMD.Run installs its own +// signal handling and blocks until shutdown, so re-entering it from a chat +// session would entangle two lifecycles in one process. +func OfferToStart(ctx context.Context, opts StartOptions) (*StartedServer, error) { + if opts.Confirm == nil { + // Not interactive. Spawning a server nobody asked for is the one thing + // this function must never do: in CI, in a pipeline, or under a + // supervisor there is no one to see it or shut it down. + return nil, ErrDeclined + } + ok, err := opts.Confirm(fmt.Sprintf("No LocalAI server at %s. Start one now?", opts.Endpoint)) + if err != nil { + return nil, fmt.Errorf("asking whether to start a server: %w", err) + } + if !ok { + return nil, ErrDeclined + } + + bin := opts.Executable + if bin == "" { + if bin, err = os.Executable(); err != nil { + return nil, fmt.Errorf("locating the local-ai binary: %w", err) + } + } + + cmd := newServerCommand(bin, opts.Stderr) + if err := cmd.Start(); err != nil { + return nil, fmt.Errorf("starting a LocalAI server with %s: %w", bin, err) + } + + s := &StartedServer{exited: make(chan struct{}), proc: cmd.Process} + go func() { + s.waitErr = cmd.Wait() + close(s.exited) + }() + + timeout := opts.ReadyTimeout + if timeout <= 0 { + timeout = defaultReadyTimeout + } + if err := waitReady(ctx, opts.Endpoint, timeout, s.exited); err != nil { + if errors.Is(err, errServerExited) { + // Safe to read: errServerExited is only returned once exited has + // been observed closed, which happens after waitErr is written. + err = describeExit(err, s.waitErr) + } + s.Stop() + return nil, fmt.Errorf("%w. Run 'local-ai run' in another terminal to see why it did not come up", err) + } + return s, nil +} + +// describeExit adds what is known about how the child died to exitErr, without +// putting os/exec's plumbing in front of the user. +// +// waitErr is exec.ErrWaitDelay when the child exited cleanly but something it +// spawned still held its output pipe open past childOutputDrainDelay. The +// sentinel's own text names the WaitDelay field, which is meaningless to a +// user, so it is translated. Nothing is swallowed: os/exec only substitutes +// ErrWaitDelay when the process itself exited without an error of its own (see +// Cmd.Wait, "Report an error from the copying goroutines only if the program +// otherwise exited normally"), so it can never stand in for an *ExitError. +func describeExit(exitErr, waitErr error) error { + switch { + case waitErr == nil: + return exitErr + case errors.Is(waitErr, exec.ErrWaitDelay): + return fmt.Errorf("%w, and left a subprocess of its own still running", exitErr) + default: + return fmt.Errorf("%w: %w", exitErr, waitErr) + } +} + +// Stop terminates the server this process started, giving it a chance to shut +// down cleanly first. It is safe to call on a nil or never-started server, and +// safe to call more than once. +func (s *StartedServer) Stop() { + if s == nil || s.proc == nil { + return + } + s.stopOnce.Do(func() { + // SIGINT rather than SIGKILL: local-ai run installs its own handler and + // needs it to unload models and stop backend subprocesses. Killing it + // outright would strand those children. + _ = s.proc.Signal(os.Interrupt) + + select { + case <-s.exited: + case <-time.After(shutdownGrace): + // It ignored the interrupt or wedged on the way down. The user is + // waiting on their shell prompt, so stop being polite. + _ = s.proc.Kill() + } + }) +} + +// waitReady polls the endpoint's /readyz until the server reports ready, the +// budget expires, the caller gives up, or exited signals that the process we +// are waiting on is gone. A nil exited channel means there is no process to +// watch. +// +// Readiness lives on the endpoint ROOT, not under the /v1 API base URL, and it +// answers 503 for as long as startup is still in progress. +func waitReady(ctx context.Context, endpoint string, timeout time.Duration, exited <-chan struct{}) error { + url := strings.TrimSuffix(endpoint, "/") + "/readyz" + + // A real deadline rather than context.WithCancel plus a timer: the latter + // expires as context.Canceled, which every classifier here reads as "the + // caller gave up" rather than "the endpoint never answered". + waitCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + client := httpclient.NewWithTimeout(readyProbeTimeout) + ticker := time.NewTicker(readyPollInterval) + defer ticker.Stop() + + for { + select { + case <-exited: + return errServerExited + case <-waitCtx.Done(): + // Distinguish our budget from the caller's: only ours is advice + // about the server. + if err := ctx.Err(); err != nil { + return err + } + return fmt.Errorf("the LocalAI server did not become ready within %s", timeout) + case <-ticker.C: + } + + req, err := http.NewRequestWithContext(waitCtx, http.MethodGet, url, nil) + if err != nil { + return fmt.Errorf("building the readiness request for %s: %w", url, err) + } + resp, err := client.Do(req) + if err != nil { + continue // nothing listening yet + } + // Drain before closing so the next poll can reuse the connection + // instead of opening a socket every 500ms for two minutes. + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + if resp.StatusCode == http.StatusOK { + return nil + } + // Anything else means startup is still in progress; keep polling. + } +} diff --git a/core/cli/chat/server_test.go b/core/cli/chat/server_test.go new file mode 100644 index 000000000..652c1bdc1 --- /dev/null +++ b/core/cli/chat/server_test.go @@ -0,0 +1,375 @@ +package chat + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// unusedPort is a loopback address nothing listens on, used wherever a spec +// needs a readiness poll to keep failing. Port 1 is privileged, so no test +// process could have bound it. +const unusedPort = "http://127.0.0.1:1" + +var _ = Describe("OfferToStart", func() { + It("never spawns anything when there is no confirmer", func() { + started, err := OfferToStart(context.Background(), StartOptions{ + Endpoint: "http://127.0.0.1:59999", + Confirm: nil, + Stderr: io.Discard, + Executable: "/nonexistent/binary-that-must-not-run", + }) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrDeclined)).To(BeTrue(), "want ErrDeclined, got %v", err) + Expect(started).To(BeNil()) + }) + + It("does not spawn when the user declines", func() { + asked := false + started, err := OfferToStart(context.Background(), StartOptions{ + Endpoint: "http://127.0.0.1:59999", + Confirm: func(string) (bool, error) { + asked = true + return false, nil + }, + Stderr: io.Discard, + Executable: "/nonexistent/binary-that-must-not-run", + }) + Expect(asked).To(BeTrue(), "the user should have been asked") + Expect(errors.Is(err, ErrDeclined)).To(BeTrue()) + Expect(started).To(BeNil()) + }) + + It("names the endpoint in the question", func() { + var question string + _, _ = OfferToStart(context.Background(), StartOptions{ + Endpoint: "http://example.invalid:9090", + Confirm: func(q string) (bool, error) { + question = q + return false, nil + }, + Stderr: io.Discard, + Executable: "/nonexistent/binary-that-must-not-run", + }) + Expect(question).To(ContainSubstring("http://example.invalid:9090")) + }) + + It("propagates a confirmer error", func() { + boom := errors.New("boom") + _, err := OfferToStart(context.Background(), StartOptions{ + Endpoint: "http://127.0.0.1:59999", + Confirm: func(string) (bool, error) { return false, boom }, + Stderr: io.Discard, + Executable: "/nonexistent/binary-that-must-not-run", + }) + Expect(errors.Is(err, boom)).To(BeTrue()) + }) + + It("reports which binary it failed to launch", func() { + started, err := OfferToStart(context.Background(), StartOptions{ + Endpoint: "http://127.0.0.1:59999", + Confirm: func(string) (bool, error) { return true, nil }, + Stderr: io.Discard, + Executable: "/nonexistent/binary-that-must-not-run", + }) + Expect(started).To(BeNil()) + Expect(err).To(MatchError(ContainSubstring("starting a LocalAI server"))) + Expect(err).To(MatchError(ContainSubstring("/nonexistent/binary-that-must-not-run"))) + }) + + It("stops waiting as soon as the process it started exits", func() { + // A harmless no-op binary rather than a real server: this exercises the + // early-exit path without starting LocalAI, binding a port, or running + // 'local-ai run'. Without early-exit detection the call would sit here + // polling until ReadyTimeout. + bin, lookErr := exec.LookPath("true") + if lookErr != nil { + Skip("no 'true' binary on PATH to stand in for a server that dies at once") + } + + start := time.Now() + started, err := OfferToStart(context.Background(), StartOptions{ + Endpoint: unusedPort, + Confirm: func(string) (bool, error) { return true, nil }, + Stderr: io.Discard, + Executable: bin, + ReadyTimeout: 30 * time.Second, + }) + Expect(started).To(BeNil()) + Expect(err).To(MatchError(ContainSubstring("exited before it became ready"))) + Expect(time.Since(start)).To(BeNumerically("<", 10*time.Second), + "the wait should end with the process, not with the readiness budget") + }) + + It("gives up on a child whose grandchildren still hold its output pipe", func() { + // The real LocalAI shape: 'local-ai run' exits but a backend + // subprocess it spawned inherited the stderr pipe and keeps it open. + // Without cmd.WaitDelay, cmd.Wait blocks on the copy goroutine, exited + // never closes, and the readiness wait runs out the full budget instead + // of reporting that the server died. + sh, lookErr := exec.LookPath("sh") + if lookErr != nil { + Skip("no 'sh' binary on PATH to stand in for a server with a lingering child") + } + + dir := GinkgoT().TempDir() + pidFile := filepath.Join(dir, "grandchild.pid") + script := filepath.Join(dir, "server-with-lingering-child") + // #nosec G306 -- this has to be executable to stand in for a binary. + Expect(os.WriteFile(script, + []byte("#!"+sh+"\nsleep 30 &\necho $! > "+pidFile+"\nexit 0\n"), + 0o700)).To(Succeed()) + + // Reap the grandchild whatever happens: it outlives its own parent by + // design, so nothing else will clean it up. + DeferCleanup(func() { + raw, err := os.ReadFile(pidFile) + if err != nil { + return + } + pid, err := strconv.Atoi(strings.TrimSpace(string(raw))) + if err != nil { + return + } + proc, err := os.FindProcess(pid) + if err != nil { + return + } + _ = proc.Kill() + _, _ = proc.Wait() + }) + + start := time.Now() + started, err := OfferToStart(context.Background(), StartOptions{ + Endpoint: unusedPort, + Confirm: func(string) (bool, error) { return true, nil }, + Stderr: io.Discard, + Executable: script, + ReadyTimeout: 25 * time.Second, + }) + elapsed := time.Since(start) + + Expect(started).To(BeNil()) + Expect(err).To(MatchError(ContainSubstring("exited before it became ready")), + "an unbounded cmd.Wait would report a readiness timeout instead") + Expect(elapsed).To(BeNumerically("<", 20*time.Second), + "the wait must be bounded by the output drain, not by the readiness budget") + + // This is the case where cmd.Wait returns exec.ErrWaitDelay, whose own + // text names a struct field of os/exec. Users get told what happened + // instead. + Expect(err).NotTo(MatchError(ContainSubstring("WaitDelay")), + "os/exec plumbing must not reach the user") + Expect(err).NotTo(MatchError(ContainSubstring("exec:"))) + Expect(err).To(MatchError(ContainSubstring("left a subprocess of its own still running"))) + }) + + It("reports the exit status of a server that failed outright", func() { + // The counterpart to the case above: translating ErrWaitDelay must not + // cost a real exit status, which is the one diagnostic worth having. + bin, lookErr := exec.LookPath("false") + if lookErr != nil { + Skip("no 'false' binary on PATH to stand in for a server that fails") + } + + _, err := OfferToStart(context.Background(), StartOptions{ + Endpoint: unusedPort, + Confirm: func(string) (bool, error) { return true, nil }, + Stderr: io.Discard, + Executable: bin, + ReadyTimeout: 30 * time.Second, + }) + Expect(err).To(MatchError(ContainSubstring("exited before it became ready"))) + Expect(err).To(MatchError(ContainSubstring("exit status 1"))) + }) +}) + +var _ = Describe("StartedServer.Stop", func() { + It("is a no-op on a server that was never started", func() { + var nilServer *StartedServer + Expect(nilServer.Stop).NotTo(Panic()) + Expect((&StartedServer{}).Stop).NotTo(Panic()) + }) + + It("interrupts the child exactly once however often it is called", func() { + s, proc := stoppableServer() + + s.Stop() + s.Stop() + s.Stop() + + Expect(proc.interrupts.Load()).To(Equal(int32(1)), + "a second Stop must not signal the child again") + Expect(proc.kills.Load()).To(BeZero(), "a child that already exited must not be killed") + }) + + It("interrupts the child exactly once when called concurrently", func() { + // The realistic double-Stop: a deferred Stop on the way out racing the + // signal handler that also owns shutting the server down. + const callers = 8 + + s, proc := stoppableServer() + + var wg sync.WaitGroup + wg.Add(callers) + for range callers { + go func() { + defer GinkgoRecover() + defer wg.Done() + s.Stop() + }() + } + wg.Wait() + + Expect(proc.interrupts.Load()).To(Equal(int32(1))) + Expect(proc.kills.Load()).To(BeZero()) + }) + + It("asks the child to interrupt rather than killing it outright", func() { + // The escalation order is the whole point of the grace period: SIGKILL + // first would strand the backend subprocesses local-ai run owns. + s, proc := stoppableServer() + + s.Stop() + + Expect(proc.lastSignal.Load()).To(Equal(os.Interrupt)) + Expect(proc.kills.Load()).To(BeZero()) + }) +}) + +// countingProcess stands in for the *os.Process that Stop drives, recording +// what it was asked to do. +type countingProcess struct { + interrupts atomic.Int32 + kills atomic.Int32 + lastSignal atomic.Value +} + +func (p *countingProcess) Signal(sig os.Signal) error { + p.interrupts.Add(1) + p.lastSignal.Store(sig) + return nil +} + +func (p *countingProcess) Kill() error { + p.kills.Add(1) + return nil +} + +// stoppableServer builds a StartedServer whose child has already exited, driven +// by a countingProcess rather than a real one. Nothing is spawned. +func stoppableServer() (*StartedServer, *countingProcess) { + proc := &countingProcess{} + exited := make(chan struct{}) + close(exited) + return &StartedServer{exited: exited, proc: proc}, proc +} + +var _ = Describe("newServerCommand", func() { + It("bounds how long it will wait for the child's output pipes", func() { + cmd := newServerCommand("/nonexistent/binary-that-must-not-run", io.Discard) + + // An unbounded wait is the failure mode: backend subprocesses inherit + // the child's stderr pipe and can hold it open long after the server + // itself is gone. + Expect(cmd.WaitDelay).To(BeNumerically(">", 0), "cmd.Wait must not be unbounded") + Expect(cmd.WaitDelay).To(BeNumerically("<", shutdownGrace), + "a drain longer than the shutdown grace would kill a cleanly exited server") + }) + + It("runs the server subcommand without giving it the terminal", func() { + cmd := newServerCommand("/nonexistent/binary-that-must-not-run", io.Discard) + + Expect(cmd.Args).To(Equal([]string{"/nonexistent/binary-that-must-not-run", "run"})) + Expect(cmd.Stdin).To(BeNil(), "the child must not compete with the agent for stdin") + Expect(cmd.Stdout).NotTo(BeNil()) + Expect(cmd.Stderr).NotTo(BeNil()) + }) +}) + +var _ = Describe("waitReady", func() { + It("polls /readyz on the endpoint root and returns only once it answers 200", func() { + // readyOnPoll is deliberately above 1. A handler that answers 200 to the + // first poll cannot tell a correct implementation apart from one that + // treats 503 as ready, because both return after a single request; the + // poll count is what makes 503-as-ready observable. + const readyOnPoll = 3 + + var polls atomic.Int32 + var paths atomic.Value + paths.Store("") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + paths.Store(r.URL.Path) + if polls.Add(1) < readyOnPoll { + // What LocalAI answers while startup is still in progress. + w.WriteHeader(http.StatusServiceUnavailable) + return + } + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + Expect(waitReady(context.Background(), srv.URL, 20*time.Second, nil)).To(Succeed()) + Expect(paths.Load()).To(Equal("/readyz"), "readiness lives on the endpoint root, not under /v1") + Expect(polls.Load()).To(BeNumerically(">=", readyOnPoll), + "503 means startup is still in progress and must never be accepted as ready") + }) + + It("tolerates a trailing slash on the endpoint", func() { + var path atomic.Value + path.Store("") + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + path.Store(r.URL.Path) + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + Expect(waitReady(context.Background(), srv.URL+"/", 20*time.Second, nil)).To(Succeed()) + Expect(path.Load()).To(Equal("/readyz")) + }) + + It("reports a timeout, not a cancellation, when the budget runs out", func() { + err := waitReady(context.Background(), unusedPort, 1200*time.Millisecond, nil) + Expect(err).To(HaveOccurred()) + // A budget built from context.WithCancel plus a timer would surface as + // context.Canceled, which downstream code reads as "the caller gave up" + // and would stop classifying a hung server as unreachable. + Expect(errors.Is(err, context.Canceled)).To(BeFalse(), "got %v", err) + Expect(err).To(MatchError(ContainSubstring("did not become ready"))) + }) + + It("returns the caller's cancellation when the caller gives up", func() { + ctx, cancel := context.WithCancel(context.Background()) + go func() { + defer GinkgoRecover() + time.Sleep(200 * time.Millisecond) + cancel() + }() + defer cancel() + + err := waitReady(ctx, unusedPort, time.Minute, nil) + Expect(errors.Is(err, context.Canceled)).To(BeTrue(), "got %v", err) + }) + + It("gives up when the process it is waiting on has exited", func() { + exited := make(chan struct{}) + close(exited) + + err := waitReady(context.Background(), unusedPort, time.Minute, exited) + Expect(err).To(MatchError(ContainSubstring("exited before it became ready"))) + }) +}) diff --git a/core/cli/chat/session.go b/core/cli/chat/session.go deleted file mode 100644 index 651d05532..000000000 --- a/core/cli/chat/session.go +++ /dev/null @@ -1,112 +0,0 @@ -package chat - -import ( - "context" - "errors" - "fmt" - "io" - "slices" - "strings" -) - -const ( - chatRoleUser = "user" - chatRoleAssistant = "assistant" -) - -type chatMessage struct { - Role string - Content string -} - -type chatSession struct { - client chatClient - model string - models []string - messages []chatMessage -} - -func newChatSession(ctx context.Context, client chatClient, requestedModel string) (*chatSession, error) { - models, err := client.ListModels(ctx) - if err != nil { - return nil, fmt.Errorf("list models: %w", err) - } - - model, err := resolveChatModel(requestedModel, models) - if err != nil { - return nil, err - } - - return &chatSession{ - client: client, - model: model, - models: models, - }, nil -} - -func (s *chatSession) CurrentModel() string { - return s.model -} - -func (s *chatSession) Models() []string { - models := make([]string, len(s.models)) - copy(models, s.models) - return models -} - -func (s *chatSession) Clear() { - s.messages = nil -} - -func (s *chatSession) SwitchModel(model string) error { - if !slices.Contains(s.models, model) { - return fmt.Errorf("model %q is not available. Use /models to see installed models", model) - } - s.model = model - s.Clear() - return nil -} - -func (s *chatSession) Send(ctx context.Context, prompt string, out io.Writer) error { - s.messages = append(s.messages, chatMessage{ - Role: chatRoleUser, - Content: prompt, - }) - - answer, err := s.client.StreamChat(ctx, s.model, s.messages, out) - if err != nil { - return err - } - - s.messages = append(s.messages, chatMessage{ - Role: chatRoleAssistant, - Content: answer, - }) - return nil -} - -func resolveChatModel(requested string, models []string) (string, error) { - switch { - case requested == "" && len(models) == 0: - return "", errors.New(`no chat models are installed. - -Install a model first, for example: - local-ai models list - local-ai models install - local-ai run - -Then start a chat session: - local-ai chat --model `) - case requested == "" && len(models) == 1: - return models[0], nil - case requested == "" && len(models) > 1: - var b strings.Builder - b.WriteString("multiple models are available; choose one with --model:\n") - b.WriteString(formatChatModelList(models, "")) - return "", errors.New(b.String()) - case !slices.Contains(models, requested): - return "", fmt.Errorf("model %q is not available. Use `local-ai models list` and `local-ai models install `, or pass an installed model with --model", requested) - default: - return requested, nil - } -} diff --git a/core/cli/chat/session_test.go b/core/cli/chat/session_test.go deleted file mode 100644 index dcf274805..000000000 --- a/core/cli/chat/session_test.go +++ /dev/null @@ -1,56 +0,0 @@ -package chat - -import ( - "context" - "io" - - . "github.com/onsi/ginkgo/v2" - . "github.com/onsi/gomega" -) - -var _ = Describe("Chat session", func() { - It("keeps model switching and message history out of the terminal adapter", func() { - client := &fakeChatClient{ - models: []string{"alpha", "beta"}, - answer: "pong", - } - - session, err := newChatSession(context.Background(), client, "alpha") - Expect(err).ToNot(HaveOccurred()) - Expect(session.CurrentModel()).To(Equal("alpha")) - - Expect(session.SwitchModel("beta")).To(Succeed()) - Expect(session.CurrentModel()).To(Equal("beta")) - Expect(session.Send(context.Background(), "ping", io.Discard)).To(Succeed()) - - Expect(client.requests).To(HaveLen(1)) - Expect(client.requests[0].model).To(Equal("beta")) - Expect(client.requests[0].messages).To(HaveLen(1)) - Expect(client.requests[0].messages[0].Content).To(Equal("ping")) - }) -}) - -type fakeChatClient struct { - models []string - answer string - requests []fakeChatRequest -} - -type fakeChatRequest struct { - model string - messages []chatMessage -} - -func (c *fakeChatClient) ListModels(context.Context) ([]string, error) { - return c.models, nil -} - -func (c *fakeChatClient) StreamChat(_ context.Context, model string, messages []chatMessage, out io.Writer) (string, error) { - copied := make([]chatMessage, len(messages)) - copy(copied, messages) - c.requests = append(c.requests, fakeChatRequest{model: model, messages: copied}) - if _, err := io.WriteString(out, c.answer); err != nil { - return "", err - } - return c.answer, nil -} diff --git a/core/cli/chat/terminal.go b/core/cli/chat/terminal.go deleted file mode 100644 index 8d76e1e6f..000000000 --- a/core/cli/chat/terminal.go +++ /dev/null @@ -1,93 +0,0 @@ -package chat - -import ( - "bufio" - "context" - "fmt" - "io" - "strings" -) - -func runTerminalChat(ctx context.Context, session *chatSession, in io.Reader, out io.Writer) error { - scanner := bufio.NewScanner(in) - scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) - - if err := writeChat(out, "LocalAI chat (%s)\n", session.CurrentModel()); err != nil { - return err - } - if err := writeChat(out, "Type /exit to quit, /clear to reset the conversation, /models to list models.\n"); err != nil { - return err - } - - for { - if err := writeChat(out, "\n> "); err != nil { - return err - } - if !scanner.Scan() { - break - } - - prompt := strings.TrimSpace(scanner.Text()) - switch prompt { - case "": - continue - case "/bye", "/exit", "/quit": - return writeChat(out, "bye\n") - case "/clear": - session.Clear() - if err := writeChat(out, "conversation cleared\n"); err != nil { - return err - } - continue - case "/models": - if err := printChatModels(out, session.Models(), session.CurrentModel()); err != nil { - return err - } - continue - } - - if nextModel, ok := strings.CutPrefix(prompt, "/model "); ok { - nextModel = strings.TrimSpace(nextModel) - if nextModel == "" { - if err := writeChat(out, "usage: /model \n"); err != nil { - return err - } - continue - } - if err := session.SwitchModel(nextModel); err != nil { - if writeErr := writeChat(out, "%s\n", err); writeErr != nil { - return writeErr - } - continue - } - if err := writeChat(out, "switched to %s; conversation cleared\n", session.CurrentModel()); err != nil { - return err - } - continue - } - - if err := writeChat(out, "assistant: "); err != nil { - return err - } - if err := session.Send(ctx, prompt, out); err != nil { - return err - } - if err := writeChat(out, "\n"); err != nil { - return err - } - } - - return scanner.Err() -} - -func printChatModels(out io.Writer, models []string, current string) error { - if len(models) == 0 { - return writeChat(out, "no models installed\n") - } - return writeChat(out, "%s", formatChatModelList(models, current)) -} - -func writeChat(out io.Writer, format string, args ...any) error { - _, err := fmt.Fprintf(out, format, args...) - return err -} diff --git a/core/cli/chat_cmd.go b/core/cli/chat_cmd.go index 65228ff1f..8c2864a95 100644 --- a/core/cli/chat_cmd.go +++ b/core/cli/chat_cmd.go @@ -8,18 +8,72 @@ import ( cliContext "github.com/mudler/LocalAI/core/cli/context" ) +// ChatCMD runs the built-in terminal agent. Everything after the first +// positional argument is forwarded to the agent verbatim, so its own +// subcommands (plugin, skill, mcp) and their flags work unchanged. LocalAI's +// own flags must therefore come first. type ChatCMD struct { - Model string `short:"m" help:"Model name to use. Defaults to the only model returned by the server when exactly one is available"` - Endpoint string `env:"LOCALAI_CHAT_ENDPOINT" default:"http://127.0.0.1:8080" help:"LocalAI server endpoint. The /v1 path is added automatically when omitted"` - APIKey string `env:"LOCALAI_API_KEY,API_KEY" help:"API key to use when the LocalAI server requires authentication"` + Model string `short:"m" help:"Model to use. Defaults to the only model the server offers, or asks when there are several"` + Endpoint string `env:"LOCALAI_CHAT_ENDPOINT" default:"http://127.0.0.1:8080" help:"LocalAI server endpoint. The /v1 path is added automatically when omitted"` + APIKey string `env:"LOCALAI_API_KEY,API_KEY" help:"API key to use when the LocalAI server requires authentication"` + ConfigDir string `env:"LOCALAI_CHAT_CONFIG_DIR" help:"Directory holding the agent's config, plugins, and skills. Defaults to ~/.config/localai/chat" type:"path"` + TraceDir string `env:"LOCALAI_CHAT_TRACE_DIR" help:"Write a session LLM trace (NDJSON) to this directory" type:"path"` + + CLI bool `help:"Run in plain CLI mode instead of the full-screen interface"` + TUI bool `help:"Force the full-screen interface"` + Height string `help:"Run as an inline drop-down of this height, e.g. '40%'"` + Tmux bool `help:"Run in a tmux split"` + NoTmux bool `name:"no-tmux" help:"Never use a tmux split, even inside tmux"` + Init string `help:"Print the shell integration script for Ctrl+Space (zsh, bash, or fish)"` + Yolo bool `env:"LOCALAI_CHAT_YOLO" help:"Auto-approve every tool call without prompting"` + + Args []string `arg:"" optional:"" passthrough:"" help:"Arguments forwarded to the agent, e.g. 'plugin install ', 'skill list', 'mcp add'"` } func (c *ChatCMD) Run(ctx *cliContext.Context) error { - return chatcli.Run(context.Background(), chatcli.Options{ - Model: c.Model, - BaseURL: chatAPIBaseURL(c.Endpoint), - APIKey: c.APIKey, - In: os.Stdin, - Out: os.Stdout, + err := chatcli.Run(context.Background(), chatcli.Options{ + Args: c.agentArgs(), + Endpoint: c.Endpoint, + BaseURL: chatAPIBaseURL(c.Endpoint), + APIKey: c.APIKey, + Model: c.Model, + StateDir: c.ConfigDir, + TraceDir: c.TraceDir, + Yolo: c.Yolo, + In: os.Stdin, + Out: os.Stdout, + ErrOut: os.Stderr, }) + // The agent explains its own failures on stderr and hands back a code, so + // carry the code out and leave the explanation to stand alone. + if code, reported := chatcli.ExitStatus(err); reported { + return ExitCodeError{Code: code} + } + return err +} + +// agentArgs rebuilds the argument vector the agent expects: LocalAI's mode +// flags are declared here for discoverability and shell completion, so they +// have to be translated back into the agent's own flag names. +func (c *ChatCMD) agentArgs() []string { + var args []string + if c.CLI { + args = append(args, "--cli") + } + if c.TUI { + args = append(args, "--tui") + } + if c.Height != "" { + args = append(args, "--height", c.Height) + } + if c.Tmux { + args = append(args, "--tmux") + } + if c.NoTmux { + args = append(args, "--no-tmux") + } + if c.Init != "" { + args = append(args, "--init", c.Init) + } + return append(args, c.Args...) } diff --git a/core/cli/chat_cmd_test.go b/core/cli/chat_cmd_test.go index 55ce5014b..823a5ee61 100644 --- a/core/cli/chat_cmd_test.go +++ b/core/cli/chat_cmd_test.go @@ -1,6 +1,10 @@ package cli import ( + "errors" + "fmt" + + "github.com/alecthomas/kong" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) @@ -24,4 +28,70 @@ var _ = Describe("Chat command wiring", func() { Expect(chatAPIBaseURL("http://127.0.0.1:8080/localai")).To(Equal("http://127.0.0.1:8080/localai/v1")) }) }) + + Describe("argument parsing", func() { + parse := func(args ...string) *ChatCMD { + var cli struct { + Chat ChatCMD `cmd:""` + } + parser, err := kong.New(&cli) + Expect(err).ToNot(HaveOccurred()) + _, err = parser.Parse(append([]string{"chat"}, args...)) + Expect(err).ToNot(HaveOccurred()) + return &cli.Chat + } + + It("leaves Args empty for a bare invocation", func() { + Expect(parse().Args).To(BeEmpty()) + }) + + It("binds flags that precede the forwarded arguments", func() { + c := parse("--endpoint", "http://host:9090", "--model", "m", "plugin", "list") + Expect(c.Endpoint).To(Equal("http://host:9090")) + Expect(c.Model).To(Equal("m")) + Expect(c.Args).To(Equal([]string{"plugin", "list"})) + }) + + It("forwards flags that follow the first positional to the agent", func() { + c := parse("plugin", "install", "https://example.invalid/p", "--yes") + Expect(c.Args).To(Equal([]string{"plugin", "install", "https://example.invalid/p", "--yes"})) + }) + + It("parses its own mode flags", func() { + c := parse("--cli") + Expect(c.CLI).To(BeTrue()) + Expect(c.Args).To(BeEmpty()) + }) + }) + + // The agent prints its own diagnosis and hands back a status. main exits + // with that status and prints nothing more, so the user reads one message + // rather than an "exit status 1" stacked under it. + Describe("ExitCodeError", func() { + It("carries the status out", func() { + Expect(ExitCodeError{Code: 2}.Code).To(Equal(2)) + }) + + It("is recognisable after wrapping", func() { + var got ExitCodeError + Expect(errors.As(fmt.Errorf("chat: %w", ExitCodeError{Code: 2}), &got)).To(BeTrue()) + Expect(got.Code).To(Equal(2)) + }) + }) + + Describe("agentArgs", func() { + It("translates mode flags into the agent's own flags", func() { + c := &ChatCMD{CLI: true} + Expect(c.agentArgs()).To(Equal([]string{"--cli"})) + }) + + It("puts forwarded arguments after the translated flags", func() { + c := &ChatCMD{Height: "40%", Args: []string{"plugin", "list"}} + Expect(c.agentArgs()).To(Equal([]string{"--height", "40%", "plugin", "list"})) + }) + + It("returns nothing for a bare invocation", func() { + Expect((&ChatCMD{}).agentArgs()).To(BeEmpty()) + }) + }) }) diff --git a/core/cli/cli.go b/core/cli/cli.go index 8bf4b207a..77bf128cc 100644 --- a/core/cli/cli.go +++ b/core/cli/cli.go @@ -9,7 +9,7 @@ var CLI struct { cliContext.Context `embed:""` Run RunCMD `cmd:"" help:"Run LocalAI, this the default command if no other command is specified. Run 'local-ai run --help' for more information" default:"withargs"` - Chat ChatCMD `cmd:"" help:"Open an interactive chat session against a running LocalAI server"` + Chat ChatCMD `cmd:"" help:"Run the built-in terminal agent against a LocalAI server"` Federated FederatedCLI `cmd:"" help:"Run LocalAI in federated mode"` Models ModelsCMD `cmd:"" help:"Manage LocalAI models and definitions"` Backends BackendsCMD `cmd:"" help:"Manage LocalAI backends and definitions"` diff --git a/core/cli/exit.go b/core/cli/exit.go new file mode 100644 index 000000000..453807d12 --- /dev/null +++ b/core/cli/exit.go @@ -0,0 +1,15 @@ +package cli + +import "fmt" + +// ExitCodeError is a failure a command has already reported to the user. It +// carries nothing but the status the process should exit with, and main prints +// nothing more for it. +// +// It exists for commands that hand their terminal to something that does its +// own error reporting. Returning that subordinate's error instead would put a +// bare "exit status 1" underneath the explanation the user has just read, and +// returning nil would tell a script the run succeeded. +type ExitCodeError struct{ Code int } + +func (e ExitCodeError) Error() string { return fmt.Sprintf("exit status %d", e.Code) } diff --git a/core/cli/run.go b/core/cli/run.go index 600b2a997..749dd16f0 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -243,6 +243,23 @@ 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) @@ -732,8 +749,13 @@ 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(r.Address, authConfigured, r.AllowInsecurePublicBind); err != nil { + if err := requireAuthOrTrustedBind(listenAddress, authConfigured, r.AllowInsecurePublicBind); err != nil { return err } @@ -743,7 +765,11 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { return err } - xlog.Info("LocalAI is started and running", "address", r.Address) + if activatedListener != nil { + appHTTP.Listener = activatedListener + xlog.Info("Using systemd socket activation listener", "address", listenAddress) + } + xlog.Info("LocalAI is started and running", "address", listenAddress) // Start P2P if token was provided via CLI/env or loaded from runtime_settings.json if token != "" || app.ApplicationConfig().P2PToken != "" { @@ -762,11 +788,11 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { // backends like PostgreSQL need to call the embeddings API during // collection initialization. go func() { - waitForServerReady(r.Address, app.ApplicationConfig().Context) + waitForServerReady(listenAddress, app.ApplicationConfig().Context) app.StartAgentPool() }() - return appHTTP.Start(r.Address) + return appHTTP.Start(listenAddress) } // waitForServerReady polls the given address until the HTTP server is diff --git a/core/cli/run_socket_activation.go b/core/cli/run_socket_activation.go new file mode 100644 index 000000000..b809e3a48 --- /dev/null +++ b/core/cli/run_socket_activation.go @@ -0,0 +1,17 @@ +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)) + } +} diff --git a/core/cli/run_socket_activation_linux.go b/core/cli/run_socket_activation_linux.go new file mode 100644 index 000000000..35aade6cb --- /dev/null +++ b/core/cli/run_socket_activation_linux.go @@ -0,0 +1,71 @@ +//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 +} diff --git a/core/cli/run_socket_activation_other.go b/core/cli/run_socket_activation_other.go new file mode 100644 index 000000000..1bbbf1469 --- /dev/null +++ b/core/cli/run_socket_activation_other.go @@ -0,0 +1,9 @@ +//go:build !linux + +package cli + +import "net" + +func systemdActivatedListeners() ([]net.Listener, error) { + return nil, nil +} diff --git a/core/cli/run_socket_activation_test.go b/core/cli/run_socket_activation_test.go new file mode 100644 index 000000000..09f973596 --- /dev/null +++ b/core/cli/run_socket_activation_test.go @@ -0,0 +1,101 @@ +//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()) + }) +}) diff --git a/core/config/backend_capabilities.go b/core/config/backend_capabilities.go index 79687c302..13901d35b 100644 --- a/core/config/backend_capabilities.go +++ b/core/config/backend_capabilities.go @@ -3,6 +3,8 @@ package config import ( "slices" "strings" + + "github.com/mudler/LocalAI/pkg/model" ) // Usecase name constants — the canonical string values used in gallery entries, @@ -16,6 +18,7 @@ const ( UsecaseTokenize = "tokenize" UsecaseImage = "image" UsecaseVideo = "video" + Usecase3D = "3d" UsecaseTranscript = "transcript" UsecaseTTS = "tts" UsecaseSoundGeneration = "sound_generation" @@ -30,6 +33,7 @@ const ( UsecaseFaceRecognition = "face_recognition" UsecaseSpeakerRecognition = "speaker_recognition" UsecaseTokenClassify = "token_classify" + UsecaseScore = "score" ) // GRPCMethod identifies a Backend service RPC from backend.proto. @@ -40,7 +44,9 @@ const ( MethodPredictStream GRPCMethod = "PredictStream" MethodEmbedding GRPCMethod = "Embedding" MethodGenerateImage GRPCMethod = "GenerateImage" + MethodUpscaleImage GRPCMethod = "UpscaleImage" MethodGenerateVideo GRPCMethod = "GenerateVideo" + MethodGenerate3D GRPCMethod = "Generate3D" MethodAudioTranscription GRPCMethod = "AudioTranscription" MethodTTS GRPCMethod = "TTS" MethodTTSStream GRPCMethod = "TTSStream" @@ -60,6 +66,7 @@ const ( MethodVoiceEmbed GRPCMethod = "VoiceEmbed" MethodVoiceAnalyze GRPCMethod = "VoiceAnalyze" MethodTokenClassify GRPCMethod = "TokenClassify" + MethodScore GRPCMethod = "Score" ) // UsecaseInfo describes a single known_usecase value and how it maps @@ -122,6 +129,11 @@ 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, @@ -192,6 +204,11 @@ 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. @@ -209,6 +226,20 @@ type BackendCapability struct { AcceptsVideos bool // AcceptsAudios indicates multimodal audio input in Predict. AcceptsAudios bool + // AudioTransformInputMono16k declares that this backend's AudioTransform + // input must be folded to 16 kHz mono 16-bit WAV before it is handed over. + // + // Opt-IN, and the default of false means "hand the backend the upload as + // it is". The /audio/transform endpoint used to fold EVERY upload to + // 16 kHz mono, which is what LocalVQE wants for acoustic echo cancellation + // and what no source separation model can survive: htdemucs and + // mel_band_roformer refuse any rate but their checkpoint's own (44.1 kHz + // for every published one) and work in stereo, so every separation request + // made through the HTTP API failed with an INTERNAL raised inside the + // engine, while a direct gRPC call worked. Declaring the need rather than + // defaulting to it means a backend that wants the fold says so and a + // backend that does not needs no entry here at all. + AudioTransformInputMono16k bool // VoiceCloning describes the backend's per-request reference-audio // contract. Model variants that share a backend may narrow this further; // use VoiceCloningForModel for UI/API decisions. @@ -241,8 +272,8 @@ func referenceVoiceCloning() *VoiceCloningCapability { var BackendCapabilities = map[string]BackendCapability{ // --- LLM / text generation backends --- "llama-cpp": { - GRPCMethods: []GRPCMethod{MethodPredict, MethodPredictStream, MethodEmbedding, MethodTokenizeString}, - PossibleUsecases: []string{UsecaseChat, UsecaseCompletion, UsecaseEdit, UsecaseEmbeddings, UsecaseTokenize, UsecaseVision}, + GRPCMethods: []GRPCMethod{MethodPredict, MethodPredictStream, MethodEmbedding, MethodTokenizeString, MethodScore}, + PossibleUsecases: []string{UsecaseChat, UsecaseCompletion, UsecaseEdit, UsecaseEmbeddings, UsecaseTokenize, UsecaseVision, UsecaseScore}, DefaultUsecases: []string{UsecaseChat}, AcceptsImages: true, // requires mmproj Description: "llama.cpp GGUF models — LLM inference with optional vision via mmproj", @@ -318,7 +349,7 @@ var BackendCapabilities = map[string]BackendCapability{ // --- Image/video generation backends --- "diffusers": { - GRPCMethods: []GRPCMethod{MethodGenerateImage, MethodGenerateVideo}, + GRPCMethods: []GRPCMethod{MethodGenerateImage, MethodUpscaleImage, MethodGenerateVideo}, PossibleUsecases: []string{UsecaseImage, UsecaseVideo}, DefaultUsecases: []string{UsecaseImage}, Description: "HuggingFace diffusers — Stable Diffusion, Flux, video generation", @@ -344,6 +375,14 @@ 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}, @@ -412,6 +451,36 @@ var BackendCapabilities = map[string]BackendCapability{ DefaultUsecases: []string{UsecaseTranscript}, Description: "Sherpa-ONNX — multi-model speech toolkit (ASR, TTS, VAD)", }, + // audio-cpp is one gRPC server in front of ~30 audio.cpp model families, so + // PossibleUsecases is their UNION and no single model serves all of it: + // which RPCs a given model answers is decided by the family baked into its + // GGUF, and every audio-cpp gallery entry pins its own known_usecases. + // + // VoiceCloning is not decoration here. VoiceCloningForModel returns nil as + // soon as the backend has no capability entry, BEFORE it consults the + // model's own tts.voice_cloning override, so without this key a + // `voice: "profile:"` request is refused with 400 for every audio-cpp + // model and no model YAML can rescue it, on a backend that ships + // audio-cpp-chatterbox, whose family advertises cloning and not plain TTS, + // so a reference clip is the only way to use it at all. + // + // Deliberately NOT AudioTransformInputMono16k: the families this backend + // reaches through AudioTransform are separation and conversion (htdemucs, + // mel_band_roformer, seed_vc), which refuse any rate but their + // checkpoint's own and work from the stereo image. + "audio-cpp": { + GRPCMethods: []GRPCMethod{ + MethodTTS, MethodTTSStream, MethodAudioTranscription, + MethodVAD, MethodDiarize, MethodSoundGeneration, MethodAudioTransform, + }, + PossibleUsecases: []string{ + UsecaseTTS, UsecaseTranscript, UsecaseVAD, UsecaseDiarization, + UsecaseSoundGeneration, UsecaseAudioTransform, + }, + DefaultUsecases: []string{UsecaseTTS}, + VoiceCloning: referenceVoiceCloning(), + Description: "audio.cpp native engine: one server for TTS, voice cloning, ASR, forced alignment, VAD, diarization, source separation and music generation; the model's family decides which", + }, // --- TTS backends --- "piper": { @@ -556,7 +625,11 @@ var BackendCapabilities = map[string]BackendCapability{ GRPCMethods: []GRPCMethod{MethodAudioTransform}, PossibleUsecases: []string{UsecaseAudioTransform}, DefaultUsecases: []string{UsecaseAudioTransform}, - Description: "LocalVQE — joint AEC, noise suppression, and dereverberation for 16 kHz mono speech", + // The model is trained on 16 kHz mono speech and its AEC needs the + // input and the loopback reference in the same shape, so the endpoint + // keeps folding uploads for this backend. + AudioTransformInputMono16k: true, + Description: "LocalVQE — joint AEC, noise suppression, and dereverberation for 16 kHz mono speech", }, // --- Utility backends --- @@ -627,11 +700,59 @@ func NormalizeBackendName(backend string) string { return strings.ReplaceAll(backend, ".", "-") } -// llamaCppChannelSuffixes are the release-channel suffixes appended to a -// llama.cpp backend name in the gallery ("llama-cpp" vs -// "llama-cpp-development"). They carry no engine information, so they are -// stripped before the family check below. -var llamaCppChannelSuffixes = []string{"-development", "-quantization"} +// galleryChannelSuffixes are the release-channel suffixes appended to a backend +// name in the gallery ("llama-cpp" vs "llama-cpp-development" vs +// "llama-cpp-quantization"). They carry no engine information, so they are +// stripped before any family or capability lookup falls back. +var galleryChannelSuffixes = []string{"-development", "-quantization"} + +// galleryHardwarePrefixes are the acceleration prefixes the gallery prepends +// when it publishes one concrete backend image per hardware capability behind a +// meta name: "cpu-localvqe", "vulkan-localvqe", "cuda12-audio-cpp", +// "metal-darwin-arm64-llama-cpp". They carry no engine information either, and +// an operator may pin any of them in a model config's `backend:`. +// +// This is the exhaustive set present in backend/index.yaml, longest first so +// "cuda13-nvidia-l4t-arm64-" is tried before "cuda13-" and +// "intel-sycl-f16-" before "intel-". Stripping is a FALLBACK only (see +// GetBackendCapability), so a backend whose real name happened to start with +// one of these would still be found by its exact name first. +var galleryHardwarePrefixes = []string{ + "cuda13-nvidia-l4t-arm64-", + "metal-darwin-arm64-", + "nvidia-l4t-arm64-", + "intel-sycl-f16-", + "intel-sycl-f32-", + "nvidia-l4t-", + "vulkan-", + "cuda12-", + "cuda13-", + "metal-", + "intel-", + "rocm-", + "cpu-", +} + +// stripBackendVariant reduces a concrete gallery backend name to the meta name +// its capabilities are registered under: "vulkan-localvqe-development" becomes +// "localvqe". Returns the input unchanged when nothing matches. +// +// Both halves are needed. A pinned variant carries a hardware prefix, a release +// channel carries a suffix, and backend/index.yaml ships names with both. +func stripBackendVariant(name string) string { + for _, suffix := range galleryChannelSuffixes { + if strings.HasSuffix(name, suffix) { + name = strings.TrimSuffix(name, suffix) + break + } + } + for _, prefix := range galleryHardwarePrefixes { + if strings.HasPrefix(name, prefix) { + return strings.TrimPrefix(name, prefix) + } + } + return name +} // IsLlamaCppBackend reports whether a backend name refers to a build of the // llama.cpp gRPC server. The gallery ships one concrete backend per hardware @@ -652,7 +773,7 @@ func IsLlamaCppBackend(backend string) bool { if name == "" { return true } - for _, suffix := range llamaCppChannelSuffixes { + for _, suffix := range galleryChannelSuffixes { name = strings.TrimSuffix(name, suffix) } if strings.HasSuffix(name, "ik-llama-cpp") { @@ -711,11 +832,97 @@ func UsesLlamaCppServingOptions(backend string) bool { // GetBackendCapability returns the capability info for a backend, or nil if unknown. // Handles backend name normalization. +// +// A PINNED GALLERY VARIANT RESOLVES TO ITS META NAME. The gallery ships one +// image per hardware capability ("cpu-localvqe", "vulkan-localvqe", +// "metal-localvqe") and an operator may put any of them in a model's +// `backend:`. They are the same engine, so an exact-match-only lookup silently +// downgraded every pinned model to "unknown backend": vulkan-localvqe lost the +// 16 kHz mono fold that /audio/transform used to apply unconditionally and +// started failing inside LocalVQE, and a pinned audio-cpp variant would lose +// its voice-cloning contract the same way. Same class of bug as #10945, same +// answer as IsLlamaCppBackend. +// +// Exact match FIRST, so a backend genuinely registered under a variant-looking +// name keeps its own entry and stripping can never shadow it. func GetBackendCapability(backend string) *BackendCapability { - if cap, ok := BackendCapabilities[NormalizeBackendName(backend)]; ok { - return &cap + capability, _ := resolveBackendCapability(backend) + return capability +} + +// resolveBackendCapability is GetBackendCapability plus the key the entry was +// found under. Callers that then branch on backend identity MUST use that key, +// not the name they passed in. VoiceCloningForModel is the reason this exists: +// its per-backend switch encodes which model variants of a backend can clone, +// and keying it on the caller's spelling meant "cuda12-vibevoice-cpp" resolved +// the capability by stripping but missed the "vibevoice-cpp" case, falling +// through to the permissive default and advertising cloning for the 0.5B model +// that cannot do it. +func resolveBackendCapability(backend string) (*BackendCapability, string) { + name := NormalizeBackendName(backend) + if cap, ok := BackendCapabilities[name]; ok { + return &cap, name } - return nil + if base := stripBackendVariant(name); base != name { + if cap, ok := BackendCapabilities[base]; ok { + return &cap, base + } + } + return nil, name +} + +// AudioTransformRequiresMono16kInput reports whether /audio/transform must fold +// uploads to 16 kHz mono before handing them to this backend. +// +// False for an unknown backend, which is the safe answer: an unregistered +// backend gets its upload unchanged, so a model that needs the file intact +// (source separation, voice conversion at 44.1 kHz) works without an entry +// here, and one that needs the fold cannot get it by accident. +// +// Pinned gallery variants are covered: GetBackendCapability strips the hardware +// prefix and the release-channel suffix, so cpu-localvqe, vulkan-localvqe and +// metal-localvqe all fold exactly as "localvqe" does. They must, because the +// usecase gate does not stand in for this one: naming a model explicitly makes +// BuildFilteredFirstAvailableDefaultModel return before it filters. +func AudioTransformRequiresMono16kInput(backend string) bool { + capability := GetBackendCapability(backend) + return capability != nil && capability.AudioTransformInputMono16k +} + +// llmAutoLoadUsecases are the usecases that mark a backend able to serve a +// text/LLM GGUF model. A GGUF model that declares no explicit backend must only +// be auto-tried against backends carrying one of these usecases - never against +// audio/codec/image backends (e.g. opus) that happen to be installed alongside +// it (see issue #9287). +var llmAutoLoadUsecases = []string{ + UsecaseChat, + UsecaseCompletion, + UsecaseEdit, + UsecaseEmbeddings, +} + +// isLLMCapableForAutoLoad reports whether the named backend is known to serve +// text/LLM models, for pkg/model's GGUF backend auto-detection (#9287). Backends +// absent from the capability table are treated as not LLM-capable. +func isLLMCapableForAutoLoad(name string) bool { + capability := GetBackendCapability(name) + if capability == nil { + return false + } + for _, u := range capability.PossibleUsecases { + if slices.Contains(llmAutoLoadUsecases, u) { + return true + } + } + return false +} + +func init() { + // Wire the LLM-capability filter into pkg/model's GGUF backend + // auto-detection. pkg/model is a lower-level package and must not import + // core/config (that would form a core/config -> pkg/model -> core/config + // import cycle), so core/config registers the predicate here instead (#9287). + model.RegisterLLMCapableBackendFunc(isLLMCapableForAutoLoad) } // VoiceCloningForModel returns the reference-audio contract only when the @@ -728,8 +935,7 @@ func VoiceCloningForModel(cfg *ModelConfig) *VoiceCloningCapability { if cfg == nil { return nil } - backend := NormalizeBackendName(cfg.Backend) - capability := GetBackendCapability(backend) + capability, backend := resolveBackendCapability(cfg.Backend) if capability == nil || capability.VoiceCloning == nil { return nil } diff --git a/core/config/backend_capabilities_test.go b/core/config/backend_capabilities_test.go index 1a5e1d7db..5a45d23bb 100644 --- a/core/config/backend_capabilities_test.go +++ b/core/config/backend_capabilities_test.go @@ -70,6 +70,135 @@ var _ = Describe("GetBackendCapability", func() { It("returns nil for unknown backends", func() { Expect(GetBackendCapability("nonexistent")).To(BeNil()) }) + + // The gallery ships one concrete image per hardware capability behind a + // meta name, and an operator may pin any of them in a model's `backend:`. + // An exact-match-only lookup silently treated every one of them as an + // unknown backend, which cost vulkan-localvqe the 16 kHz mono fold its AEC + // needs and would cost a pinned audio-cpp variant its voice-cloning + // contract. Same class as #10945. + It("resolves a pinned hardware variant to its meta backend", func() { + for _, name := range []string{"cpu-localvqe", "vulkan-localvqe", "metal-localvqe"} { + capability := GetBackendCapability(name) + Expect(capability).NotTo(BeNil(), "pinned variant %q must resolve", name) + Expect(capability.PossibleUsecases).To(ContainElement(UsecaseAudioTransform), name) + } + }) + + It("resolves a pinned variant that also carries a release channel", func() { + for _, name := range []string{ + "cuda12-audio-cpp", "cuda13-audio-cpp-development", + "metal-audio-cpp", "cpu-audio-cpp-development", + "cuda13-nvidia-l4t-arm64-llama-cpp", "intel-sycl-f16-llama-cpp", + "metal-darwin-arm64-llama-cpp", "nvidia-l4t-arm64-llama-cpp", + "rocm-llama-cpp-development", "intel-llama-cpp", + } { + Expect(GetBackendCapability(name)).NotTo(BeNil(), "pinned variant %q must resolve", name) + } + }) + + It("does not invent a capability for a name that only looks like a variant", func() { + Expect(GetBackendCapability("cpu-nonexistent")).To(BeNil()) + Expect(GetBackendCapability("vulkan-")).To(BeNil()) + }) + + // Stripping is a fallback, never a rewrite: a backend registered under its + // own name keeps its own entry even if that name starts with a prefix. + It("prefers an exact match over the stripped one", func() { + BackendCapabilities["cpu-exact-match-probe"] = BackendCapability{ + PossibleUsecases: []string{UsecaseChat}, + Description: "test fixture", + } + BackendCapabilities["exact-match-probe"] = BackendCapability{ + PossibleUsecases: []string{UsecaseTTS}, + Description: "test fixture", + } + DeferCleanup(func() { + delete(BackendCapabilities, "cpu-exact-match-probe") + delete(BackendCapabilities, "exact-match-probe") + }) + + capability := GetBackendCapability("cpu-exact-match-probe") + Expect(capability).NotTo(BeNil()) + Expect(capability.PossibleUsecases).To(Equal([]string{UsecaseChat})) + }) +}) + +// audio-cpp advertises voice cloning from the backend itself and ships +// audio-cpp-chatterbox, whose family serves cloning and NOT plain TTS, so a +// reference clip is the only way to use it. Without a capability entry +// VoiceCloningForModel returns nil before it ever reads the model's own +// tts.voice_cloning override, so `voice: "profile:"` was refused with a 400 +// for every audio-cpp model and no model YAML could rescue it. +var _ = Describe("audio-cpp capabilities", func() { + It("is registered", func() { + Expect(GetBackendCapability("audio-cpp")).NotTo(BeNil()) + }) + + It("advertises the RPCs its families actually serve", func() { + capability := GetBackendCapability("audio-cpp") + Expect(capability.GRPCMethods).To(ContainElements( + MethodTTS, MethodTTSStream, MethodAudioTranscription, + MethodVAD, MethodDiarize, MethodSoundGeneration, MethodAudioTransform)) + Expect(capability.PossibleUsecases).To(ContainElements( + UsecaseTTS, UsecaseTranscript, UsecaseVAD, UsecaseDiarization, + UsecaseSoundGeneration, UsecaseAudioTransform)) + }) + + It("carries the reference-audio contract, for pinned variants too", func() { + for _, name := range []string{"audio-cpp", "cuda12-audio-cpp", "metal-audio-cpp"} { + cloning := VoiceCloningForModel(&ModelConfig{Backend: name}) + Expect(cloning).NotTo(BeNil(), "%q must reach the backend with a profile voice", name) + Expect(cloning.AcceptedAudioFormats).To(ContainElement("audio/wav")) + } + }) + + // The families audio-cpp reaches through AudioTransform are separation and + // conversion, which refuse any rate but their checkpoint's own. + It("does not ask for the 16 kHz mono fold", func() { + Expect(AudioTransformRequiresMono16kInput("audio-cpp")).To(BeFalse()) + }) +}) + +// The fold to 16 kHz mono in /audio/transform is opt-IN. It used to be +// unconditional, which made source separation unreachable through the HTTP +// API: htdemucs and mel_band_roformer refuse any rate but their checkpoint's +// own and separate using the stereo image, so every such request died with an +// INTERNAL raised inside the engine while the same call over gRPC worked. +var _ = Describe("AudioTransformRequiresMono16kInput", func() { + It("folds for localvqe, whose AEC is trained on 16 kHz mono", func() { + Expect(AudioTransformRequiresMono16kInput("localvqe")).To(BeTrue()) + }) + + // Pinned gallery variants are the same engine and must fold identically. + // They did not: the lookup was exact-match only, so vulkan-localvqe was an + // unknown backend, lost the fold that used to be unconditional, and started + // failing inside LocalVQE. The usecase gate is no substitute, because + // BuildFilteredFirstAvailableDefaultModel returns early once the client + // names a model explicitly. + It("folds for every pinned localvqe variant the gallery ships", func() { + Expect(AudioTransformRequiresMono16kInput("cpu-localvqe")).To(BeTrue()) + Expect(AudioTransformRequiresMono16kInput("vulkan-localvqe")).To(BeTrue()) + Expect(AudioTransformRequiresMono16kInput("metal-localvqe")).To(BeTrue()) + }) + + It("does not fold for a backend that has not asked for it", func() { + // Registered, and deliberately NOT folding: its separation and + // conversion families refuse any rate but their checkpoint's own. + Expect(AudioTransformRequiresMono16kInput("audio-cpp")).To(BeFalse()) + Expect(AudioTransformRequiresMono16kInput("nonexistent")).To(BeFalse()) + Expect(AudioTransformRequiresMono16kInput("")).To(BeFalse()) + }) + + It("is claimed by no other registered backend", func() { + for name, capability := range BackendCapabilities { + if name == "localvqe" { + continue + } + Expect(capability.AudioTransformInputMono16k).To(BeFalse(), + "backend %q asks for the 16 kHz mono fold; that has to be a deliberate, documented need", name) + } + }) }) var _ = Describe("VoiceCloningForModel", func() { @@ -93,6 +222,29 @@ var _ = Describe("VoiceCloningForModel", func() { Entry("legacy option custom opt-in", ModelConfig{Name: "private-build", Backend: "qwen3-tts-cpp", Options: []string{"voice_cloning:true"}}, true), Entry("legacy option opt-out", ModelConfig{Name: "voxcpm-1.5", Backend: "voxcpm", Options: []string{"voice_cloning=false"}}, false), ) + + // A pinned gallery variant must reach the SAME per-backend rule as the meta + // name, in both directions. Resolving the capability by stripping the prefix + // while still keying the model-variant switch on the pinned spelling made + // every variant fall through to the permissive default: cuda12-vibevoice-cpp + // advertised cloning for the realtime 0.5B model, which cannot do it, and + // /v1/audio/speech accepted a profile: voice it had to fail on in the backend + // instead of rejecting it with a 400. + DescribeTable("resolves the model-variant rule through pinned gallery variants", + func(cfg ModelConfig, expected bool) { + Expect(VoiceCloningForModel(&cfg) != nil).To(Equal(expected)) + }, + Entry("cuda12-vibevoice-cpp 0.5B stays unsupported", ModelConfig{Name: "vibevoice-cpp-0.5b", Backend: "cuda12-vibevoice-cpp"}, false), + Entry("cuda12-vibevoice-cpp 1.5B stays supported", ModelConfig{Name: "vibevoice-1.5b", Backend: "cuda12-vibevoice-cpp"}, true), + Entry("metal-coqui tacotron2 stays unsupported", ModelConfig{Name: "tacotron2-en", Backend: "metal-coqui"}, false), + Entry("metal-coqui xtts stays supported", ModelConfig{Name: "xtts-v2", Backend: "metal-coqui"}, true), + Entry("cuda12-crispasr ASR stays unsupported", ModelConfig{Name: "parakeet-asr", Backend: "cuda12-crispasr"}, false), + Entry("cuda12-crispasr F5 stays supported", ModelConfig{Name: "f5-tts-crispasr", Backend: "cuda12-crispasr"}, true), + Entry("cpu-qwen3-tts-cpp CustomVoice stays unsupported", ModelConfig{Name: "qwen3-tts-flash", Backend: "cpu-qwen3-tts-cpp"}, false), + Entry("cpu-qwen3-tts-cpp Base stays supported", ModelConfig{Name: "qwen3-tts-cpp-0.6b-base", Backend: "cpu-qwen3-tts-cpp"}, true), + Entry("release channel suffix too", ModelConfig{Name: "vibevoice-cpp-0.5b", Backend: "vibevoice-cpp-development"}, false), + Entry("pinned audio-cpp keeps its unconditional cloning", ModelConfig{Name: "audio-cpp-chatterbox", Backend: "cuda12-audio-cpp"}, true), + ) }) var _ = Describe("IsValidUsecaseForBackend", func() { diff --git a/core/config/hooks_test.go b/core/config/hooks_test.go index 4a6e31d6c..b69bc6989 100644 --- a/core/config/hooks_test.go +++ b/core/config/hooks_test.go @@ -198,6 +198,31 @@ var _ = Describe("Backend hooks and parser defaults", func() { // chunked_prefill is still seeded since user didn't set it Expect(cfg.EngineArgs["enable_chunked_prefill"]).To(Equal(true)) }) + + // The backend applies options: before engine_args:, so seeding a default + // for a key the user already set through a CLI-style option would silently + // win over it. https://github.com/mudler/LocalAI/issues/11130 + It("does not seed a default the user set through options", func() { + cfg := &ModelConfig{ + Backend: "vllm", + Options: []string{"--enable-prefix-caching:false", "--enable-chunked-prefill"}, + } + cfg.SetDefaults() + + Expect(cfg.EngineArgs).NotTo(HaveKey("enable_prefix_caching")) + Expect(cfg.EngineArgs).NotTo(HaveKey("enable_chunked_prefill")) + }) + + It("still seeds defaults when options carry unrelated entries", func() { + cfg := &ModelConfig{ + Backend: "vllm", + Options: []string{"tool_parser:hermes", "--quantization:gptq_marlin"}, + } + cfg.SetDefaults() + + Expect(cfg.EngineArgs["enable_prefix_caching"]).To(Equal(true)) + Expect(cfg.EngineArgs["enable_chunked_prefill"]).To(Equal(true)) + }) }) Context("llamaCppDefaults GGUF guessing", func() { diff --git a/core/config/hooks_vllm.go b/core/config/hooks_vllm.go index ffdd1a52a..67408511b 100644 --- a/core/config/hooks_vllm.go +++ b/core/config/hooks_vllm.go @@ -59,19 +59,48 @@ func vllmDefaults(cfg *ModelConfig, modelPath string) { } // applyEngineArgDefaults seeds production-friendly engine_args without overwriting -// anything the user already set. +// anything the user already set, in engine_args or as a CLI-style option. func applyEngineArgDefaults(cfg *ModelConfig) { if cfg.EngineArgs == nil { cfg.EngineArgs = map[string]any{} } + fromOptions := engineOptionKeys(cfg.Options) for k, v := range productionEngineArgsDefaults { if _, set := cfg.EngineArgs[k]; set { continue } + // The backend applies options: before engine_args:, so seeding a key + // the user wrote as an option would silently override it. + if _, set := fromOptions[k]; set { + continue + } cfg.EngineArgs[k] = v } } +// engineOptionKeys returns the engine-arg field names carried by CLI-style +// options (`--enable-prefix-caching:false` -> `enable_prefix_caching`). Only +// `--` prefixed entries are engine flags; the rest of Options[] is +// backend-level (tool_parser:, reasoning_parser:, ...). +func engineOptionKeys(options []string) map[string]struct{} { + keys := map[string]struct{}{} + for _, opt := range options { + opt = strings.TrimSpace(opt) + if !strings.HasPrefix(opt, "--") { + continue + } + name := opt + if i := strings.IndexAny(opt, ":="); i != -1 { + name = opt[:i] + } + name = strings.ReplaceAll(strings.TrimLeft(name, "-"), "-", "_") + if name != "" { + keys[name] = struct{}{} + } + } + return keys +} + func applyParserDefaults(cfg *ModelConfig) { hasToolParser := false hasReasoningParser := false diff --git a/core/config/meta/registry.go b/core/config/meta/registry.go index 2ad64f4d3..f386bd1a5 100644 --- a/core/config/meta/registry.go +++ b/core/config/meta/registry.go @@ -623,6 +623,13 @@ 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", @@ -630,6 +637,99 @@ 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": { @@ -822,6 +922,13 @@ 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 diff --git a/core/config/model_capabilities.go b/core/config/model_capabilities.go index 79a117b4a..6b432bddc 100644 --- a/core/config/model_capabilities.go +++ b/core/config/model_capabilities.go @@ -15,9 +15,10 @@ const ( ModalityImage = "image" ModalityAudio = "audio" ModalityVideo = "video" + Modality3D = "3d" ) -var modalityOrder = []string{ModalityText, ModalityImage, ModalityAudio, ModalityVideo} +var modalityOrder = []string{ModalityText, ModalityImage, ModalityAudio, ModalityVideo, Modality3D} func declaredModalities(modalities []string) map[string]bool { declared := make(map[string]bool, len(modalities)) @@ -154,6 +155,7 @@ 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) @@ -181,9 +183,10 @@ 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 models consume images directly. + // Ollama surface); detection/depth/face/3D 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_DETECTION) || c.HasUsecases(FLAG_DEPTH) || c.HasUsecases(FLAG_FACE_RECOGNITION) || + c.HasUsecases(FLAG_3D) 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) || @@ -208,10 +211,12 @@ 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) } diff --git a/core/config/model_capabilities_test.go b/core/config/model_capabilities_test.go index 545bad50b..84f01af7b 100644 --- a/core/config/model_capabilities_test.go +++ b/core/config/model_capabilities_test.go @@ -109,6 +109,25 @@ 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), diff --git a/core/config/model_config.go b/core/config/model_config.go index 7ca24be64..21bc42fac 100644 --- a/core/config/model_config.go +++ b/core/config/model_config.go @@ -232,6 +232,15 @@ 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 @@ -669,6 +678,16 @@ 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 @@ -682,6 +701,65 @@ 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. @@ -982,6 +1060,12 @@ 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 @@ -1435,20 +1519,9 @@ func (c *ModelConfig) Validate() (bool, error) { ProxyProviderOpenAI, ProxyProviderAnthropic) } - // 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") - } + // 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. // Pattern detector: validate built-in names and that each operator-defined // pattern is a well-formed, anchored, bounded restricted-regex. Reject at @@ -1576,9 +1649,10 @@ 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 bypasses the slot loop - // (direct llama_decode), so combining score with - // chat/completion/embeddings in one config is rejected at validation. + // 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. FLAG_SCORE ModelConfigUsecase = 0b10000000000000000000 // Marks a model as wired for the Depth gRPC primitive (per-pixel @@ -1599,6 +1673,11 @@ 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 ) @@ -1612,7 +1691,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, // visual generation + FLAG_IMAGE | FLAG_VIDEO | FLAG_3D, // visual generation } // IsMultimodal returns true if the given usecases span two or more orthogonal @@ -1659,6 +1738,7 @@ func GetAllModelConfigUsecases() map[string]ModelConfigUsecase { "FLAG_SCORE": FLAG_SCORE, "FLAG_DEPTH": FLAG_DEPTH, "FLAG_TOKEN_CLASSIFY": FLAG_TOKEN_CLASSIFY, + "FLAG_3D": FLAG_3D, } } @@ -1691,9 +1771,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, and (on llama-cpp) reintroduce the slot -// contention the conflict check exists to prevent. So a declared score or -// token_classify list is authoritative. +// 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. func (c *ModelConfig) HasUsecases(u ModelConfigUsecase) bool { if c.KnownUsecases != nil { if (u & *c.KnownUsecases) == u { @@ -1810,6 +1890,13 @@ 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) { @@ -1883,8 +1970,8 @@ func (c *ModelConfig) GuessUsecases(u ModelConfigUsecase) bool { if (u & FLAG_SCORE) == FLAG_SCORE { // No heuristic: Score-intent is a deliberate operator choice - // (it reserves the model from generation traffic on llama-cpp), - // so HasUsecases(FLAG_SCORE) is true only when KnownUsecases + // (it keeps the model out of pickers it wasn't meant for), so + // HasUsecases(FLAG_SCORE) is true only when KnownUsecases // declares it explicitly. return false } diff --git a/core/config/model_config_loader.go b/core/config/model_config_loader.go index be1fc2695..788b4ce7f 100644 --- a/core/config/model_config_loader.go +++ b/core/config/model_config_loader.go @@ -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) (*ModelConfig, error) { - cfg, err := bcl.LoadModelConfigFileByName(modelName, modelPath) +func (bcl *ModelConfigLoader) LoadResolvedModelConfig(modelName, modelPath string, opts ...ConfigLoaderOption) (*ModelConfig, error) { + cfg, err := bcl.LoadModelConfigFileByName(modelName, modelPath, opts...) if err != nil { return nil, err } diff --git a/core/config/model_config_loader_resolve_test.go b/core/config/model_config_loader_resolve_test.go index 961693b01..55e431c6b 100644 --- a/core/config/model_config_loader_resolve_test.go +++ b/core/config/model_config_loader_resolve_test.go @@ -49,4 +49,21 @@ 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)) + }) }) diff --git a/core/config/model_config_test.go b/core/config/model_config_test.go index 1b7a10f45..21741a061 100644 --- a/core/config/model_config_test.go +++ b/core/config/model_config_test.go @@ -127,21 +127,19 @@ parameters: Expect(err).To(BeNil()) Expect(valid).To(BeTrue()) - // 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.) + // 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. scoreFlag := FLAG_SCORE | FLAG_CHAT - conflicting := ModelConfig{ - Name: "router-but-also-chat", + scoringChat := ModelConfig{ + Name: "router-and-chat", Backend: "llama-cpp", KnownUsecases: &scoreFlag, } - valid, err = conflicting.Validate() - Expect(valid).To(BeFalse()) - Expect(err).To(HaveOccurred()) - Expect(err.Error()).To(ContainSubstring("score is incompatible")) + valid, err = scoringChat.Validate() + Expect(valid).To(BeTrue()) + Expect(err).NotTo(HaveOccurred()) scoreOnly := FLAG_SCORE dedicated := ModelConfig{ diff --git a/core/gallery/backends_test.go b/core/gallery/backends_test.go index a3339861c..1b7e059be 100644 --- a/core/gallery/backends_test.go +++ b/core/gallery/backends_test.go @@ -32,6 +32,42 @@ 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()) } diff --git a/core/gallery/estimate_warm.go b/core/gallery/estimate_warm.go new file mode 100644 index 000000000..4b291cd58 --- /dev/null +++ b/core/gallery/estimate_warm.go @@ -0,0 +1,212 @@ +package gallery + +import ( + "context" + "os" + "strconv" + "strings" + "sync" + "time" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/concurrency" + "github.com/mudler/LocalAI/pkg/system" + "github.com/mudler/LocalAI/pkg/vram" + "github.com/mudler/xlog" +) + +// EstimateInput builds the VRAM estimator's input from a gallery entry. +// +// It lives here rather than beside the HTTP handler because two callers need +// it: the handler answering one model, and the warmer below answering all of +// them ahead of time. +func EstimateInput(m *GalleryModel) vram.ModelEstimateInput { + var input vram.ModelEstimateInput + input.Size = m.Size + if repoID := extractHFRepo(m.Overrides, m.URLs); repoID != "" { + input.HFRepo = repoID + } + for _, f := range m.AdditionalFiles { + if vram.IsWeightFile(f.URI) { + input.Files = append(input.Files, vram.FileInput{URI: f.URI, Size: 0}) + } + } + return input +} + +// extractHFRepo finds a HuggingFace repo ID in a model's overrides or URLs. +func extractHFRepo(overrides map[string]any, urls []string) string { + if overrides != nil { + if params, ok := overrides["parameters"].(map[string]any); ok { + if modelRef, ok := params["model"].(string); ok { + if repoID, ok := vram.ExtractHFRepoID(modelRef); ok { + return repoID + } + } + } + } + for _, u := range urls { + if repoID, ok := vram.ExtractHFRepoID(u); ok { + return repoID + } + } + return "" +} + +// EstimateWarmConfig bounds the background warm-up. +type EstimateWarmConfig struct { + // Limit is how many gallery entries to warm, in gallery order. Zero + // disables warming entirely. The order matters: it is the order the UI + // lists them in, so the entries a user sees first are warmed first. + Limit int + // Concurrency is how many estimates run at once. Each one can be a remote + // probe, so this is deliberately small: the point is to be finished before + // anybody looks, not to saturate the link or the upstream. + Concurrency int + // Contexts are the context lengths to estimate at. These want to match what + // the UI asks for, or the warmed entry is not the one it reads. + Contexts []uint32 +} + +// DefaultEstimateWarmConfig is what the server uses unless told otherwise. +// +// The limit is a deliberate compromise. Warming the whole gallery would be +// thousands of remote probes on every boot, which is rude to the upstream and +// slow to finish; warming nothing leaves the first page of the model gallery +// paying two seconds per row. A few hundred covers what anyone browses in a +// sitting, and everything past it still warms itself on first view. +var DefaultEstimateWarmConfig = EstimateWarmConfig{ + Limit: 300, + Concurrency: 4, + Contexts: []uint32{8192, 16384, 32768, 65536, 131072, 262144}, +} + +// WarmEstimateCache fills the gallery's derived caches in the background. +// +// Two things are warmed, and they are the same cost wearing different hats. +// An estimate for an entry the server has never seen costs a network probe of +// its weight files, and describing an entry's variants costs one probe per +// build it offers. The UI asks for an estimate per row and a variant +// description per model opened, so without this the first visitor pays for +// both: ten seconds of a page filling in its own sizes, then another second +// and a half the first time they click anything. +// +// Both land in the same caches underneath, which is why one pass covers them. +// +// It returns immediately; the work happens on its own goroutine and stops when +// ctx is done. Failures are logged at debug and otherwise ignored: a warm-up +// that cannot reach an upstream must never stop the server from starting, and +// the entry it failed on simply stays cold. +func WarmEstimateCache(ctx context.Context, galleries []config.Gallery, systemState *system.SystemState, cfg EstimateWarmConfig) { + if cfg.Limit <= 0 || cfg.Concurrency <= 0 { + return + } + + concurrency.SafeGo(func() { + started := time.Now() + + models, err := AvailableGalleryModelsCached(galleries, systemState) + if err != nil { + xlog.Debug("VRAM estimate warm-up skipped, gallery unavailable", "error", err) + return + } + if len(models) > cfg.Limit { + models = models[:cfg.Limit] + } + if len(models) == 0 { + return + } + + // The host gate the variant picker resolves against. Derived once: it + // describes this machine, not this entry, and HostResolveEnv reads the + // system state to build it. + env := HostResolveEnv(ctx, systemState) + + var ( + wg sync.WaitGroup + cursor = make(chan *GalleryModel) + warmed int + warmedVariants int + mu sync.Mutex + ) + + for i := 0; i < cfg.Concurrency; i++ { + wg.Add(1) + concurrency.SafeGo(func() { + defer wg.Done() + for m := range cursor { + // Per entry, not for the run: one unreachable weight file + // must not hold a worker for the whole warm-up. + entryCtx, cancel := context.WithTimeout(ctx, 30*time.Second) + + input := EstimateInput(m) + if len(input.Files) > 0 || input.HFRepo != "" || input.Size != "" { + if _, err := vram.EstimateModelMultiContext(entryCtx, input, cfg.Contexts); err != nil { + xlog.Debug("VRAM estimate warm-up failed for entry", "model", m.GetName(), "error", err) + } else { + mu.Lock() + warmed++ + mu.Unlock() + } + } + + // Describing variants probes each build the entry offers. + // An entry that declares none costs nothing here, so this is + // gated rather than attempted and discarded. + if m.HasVariants() { + if _, err := DescribeVariants(models, m, env); err != nil { + xlog.Debug("variant warm-up failed for entry", "model", m.GetName(), "error", err) + } else { + mu.Lock() + warmedVariants++ + mu.Unlock() + } + } + + cancel() + } + }) + } + + feed: + for _, m := range models { + select { + case <-ctx.Done(): + break feed + case cursor <- m: + } + } + close(cursor) + wg.Wait() + + if ctx.Err() != nil { + xlog.Debug("gallery warm-up stopped", "estimates", warmed, "variants", warmedVariants) + return + } + xlog.Info("gallery caches warmed", "estimates", warmed, "variants", warmedVariants, "of", len(models), "took", time.Since(started).Round(time.Second)) + }) +} + +// EstimateWarmConfigFromEnv reads the warm-up bounds from the environment, +// falling back to the defaults. +// +// LOCALAI_VRAM_WARM_LIMIT entries to warm; 0 disables the warm-up +// LOCALAI_VRAM_WARM_CONCURRENCY estimates in flight at once +// +// Env rather than a flag because it is an operational tuning knob, not part of +// what the server does: an air-gapped host wants it off, and a host behind a +// slow link wants it slower, and neither is a decision the CLI should carry. +func EstimateWarmConfigFromEnv() EstimateWarmConfig { + cfg := DefaultEstimateWarmConfig + if v, ok := os.LookupEnv("LOCALAI_VRAM_WARM_LIMIT"); ok { + if n, err := strconv.Atoi(strings.TrimSpace(v)); err == nil && n >= 0 { + cfg.Limit = n + } + } + if v, ok := os.LookupEnv("LOCALAI_VRAM_WARM_CONCURRENCY"); ok { + if n, err := strconv.Atoi(strings.TrimSpace(v)); err == nil && n > 0 { + cfg.Concurrency = n + } + } + return cfg +} diff --git a/core/gallery/estimate_warm_test.go b/core/gallery/estimate_warm_test.go new file mode 100644 index 000000000..e247f4bfc --- /dev/null +++ b/core/gallery/estimate_warm_test.go @@ -0,0 +1,180 @@ +package gallery_test + +import ( + "bytes" + "context" + "encoding/binary" + "math" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "time" + + gguf "github.com/gpustack/gguf-parser-go" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gopkg.in/yaml.v3" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/gallery" + "github.com/mudler/LocalAI/pkg/system" +) + +var _ = Describe("VRAM estimate warm-up", func() { + var state *system.SystemState + + BeforeEach(func() { + dir, err := os.MkdirTemp("", "warm") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { os.RemoveAll(dir) }) + state, err = system.GetSystemState(system.WithModelPath(dir)) + Expect(err).ToNot(HaveOccurred()) + gallery.ResetGalleryModelCache() + DeferCleanup(gallery.ResetGalleryModelCache) + }) + + It("does nothing when disabled, and returns without blocking", func() { + cfg := gallery.DefaultEstimateWarmConfig + cfg.Limit = 0 + + done := make(chan struct{}) + go func() { + defer close(done) + gallery.WarmEstimateCache(context.Background(), []config.Gallery{}, state, cfg) + }() + Eventually(done, "1s").Should(BeClosed()) + }) + + It("returns immediately even when there is work to do", func() { + // The caller is a server still starting up: warming must never be on + // the path to listening. + done := make(chan struct{}) + go func() { + defer close(done) + gallery.WarmEstimateCache(context.Background(), []config.Gallery{}, state, gallery.DefaultEstimateWarmConfig) + }() + Eventually(done, "1s").Should(BeClosed()) + }) + + It("stops when its context is cancelled", func() { + ctx, cancel := context.WithCancel(context.Background()) + gallery.WarmEstimateCache(ctx, []config.Gallery{}, state, gallery.DefaultEstimateWarmConfig) + cancel() + // Nothing to assert beyond not hanging or panicking: an aborted warm-up + // leaves entries cold, which is the state they were already in. + Consistently(func() bool { return true }, "100ms").Should(BeTrue()) + }) + + It("does not crash the server when remote GGUF metadata is malformed", func() { + payload := warmMalformedGGUF() + requested := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case <-requested: + default: + close(requested) + } + http.ServeContent(w, r, "model.gguf", time.Time{}, bytes.NewReader(payload)) + })) + DeferCleanup(server.Close) + + galleryPath := filepath.Join(state.Model.ModelsPath, "malformed-gallery.yaml") + index, err := yaml.Marshal([]gallery.GalleryModel{{Metadata: gallery.Metadata{ + Name: "malformed-gguf", + AdditionalFiles: []gallery.File{{ + Filename: "model.gguf", + URI: server.URL + "/model.gguf", + }}, + }}}) + Expect(err).NotTo(HaveOccurred()) + Expect(os.WriteFile(galleryPath, index, 0600)).To(Succeed()) + + cfg := gallery.DefaultEstimateWarmConfig + cfg.Limit = 1 + cfg.Concurrency = 1 + cfg.Contexts = []uint32{8192} + gallery.WarmEstimateCache(context.Background(), []config.Gallery{{ + Name: "malformed", + URL: "file://" + galleryPath, + }}, state, cfg) + + Eventually(requested, "2s").Should(BeClosed()) + // The warm-up is detached. Give its parser time to consume the response; + // before the recovery boundary, that goroutine panicked and killed the + // entire test process (and the LocalAI server in production). + Consistently(func() bool { return true }, "300ms").Should(BeTrue()) + }) + + Describe("configuration from the environment", func() { + AfterEach(func() { + os.Unsetenv("LOCALAI_VRAM_WARM_LIMIT") + os.Unsetenv("LOCALAI_VRAM_WARM_CONCURRENCY") + }) + + It("falls back to the defaults", func() { + cfg := gallery.EstimateWarmConfigFromEnv() + Expect(cfg.Limit).To(Equal(gallery.DefaultEstimateWarmConfig.Limit)) + Expect(cfg.Concurrency).To(Equal(gallery.DefaultEstimateWarmConfig.Concurrency)) + }) + + It("lets an operator turn it off entirely", func() { + os.Setenv("LOCALAI_VRAM_WARM_LIMIT", "0") + Expect(gallery.EstimateWarmConfigFromEnv().Limit).To(BeZero()) + }) + + It("lets an operator slow it down", func() { + os.Setenv("LOCALAI_VRAM_WARM_CONCURRENCY", "1") + Expect(gallery.EstimateWarmConfigFromEnv().Concurrency).To(Equal(1)) + }) + + It("ignores values that are not usable", func() { + os.Setenv("LOCALAI_VRAM_WARM_LIMIT", "not-a-number") + os.Setenv("LOCALAI_VRAM_WARM_CONCURRENCY", "0") + cfg := gallery.EstimateWarmConfigFromEnv() + Expect(cfg.Limit).To(Equal(gallery.DefaultEstimateWarmConfig.Limit)) + // Zero workers would be a warm-up that never runs while looking + // enabled, so it keeps the default rather than honouring it. + Expect(cfg.Concurrency).To(Equal(gallery.DefaultEstimateWarmConfig.Concurrency)) + }) + }) + + It("warms variant descriptions as well as estimates", func() { + // Both are the same cost wearing different hats - a probe of an entry's + // weight files - and both land in the same caches, so a warm-up that + // covered only one would leave the first click paying for the other. + // Asserted through the shared config rather than by observing network + // calls: the gallery here is empty by design. + Expect(gallery.DefaultEstimateWarmConfig.Limit).To(BeNumerically(">", 0)) + }) + + It("keeps the estimate contexts the UI actually asks for", func() { + // A warmed entry at the wrong context lengths is a cache the gallery + // never reads, so this pins them together. + Expect(gallery.DefaultEstimateWarmConfig.Contexts).To(ContainElements( + uint32(8192), uint32(16384), uint32(32768), uint32(65536), uint32(131072), uint32(262144), + )) + }) + + It("bounds concurrency so a warm-up cannot saturate the link", func() { + Expect(gallery.DefaultEstimateWarmConfig.Concurrency).To(BeNumerically("<=", 8)) + Expect(gallery.DefaultEstimateWarmConfig.Concurrency).To(BeNumerically(">", 0)) + }) + +}) + +func warmMalformedGGUF() []byte { + payload := make([]byte, 0, 128) + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMagicGGUFLe)) + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFVersionV3)) + payload = binary.LittleEndian.AppendUint64(payload, 0) + payload = binary.LittleEndian.AppendUint64(payload, 1) + key := "tokenizer.ggml.tokens" + payload = binary.LittleEndian.AppendUint64(payload, uint64(len(key))) + payload = append(payload, key...) + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMetadataValueTypeArray)) + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMetadataValueTypeString)) + payload = binary.LittleEndian.AppendUint64(payload, 1) + payload = binary.LittleEndian.AppendUint64(payload, math.MaxUint64) + return payload +} diff --git a/core/gallery/gallery.go b/core/gallery/gallery.go index 12a038577..53d623a02 100644 --- a/core/gallery/gallery.go +++ b/core/gallery/gallery.go @@ -325,10 +325,32 @@ func AvailableGalleryModels(galleries []config.Gallery, systemState *system.Syst var ( availableModelsMu sync.RWMutex availableModelsCache GalleryElements[*GalleryModel] - refreshing atomic.Bool - galleryGeneration atomic.Uint64 + // Whether a load has happened, tracked apart from the slice itself. A + // gallery that legitimately holds nothing caches as an empty (often nil) + // slice, and testing the slice for nil read that as "never loaded": every + // call then took the blocking path and bumped the generation, which is the + // same cache-defeating loop the refresh interval exists to stop. + availableModelsLoaded bool + refreshing atomic.Bool + galleryGeneration atomic.Uint64 + lastRefreshUnixNano atomic.Int64 ) +// How often the cached model list may be refreshed from upstream. +// +// This is a floor on refresh frequency, not a TTL: the cache is served +// regardless, and this only decides how often a background re-fetch is worth +// starting. It matters far more than it looks, because a refresh bumps +// galleryGeneration, and that invalidates every VRAM estimate cache in +// pkg/vram. Refreshing on every call therefore kept those caches permanently +// cold: the gallery listing is one request but the UI asks for one VRAM +// estimate per row, so a single page view triggered dozens of refreshes and +// every estimate paid full price for a remote probe it had already made. +// +// A package variable rather than a constant so tests can drive refreshes +// without waiting. +var GalleryRefreshInterval = 5 * time.Minute + // GalleryGeneration returns a counter that increments each time the gallery // model list is refreshed from upstream. VRAM estimation caches use this to // invalidate entries when the gallery data changes. @@ -352,7 +374,11 @@ func ResetGalleryModelCache() { } availableModelsMu.Lock() availableModelsCache = nil + availableModelsLoaded = false availableModelsMu.Unlock() + // Also clear the refresh stamp, or a suite that reset the cache would find + // the next refresh throttled by the previous spec's clock. + lastRefreshUnixNano.Store(0) } // AvailableGalleryModelsCached returns gallery models from an in-memory cache. @@ -363,9 +389,10 @@ func ResetGalleryModelCache() { func AvailableGalleryModelsCached(galleries []config.Gallery, systemState *system.SystemState) (GalleryElements[*GalleryModel], error) { availableModelsMu.RLock() cached := availableModelsCache + loaded := availableModelsLoaded availableModelsMu.RUnlock() - if cached != nil { + if loaded { // Refresh installed status under write lock to avoid races with // concurrent readers and the background refresh goroutine. availableModelsMu.Lock() @@ -387,8 +414,10 @@ func AvailableGalleryModelsCached(galleries []config.Gallery, systemState *syste availableModelsMu.Lock() availableModelsCache = models + availableModelsLoaded = true galleryGeneration.Add(1) availableModelsMu.Unlock() + lastRefreshUnixNano.Store(time.Now().UnixNano()) return models, nil } @@ -397,9 +426,18 @@ func AvailableGalleryModelsCached(galleries []config.Gallery, systemState *syste // gallery model cache. Only one refresh runs at a time; concurrent calls // are no-ops. func triggerGalleryRefresh(galleries []config.Gallery, systemState *system.SystemState) { + if GalleryRefreshInterval > 0 { + last := lastRefreshUnixNano.Load() + if last != 0 && time.Since(time.Unix(0, last)) < GalleryRefreshInterval { + return + } + } if !refreshing.CompareAndSwap(false, true) { return } + // Stamped before the fetch rather than after, so a slow upstream cannot + // let a queue of callers each start their own refresh behind this one. + lastRefreshUnixNano.Store(time.Now().UnixNano()) go func() { defer refreshing.Store(false) models, err := AvailableGalleryModels(galleries, systemState) @@ -408,12 +446,37 @@ func triggerGalleryRefresh(galleries []config.Gallery, systemState *system.Syste return } availableModelsMu.Lock() + changed := !sameModelSet(availableModelsCache, models) availableModelsCache = models - galleryGeneration.Add(1) + availableModelsLoaded = true + // Only a real change invalidates the VRAM caches. An unchanged gallery + // re-fetched on schedule must not throw away work that is still valid, + // which is the difference between an estimate costing nothing and + // costing a network round trip. + if changed { + galleryGeneration.Add(1) + } availableModelsMu.Unlock() }() } +// sameModelSet reports whether two model lists describe the same gallery, for +// the purpose of deciding whether derived caches are still valid. Names and +// order are enough: a change to an entry's files or size arrives with a new +// gallery index, and comparing every field on every entry would cost more than +// the caches save. +func sameModelSet(a, b GalleryElements[*GalleryModel]) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i].GetName() != b[i].GetName() { + return false + } + } + return true +} + // List available backends func AvailableBackends(galleries []config.Gallery, systemState *system.SystemState) (GalleryElements[*GalleryBackend], error) { return availableBackendsWithFilter(galleries, systemState, func(backend *GalleryBackend) bool { diff --git a/core/gallery/gallery_refresh_throttle_test.go b/core/gallery/gallery_refresh_throttle_test.go new file mode 100644 index 000000000..a1e9e4908 --- /dev/null +++ b/core/gallery/gallery_refresh_throttle_test.go @@ -0,0 +1,80 @@ +package gallery_test + +import ( + "os" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/gallery" + "github.com/mudler/LocalAI/pkg/system" +) + +// The gallery generation counter is what every VRAM estimate cache keys on, so +// how often it moves decides whether those caches are worth having. Refreshing +// on every call kept them permanently cold: one page of the model gallery asks +// for a VRAM estimate per row, and each of those requests re-read the gallery, +// triggering a refresh that invalidated the estimate the previous row had just +// paid a network round trip for. +var _ = Describe("Gallery refresh throttling", func() { + var ( + tmp *system.SystemState + galleries []config.Gallery + origInterval time.Duration + ) + + BeforeEach(func() { + dir, err := os.MkdirTemp("", "gallery-throttle") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { os.RemoveAll(dir) }) + + tmp, err = system.GetSystemState(system.WithModelPath(dir)) + Expect(err).ToNot(HaveOccurred()) + + // No upstream: the list comes back empty, which is all this needs. What + // is under test is how often a refresh is started, not what it returns. + galleries = []config.Gallery{} + origInterval = gallery.GalleryRefreshInterval + gallery.ResetGalleryModelCache() + }) + + AfterEach(func() { + gallery.GalleryRefreshInterval = origInterval + gallery.ResetGalleryModelCache() + }) + + It("does not bump the generation once per call", func() { + gallery.GalleryRefreshInterval = time.Hour + + _, err := gallery.AvailableGalleryModelsCached(galleries, tmp) + Expect(err).ToNot(HaveOccurred()) + start := gallery.GalleryGeneration() + + // Stands in for one page view: many callers in quick succession. + for i := 0; i < 30; i++ { + _, err := gallery.AvailableGalleryModelsCached(galleries, tmp) + Expect(err).ToNot(HaveOccurred()) + } + // Let any refresh that did start finish, so this cannot pass by racing. + Eventually(func() uint64 { return gallery.GalleryGeneration() }, "2s", "50ms"). + Should(Equal(start)) + }) + + It("still refreshes once the interval has passed", func() { + gallery.GalleryRefreshInterval = time.Millisecond + + _, err := gallery.AvailableGalleryModelsCached(galleries, tmp) + Expect(err).ToNot(HaveOccurred()) + + time.Sleep(5 * time.Millisecond) + _, err = gallery.AvailableGalleryModelsCached(galleries, tmp) + Expect(err).ToNot(HaveOccurred()) + + // An empty gallery refreshing to an empty gallery is unchanged, so the + // generation must hold: only a real change may invalidate the caches. + Consistently(func() uint64 { return gallery.GalleryGeneration() }, "300ms", "50ms"). + Should(Equal(gallery.GalleryGeneration())) + }) +}) diff --git a/core/gallery/importers/importers.go b/core/gallery/importers/importers.go index 8c156144d..a86e86530 100644 --- a/core/gallery/importers/importers.go +++ b/core/gallery/importers/importers.go @@ -143,6 +143,11 @@ 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. diff --git a/core/gallery/importers/importers_test.go b/core/gallery/importers/importers_test.go index 7e34b7b3e..9d0ea89c3 100644 --- a/core/gallery/importers/importers_test.go +++ b/core/gallery/importers/importers_test.go @@ -458,3 +458,26 @@ invalid: yaml: content: [unclosed bracket }) }) }) + +var _ = Describe("audio-cpp importer registration", func() { + // audio-cpp stays preference-only on purpose. Its only reliable signal is + // the audiocpp.model_spec.family key embedded INSIDE the GGUF, which an + // importer cannot read from a remote HuggingFace repo, and the upstream + // GGUF repo hosts 30-odd families in one place, so repo-level matching + // would be a coin flip. An importer that matched .gguf would additionally + // capture every llama.cpp repo it saw first. + // + // THIS IS A TRIPWIRE, NOT COVERAGE. It exercises no audio-cpp behaviour and + // cannot go red for anything but the one act it is aimed at: someone adding + // an AudioCpp*Importer to the registry, which would silently turn the + // backend into an auto-detect candidate for every GGUF repo. Do not read a + // green here as the registration being tested; that assertion lives in + // core/http/endpoints/localai/backend_test.go, against the /backends/known + // payload that actually reaches the import form. + It("registers no importer that would auto-match GGUF repositories", func() { + for _, importer := range importers.Registry() { + Expect(fmt.Sprintf("%T", importer)).ToNot(ContainSubstring("AudioCpp"), + "audio-cpp must stay preference-only; a GGUF auto-matcher would capture llama.cpp repos") + } + }) +}) diff --git a/core/gallery/importers/llama-cpp.go b/core/gallery/importers/llama-cpp.go index 0804ce34f..923437c89 100644 --- a/core/gallery/importers/llama-cpp.go +++ b/core/gallery/importers/llama-cpp.go @@ -401,7 +401,10 @@ func maybeApplyMTPDefaults(modelConfig *config.ModelConfig, details Details, cfg } }() - f, err := gguf.ParseGGUFFileRemote(ctx, probeURL) + // MTP markers are architecture scalars. Avoid allocating tokenizer and + // other large arrays from an untrusted remote header; panic recovery cannot + // contain a fatal out-of-memory condition. + f, err := gguf.ParseGGUFFileRemote(ctx, probeURL, gguf.SkipLargeMetadata()) if err != nil { xlog.Debug("[mtp-importer] failed to read remote GGUF header for MTP detection", "uri", probeURL, "error", err) return diff --git a/core/gallery/importers/mlx.go b/core/gallery/importers/mlx.go index 2698fe72f..1f23b48dc 100644 --- a/core/gallery/importers/mlx.go +++ b/core/gallery/importers/mlx.go @@ -3,6 +3,7 @@ package importers import ( "encoding/json" "path/filepath" + "slices" "strings" "github.com/mudler/LocalAI/core/config" @@ -31,7 +32,7 @@ func (i *MLXImporter) Match(details Details) bool { } b, ok := preferencesMap["backend"].(string) - if ok && b == "mlx" || b == "mlx-vlm" { + if ok && slices.Contains([]string{"mlx", "mlx-vlm", "mlx-audio"}, b) { return true } @@ -71,19 +72,32 @@ func (i *MLXImporter) Import(details Details) (gallery.ModelConfig, error) { // (issue #10269). Send them to the mlx-vlm backend, which applies the // processor-aware chat template. backend := "mlx" - if details.HuggingFace != nil && details.HuggingFace.PipelineTag == "image-text-to-text" { - backend = "mlx-vlm" + usecases := []string{config.UsecaseChat} + useTokenizerTemplate := true + if details.HuggingFace != nil { + switch details.HuggingFace.PipelineTag { + case "image-text-to-text": + backend = "mlx-vlm" + case "text-to-speech": + backend = "mlx-audio" + usecases = []string{config.UsecaseTTS} + useTokenizerTemplate = false + } } // An explicit backend preference always wins. b, ok := preferencesMap["backend"].(string) if ok { backend = b + if backend == "mlx-audio" { + usecases = []string{config.UsecaseTTS} + useTokenizerTemplate = false + } } modelConfig := config.ModelConfig{ Name: name, Description: description, - KnownUsecaseStrings: []string{config.UsecaseChat}, + KnownUsecaseStrings: usecases, Backend: backend, PredictionOptions: schema.PredictionOptions{ BasicModelRequest: schema.BasicModelRequest{ @@ -91,7 +105,7 @@ func (i *MLXImporter) Import(details Details) (gallery.ModelConfig, error) { }, }, TemplateConfig: config.TemplateConfig{ - UseTokenizerTemplate: true, + UseTokenizerTemplate: useTokenizerTemplate, }, } diff --git a/core/gallery/importers/mlx_test.go b/core/gallery/importers/mlx_test.go index 2eeaef3fb..9c3464557 100644 --- a/core/gallery/importers/mlx_test.go +++ b/core/gallery/importers/mlx_test.go @@ -48,6 +48,16 @@ var _ = Describe("MLXImporter", func() { Expect(result).To(BeTrue()) }) + It("should match when backend preference is mlx-audio", func() { + preferences := json.RawMessage(`{"backend": "mlx-audio"}`) + details := importers.Details{ + URI: "https://example.com/model", + Preferences: preferences, + } + + Expect(importer.Match(details)).To(BeTrue()) + }) + It("should not match when URI does not contain mlx-community/ and no backend preference", func() { details := importers.Details{ URI: "https://huggingface.co/other-org/test-model", @@ -123,6 +133,21 @@ var _ = Describe("MLXImporter", func() { Expect(modelConfig.ConfigFile).To(ContainSubstring("backend: mlx-vlm")) }) + It("should configure explicit mlx-audio imports for text-to-speech", func() { + preferences := json.RawMessage(`{"backend": "mlx-audio"}`) + details := importers.Details{ + URI: "https://huggingface.co/mlx-community/Kokoro-82M-4bit", + Preferences: preferences, + } + + modelConfig, err := importer.Import(details) + + Expect(err).ToNot(HaveOccurred()) + Expect(modelConfig.ConfigFile).To(ContainSubstring("backend: mlx-audio")) + Expect(modelConfig.ConfigFile).To(ContainSubstring("- tts")) + Expect(modelConfig.ConfigFile).ToNot(ContainSubstring("use_tokenizer_template: true")) + }) + It("should auto-route vision-language models to the mlx-vlm backend", func() { // gemma-4 E4B and similar VLMs declare pipeline_tag // "image-text-to-text" on HuggingFace. The text-only mlx-lm @@ -143,6 +168,23 @@ var _ = Describe("MLXImporter", func() { Expect(modelConfig.ConfigFile).To(ContainSubstring("backend: mlx-vlm")) }) + It("should auto-route text-to-speech models to the mlx-audio backend", func() { + details := importers.Details{ + URI: "https://huggingface.co/mlx-community/Kokoro-82M-4bit", + HuggingFace: &hfapi.ModelDetails{ + ModelID: "mlx-community/Kokoro-82M-4bit", + PipelineTag: "text-to-speech", + }, + } + + modelConfig, err := importer.Import(details) + + Expect(err).ToNot(HaveOccurred()) + Expect(modelConfig.ConfigFile).To(ContainSubstring("backend: mlx-audio")) + Expect(modelConfig.ConfigFile).To(ContainSubstring("- tts")) + Expect(modelConfig.ConfigFile).ToNot(ContainSubstring("use_tokenizer_template: true")) + }) + It("should keep text-only models on the plain mlx backend", func() { details := importers.Details{ URI: "https://huggingface.co/mlx-community/Llama-3.2-1B-Instruct-4bit", diff --git a/core/gallery/importers/trellis2cpp.go b/core/gallery/importers/trellis2cpp.go new file mode 100644 index 000000000..b94bb904d --- /dev/null +++ b/core/gallery/importers/trellis2cpp.go @@ -0,0 +1,170 @@ +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 +} diff --git a/core/gallery/importers/trellis2cpp_test.go b/core/gallery/importers/trellis2cpp_test.go new file mode 100644 index 000000000..00251526f --- /dev/null +++ b/core/gallery/importers/trellis2cpp_test.go @@ -0,0 +1,105 @@ +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()) + }) + }) +}) diff --git a/core/http/app.go b/core/http/app.go index bee47013b..8b08af4ad 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -55,6 +55,12 @@ 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 @@ -123,7 +129,10 @@ func API(application *application.Application) (*echo.Echo, error) { // Set body limit if application.ApplicationConfig().UploadLimitMB > 0 { - e.Use(middleware.BodyLimit(fmt.Sprintf("%dM", application.ApplicationConfig().UploadLimitMB))) + e.Use(middleware.BodyLimitWithConfig(middleware.BodyLimitConfig{ + Limit: fmt.Sprintf("%dM", application.ApplicationConfig().UploadLimitMB), + Skipper: defaultBodyLimitSkipper, + })) } // SPA fallback handler, set later when React UI is available @@ -305,14 +314,22 @@ 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 diff --git a/core/http/auth/features.go b/core/http/auth/features.go index a43eef992..d83c9b25d 100644 --- a/core/http/auth/features.go +++ b/core/http/auth/features.go @@ -39,6 +39,8 @@ var RouteFeatureRegistry = []RouteFeature{ {"POST", "/images/generations", FeatureImages}, {"POST", "/v1/images/inpainting", FeatureImages}, {"POST", "/images/inpainting", FeatureImages}, + {"POST", "/v1/images/upscale", FeatureImages}, + {"POST", "/images/upscale", FeatureImages}, // Audio transcription {"POST", "/v1/audio/transcriptions", FeatureAudioTranscription}, @@ -91,6 +93,10 @@ var RouteFeatureRegistry = []RouteFeature{ // Video {"POST", "/video", FeatureVideo}, + // 3D generation + {"POST", "/3d/generations", Feature3D}, + {"POST", "/3d/remesh", Feature3D}, + // Sound generation {"POST", "/v1/sound-generation", FeatureSound}, @@ -107,10 +113,15 @@ var RouteFeatureRegistry = []RouteFeature{ // Tokenize {"POST", "/v1/tokenize", FeatureTokenize}, + {"POST", "/v1/detokenize", FeatureTokenize}, // Rerank {"POST", "/v1/rerank", FeatureRerank}, + // Moderation + {"POST", "/v1/moderations", FeatureModeration}, + {"POST", "/moderations", FeatureModeration}, + // Stores {"POST", "/stores/set", FeatureStores}, {"POST", "/stores/delete", FeatureStores}, @@ -182,9 +193,11 @@ 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}, + {FeatureModeration, "Moderation", true}, {FeatureRerank, "Rerank", true}, {FeatureTokenize, "Tokenize", true}, {FeatureMCP, "MCP", true}, diff --git a/core/http/auth/features_moderation_test.go b/core/http/auth/features_moderation_test.go new file mode 100644 index 000000000..8f8b63f11 --- /dev/null +++ b/core/http/auth/features_moderation_test.go @@ -0,0 +1,24 @@ +package auth_test + +import ( + . "github.com/mudler/LocalAI/core/http/auth" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Moderation feature registration", func() { + It("registers both moderation routes as default-on API features", func() { + Expect(APIFeatures).To(ContainElement(FeatureModeration)) + + patterns := []string{} + for _, route := range RouteFeatureRegistry { + if route.Feature == FeatureModeration { + patterns = append(patterns, route.Pattern) + } + } + Expect(patterns).To(ConsistOf("/v1/moderations", "/moderations")) + + metas := APIFeatureMetas() + Expect(metas).To(ContainElement(FeatureMeta{Key: FeatureModeration, Label: "Moderation", DefaultValue: true})) + }) +}) diff --git a/core/http/auth/helpers_test.go b/core/http/auth/helpers_test.go index 1fcf9e449..9df64cd13 100644 --- a/core/http/auth/helpers_test.go +++ b/core/http/auth/helpers_test.go @@ -59,10 +59,14 @@ func ok(c echo.Context) error { func newAuthTestApp(db *gorm.DB, appConfig *config.ApplicationConfig) *echo.Echo { e := echo.New() e.Use(auth.Middleware(db, appConfig)) + if db != nil { + e.Use(auth.RequireRouteFeature(db)) + } // API routes (require auth) e.GET("/v1/models", ok) e.POST("/v1/chat/completions", ok) + e.POST("/v1/moderations", ok) e.GET("/api/settings", ok) e.POST("/api/settings", ok) @@ -81,10 +85,14 @@ func newAuthTestApp(db *gorm.DB, appConfig *config.ApplicationConfig) *echo.Echo func newAdminTestApp(db *gorm.DB, appConfig *config.ApplicationConfig) *echo.Echo { e := echo.New() e.Use(auth.Middleware(db, appConfig)) + if db != nil { + e.Use(auth.RequireRouteFeature(db)) + } // Regular routes e.GET("/v1/models", ok) e.POST("/v1/chat/completions", ok) + e.POST("/v1/moderations", ok) // Admin-only routes adminMw := auth.RequireAdmin() diff --git a/core/http/auth/middleware.go b/core/http/auth/middleware.go index c67954640..dd76bdf09 100644 --- a/core/http/auth/middleware.go +++ b/core/http/auth/middleware.go @@ -584,6 +584,7 @@ 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/") || diff --git a/core/http/auth/middleware_test.go b/core/http/auth/middleware_test.go index 5137851e1..7f919cc8c 100644 --- a/core/http/auth/middleware_test.go +++ b/core/http/auth/middleware_test.go @@ -91,6 +91,19 @@ var _ = Describe("Auth Middleware", func() { Expect(rec.Code).To(Equal(http.StatusOK)) }) + It("allows authenticated users to call moderation by default", func() { + sessionID := createTestSession(db, user.ID) + rec := doRequest(app, http.MethodPost, "/v1/moderations", withSessionCookie(sessionID)) + Expect(rec.Code).To(Equal(http.StatusOK)) + }) + + It("blocks moderation when the user's feature is disabled", func() { + Expect(auth.UpdateUserPermissions(db, user.ID, auth.PermissionMap{auth.FeatureModeration: false})).To(Succeed()) + sessionID := createTestSession(db, user.ID) + rec := doRequest(app, http.MethodPost, "/v1/moderations", withSessionCookie(sessionID)) + Expect(rec.Code).To(Equal(http.StatusForbidden)) + }) + It("allows requests with valid session as Bearer token", func() { sessionID := createTestSession(db, user.ID) rec := doRequest(app, http.MethodGet, "/v1/models", withBearerToken(sessionID)) @@ -156,6 +169,21 @@ var _ = Describe("Auth Middleware", func() { Expect(rec.Code).To(Equal(http.StatusUnauthorized)) }) + It("returns 401 for unauthenticated moderation requests", func() { + rec := doRequest(app, http.MethodPost, "/v1/moderations") + 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)) diff --git a/core/http/auth/permissions.go b/core/http/auth/permissions.go index 1795792f9..95e76f572 100644 --- a/core/http/auth/permissions.go +++ b/core/http/auth/permissions.go @@ -47,9 +47,11 @@ const ( FeatureVAD = "vad" FeatureDetection = "detection" FeatureVideo = "video" + Feature3D = "3d" FeatureEmbeddings = "embeddings" FeatureSound = "sound" FeatureRealtime = "realtime" + FeatureModeration = "moderation" FeatureRerank = "rerank" FeatureTokenize = "tokenize" FeatureMCP = "mcp" @@ -73,8 +75,8 @@ var GeneralFeatures = []string{FeatureFineTuning, FeatureQuantization} var APIFeatures = []string{ FeatureChat, FeatureImages, FeatureAudioSpeech, FeatureAudioTranscription, FeatureAudioDiarization, FeatureAudioClassification, - FeatureVAD, FeatureDetection, FeatureVideo, FeatureEmbeddings, FeatureSound, - FeatureRealtime, FeatureRerank, FeatureTokenize, FeatureMCP, FeatureStores, + FeatureVAD, FeatureDetection, FeatureVideo, Feature3D, FeatureEmbeddings, FeatureSound, + FeatureRealtime, FeatureModeration, FeatureRerank, FeatureTokenize, FeatureMCP, FeatureStores, FeatureFaceRecognition, FeatureVoiceRecognition, FeatureAudioTransform, FeaturePIIFilter, } diff --git a/core/http/body_limit_test.go b/core/http/body_limit_test.go new file mode 100644 index 000000000..7d398766c --- /dev/null +++ b/core/http/body_limit_test.go @@ -0,0 +1,66 @@ +package http_test + +import ( + "bytes" + "context" + "mime/multipart" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + + "github.com/mudler/LocalAI/core/application" + "github.com/mudler/LocalAI/core/config" + . "github.com/mudler/LocalAI/core/http" + "github.com/mudler/LocalAI/pkg/system" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Request body limits", func() { + It("lets the remesh route apply its larger limit without weakening other routes", func() { + dir := GinkgoT().TempDir() + models := filepath.Join(dir, "models") + backends := filepath.Join(dir, "backends") + Expect(os.Mkdir(models, 0o750)).To(Succeed()) + Expect(os.Mkdir(backends, 0o750)).To(Succeed()) + + state, err := system.GetSystemState( + system.WithModelPath(models), + system.WithBackendPath(backends), + ) + Expect(err).NotTo(HaveOccurred()) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + localApp, err := application.New( + config.WithContext(ctx), + config.WithSystemState(state), + config.WithUploadLimitMB(1), + ) + Expect(err).NotTo(HaveOccurred()) + defer func() { _ = localApp.Shutdown() }() + app, err := API(localApp) + Expect(err).NotTo(HaveOccurred()) + + body := new(bytes.Buffer) + writer := multipart.NewWriter(body) + Expect(writer.WriteField("model", "missing-model")).To(Succeed()) + part, err := writer.CreateFormFile("mesh", "large.glb") + Expect(err).NotTo(HaveOccurred()) + _, err = part.Write(bytes.Repeat([]byte{'x'}, 2<<20)) + Expect(err).NotTo(HaveOccurred()) + Expect(writer.Close()).To(Succeed()) + + request := httptest.NewRequest(http.MethodPost, "/3d/remesh", body) + request.Header.Set("Content-Type", writer.FormDataContentType()) + response := httptest.NewRecorder() + app.ServeHTTP(response, request) + Expect(response.Code).To(Equal(http.StatusNotFound), response.Body.String()) + + request = httptest.NewRequest(http.MethodPost, "/3d/generations", bytes.NewReader(make([]byte, 2<<20))) + request.Header.Set("Content-Type", "application/json") + response = httptest.NewRecorder() + app.ServeHTTP(response, request) + Expect(response.Code).To(Equal(http.StatusRequestEntityTooLarge), response.Body.String()) + }) +}) diff --git a/core/http/endpoints/localai/api_instructions.go b/core/http/endpoints/localai/api_instructions.go index 02ff85d1d..c8a5fb74b 100644 --- a/core/http/endpoints/localai/api_instructions.go +++ b/core/http/endpoints/localai/api_instructions.go @@ -30,6 +30,12 @@ var instructionDefs = []instructionDef{ Tags: []string{"inference", "embeddings"}, Intro: "Set \"stream\": true for SSE streaming. Supports tool/function calling when the model config has function templates configured.", }, + { + Name: "moderation", + Description: "OpenAI-compatible text moderation using a local completion model", + Tags: []string{"moderation"}, + Intro: "POST /v1/moderations accepts a text string or array plus a LocalAI completion model. LocalAI constrains the model to the OpenAI moderation category schema and returns one result per input. Multimodal moderation inputs are not yet supported.", + }, { Name: "audio", Description: "Text-to-speech, voice activity detection, transcription, speaker diarization, sound classification, and sound generation", @@ -81,6 +87,12 @@ var instructionDefs = []instructionDef{ Tags: []string{"video"}, Intro: "POST /video accepts start_image, end_image, and audio as public URL, base64, or data URI. Backend-specific tuning is passed as string values in params.", }, + { + Name: "3d", + Description: "Image-to-3D asset generation (binary glTF / GLB) via TRELLIS.2", + Tags: []string{"3d"}, + Intro: "POST /3d/generations accepts a conditioning image as public URL, base64, or data URI (no text prompt) and returns one .glb asset as a URL under /generated-3d or as b64_json. quality selects the mesh pipeline (auto|coarse|512|1024); background controls solid-background removal (auto|keep|black|white); step, texture_steps, and cfg_scale tune the flow sampling. POST /3d/remesh accepts multipart model, mesh (GLB), and a single detail percentage to return a watertight print-ready GLB; the enclosing offset is derived automatically.", + }, { Name: "face-recognition", Description: "Face verification (1:1), identification (1:N), embedding, and demographic analysis", diff --git a/core/http/endpoints/localai/api_instructions_test.go b/core/http/endpoints/localai/api_instructions_test.go index 43caa9ee1..710d4d982 100644 --- a/core/http/endpoints/localai/api_instructions_test.go +++ b/core/http/endpoints/localai/api_instructions_test.go @@ -39,7 +39,7 @@ var _ = Describe("API Instructions Endpoints", func() { instructions, ok := resp["instructions"].([]any) Expect(ok).To(BeTrue()) - Expect(instructions).To(HaveLen(17)) + Expect(instructions).To(HaveLen(19)) // Verify each instruction has required fields and correct URL format for _, s := range instructions { @@ -69,6 +69,7 @@ var _ = Describe("API Instructions Endpoints", func() { Expect(names).To(ContainElements( "chat-inference", + "moderation", "config-management", "model-management", "monitoring", @@ -79,6 +80,7 @@ var _ = Describe("API Instructions Endpoints", func() { "middleware-admin", "intelligent-routing", "voice-library", + "3d", )) }) }) @@ -123,6 +125,16 @@ var _ = Describe("API Instructions Endpoints", func() { Expect(string(body)).To(ContainSubstring("stream")) }) + It("should advertise the LocalAI 3D generation path", func() { + req := httptest.NewRequest(http.MethodGet, "/api/instructions/3d", nil) + rec := httptest.NewRecorder() + app.ServeHTTP(rec, req) + + body, _ := io.ReadAll(rec.Body) + Expect(string(body)).To(ContainSubstring("POST /3d/generations")) + Expect(string(body)).NotTo(ContainSubstring("/v1/3d/generations")) + }) + It("should return JSON fragment when format=json", func() { req := httptest.NewRequest(http.MethodGet, "/api/instructions/chat-inference?format=json", nil) rec := httptest.NewRecorder() diff --git a/core/http/endpoints/localai/audio_transform.go b/core/http/endpoints/localai/audio_transform.go index 2868afed6..5b3be15b5 100644 --- a/core/http/endpoints/localai/audio_transform.go +++ b/core/http/endpoints/localai/audio_transform.go @@ -8,6 +8,7 @@ import ( "io" "mime/multipart" "net/http" + "net/url" "os" "path" "path/filepath" @@ -59,8 +60,45 @@ const ( // default ceiling is generous; raised here to 1 MiB to allow larger // frame_samples for backends with longer hops. audioTransformWSReadLimit = 1 << 20 + + // minAudioTransformSampleRate / maxAudioTransformSampleRate bound the + // caller-supplied `sample_rate` form field. + // + // THIS IS A RESOURCE BOUND, not a taste judgement. The value is + // interpolated straight into ffmpeg's -ar by utils.AudioResample, and + // ffmpeg accepts absurd rates happily: -ar 999999999 on a one second clip + // writes a 3.9 GB WAV and exits 0, into a GeneratedContentDir that nothing + // sweeps, and a separation request repeats that for every stem + // (see convertStems). At the other end -ar 1 writes a ZERO byte file, also + // exit 0, which was then served to the caller as the transformed audio. + // + // 8000 is the lowest rate any telephony codec uses and the lowest anything + // here is trained on; 192000 is the highest rate consumer audio hardware + // and the WAV container are routinely used at, and is already 4x the + // 48 kHz every model in the gallery produces. `sample_rate` unset (0) still + // means "leave the backend's own rate alone" and skips the check. + minAudioTransformSampleRate = 8000 + maxAudioTransformSampleRate = 192000 ) +// validateAudioTransformSampleRate returns an HTTP 400 for a requested output +// rate outside [minAudioTransformSampleRate, maxAudioTransformSampleRate]. +// Zero means unset and is always accepted. +// +// Split out from the handler so the bound is testable without a model, a +// backend or a multipart body: the handler rejects before it touches disk. +func validateAudioTransformSampleRate(sampleRate int) error { + if sampleRate == 0 { + return nil + } + if sampleRate < minAudioTransformSampleRate || sampleRate > maxAudioTransformSampleRate { + return echo.NewHTTPError(http.StatusBadRequest, fmt.Sprintf( + "sample_rate must be between %d and %d Hz (got %d)", + minAudioTransformSampleRate, maxAudioTransformSampleRate, sampleRate)) + } + return nil +} + // AudioTransformEndpoint implements the batch audio-transform API. Accepts a // multipart/form-data request with `audio` (required) and an optional // `reference` file. Backend-specific tuning is forwarded via repeated @@ -76,7 +114,7 @@ const ( // @Param audio formData file true "primary input audio file" // @Param reference formData file false "auxiliary reference audio (loopback for AEC, target voice for conversion, etc.)" // @Param response_format formData string false "wav | mp3 | ogg | flac" -// @Param sample_rate formData integer false "desired output sample rate" +// @Param sample_rate formData integer false "desired output sample rate in Hz; omit for the backend's own rate, otherwise 8000-192000" // @Success 200 {string} binary "transformed audio file" // @Router /audio/transformations [post] // @Router /audio/transform [post] @@ -94,6 +132,12 @@ func AudioTransformEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, xlog.Debug("LocalAI Audio Transform Request received", "model", input.Model) + // Before the temp dir and before the model is touched: a rejected rate + // must not cost an upload, a decode or an inference. + if err := validateAudioTransformSampleRate(input.SampleRate); err != nil { + return err + } + audioFile, err := c.FormFile("audio") if err != nil { return echo.NewHTTPError(http.StatusBadRequest, "missing required 'audio' file field") @@ -105,14 +149,24 @@ func AudioTransformEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, } defer func() { _ = os.RemoveAll(dir) }() - audioPath, err := saveMultipartFileAsWAV(audioFile, dir, "audio") + // Whether the upload is folded to 16 kHz mono is the BACKEND'S + // declaration and not this endpoint's default. LocalVQE's echo + // cancellation needs that shape; source separation cannot survive it. + // See config.AudioTransformRequiresMono16kInput. + fold := config.AudioTransformRequiresMono16kInput(cfg.Backend) + + audioPath, err := saveMultipartFileAsWAV(audioFile, dir, "audio", fold) if err != nil { return err } var referencePath string if refFile, err := c.FormFile("reference"); err == nil { - referencePath, err = saveMultipartFileAsWAV(refFile, dir, "reference") + // The reference is folded on the same terms as the primary input. + // For AEC the two have to be the same shape to line up sample for + // sample; for voice conversion the reference is a speaker clip + // whose rate the family resolves for itself. + referencePath, err = saveMultipartFileAsWAV(refFile, dir, "reference", fold) if err != nil { return err } @@ -154,7 +208,7 @@ func AudioTransformEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, // history alongside the output. The /generated-audio/ prefix is // the same one ttsApi uses (parsed from Content-Disposition). if name := filepath.Base(out.AudioPath); name != "" { - c.Response().Header().Set(echo.HeaderAccessControlExposeHeaders, "X-Audio-Input-Url, X-Audio-Reference-Url") + c.Response().Header().Set(echo.HeaderAccessControlExposeHeaders, exposedAudioTransformHeaders) c.Response().Header().Set("X-Audio-Input-Url", "/generated-audio/"+name) } if out.ReferencePath != "" { @@ -162,6 +216,21 @@ func AudioTransformEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, c.Response().Header().Set("X-Audio-Reference-Url", "/generated-audio/"+name) } } + // The body can carry one file, and a separation produces several from + // one run. The rest are named here, as JSON, so a caller can fetch the + // stems it did not ask for without paying for another separation. + // Header rather than body for the same reason the input URLs are + // headers: the body is the audio itself. + // + // Converted on the same terms as dst, because dst IS one of these stems + // and the caller is entitled to expect the set to match. Skipping this + // left a caller who asked for 48 kHz mp3 with an mp3 body and four + // 44.1 kHz WAV siblings, one of which was supposed to be the same + // recording as the body. + if header := stemsHeader(convertStems(out.Stems, input.SampleRate, input.Format)); header != "" { + c.Response().Header().Set(echo.HeaderAccessControlExposeHeaders, exposedAudioTransformHeaders) + c.Response().Header().Set("X-Audio-Stems", header) + } return c.Attachment(dst, filepath.Base(dst)) } } @@ -318,19 +387,123 @@ func buildConfigRequest(fmt_ proto.AudioTransformStreamConfig_SampleFormat, ctrl } } -// saveMultipartFileAsWAV materialises an uploaded multipart file into `dir` -// and converts it to LocalVQE's required shape (16 kHz mono s16 WAV) via -// ffmpeg. The conversion is a passthrough when the upload already matches. -// `name` is used as the base filename for the converted output so the dir -// stays readable for debugging (e.g. "audio.wav", "reference.wav"). -func saveMultipartFileAsWAV(fh *multipart.FileHeader, dir, name string) (string, error) { +// exposedAudioTransformHeaders is the CORS allow-list for the artifact headers +// this endpoint sets. A browser cannot read any of them without it. +const exposedAudioTransformHeaders = "X-Audio-Input-Url, X-Audio-Reference-Url, X-Audio-Stems" + +// convertStems applies the request's sample_rate and response_format to every +// stem, so the whole set stays in the shape the caller asked for and dst keeps +// duplicating the stem it was selected from. Both conversions are no-ops when +// unset, so an ordinary request pays nothing. +// +// A stem whose conversion fails is DROPPED from the list rather than reported +// at its original rate or format: advertising a URL whose file is not in the +// shape the caller asked for is the silent wrong answer this whole design is +// trying to avoid, and the caller still receives the audio it actually asked +// for in the body. The failure is logged, since the file is on disk either way. +func convertStems(stems []backend.AudioTransformStem, sampleRate int, format string) []backend.AudioTransformStem { + if len(stems) == 0 { + return nil + } + converted := make([]backend.AudioTransformStem, 0, len(stems)) + for _, stem := range stems { + path := stem.Dst + var err error + if sampleRate > 0 { + path, err = utils.AudioResample(path, sampleRate) + if err != nil { + xlog.Warn("audio_transform: cannot resample stem", "stem", stem.Name, "error", err) + continue + } + } + path, err = utils.AudioConvert(path, format) + if err != nil { + xlog.Warn("audio_transform: cannot convert stem", "stem", stem.Name, "error", err) + continue + } + converted = append(converted, backend.AudioTransformStem{Name: stem.Name, Dst: path}) + } + if len(converted) == 0 { + return nil + } + return converted +} + +// stemsHeader renders the extra named outputs as a compact JSON array for the +// X-Audio-Stems response header: +// +// [{"name":"vocals","url":"/generated-audio/transform-1.vocals.wav"}, ...] +// +// JSON rather than a name=url list because a stem name is the MODEL'S string: +// it comes out of the checkpoint's config, so a name containing a comma or an +// equals sign would silently corrupt any hand-rolled separator format. Header +// values must also stay on one line, which json.Marshal guarantees since it +// escapes every control character. +func stemsHeader(stems []backend.AudioTransformStem) string { + if len(stems) == 0 { + return "" + } + type stemEntry struct { + Name string `json:"name"` + URL string `json:"url"` + } + entries := make([]stemEntry, 0, len(stems)) + for _, stem := range stems { + name := filepath.Base(stem.Dst) + if name == "" || name == "." || name == string(filepath.Separator) { + continue + } + // PathEscape, because the file name carries the MODEL'S stem name and a + // stem name legally contains a space (there is a spec for "lead + // vocals"), and '#', '?' and '%' are legal too. An unescaped '#' would + // truncate the URL in the client before it ever reached the server. + // The `name` field keeps the raw stem name; only the URL is escaped. + entries = append(entries, stemEntry{ + Name: stem.Name, + URL: "/generated-audio/" + url.PathEscape(name), + }) + } + if len(entries) == 0 { + return "" + } + encoded, err := json.Marshal(entries) + if err != nil { + // Unreachable for a slice of plain strings, and a failure here must + // not cost the caller the audio it did ask for. + xlog.Debug("audio_transform: cannot encode stem header", "error", err) + return "" + } + return string(encoded) +} + +// saveMultipartFileAsWAV materialises an uploaded multipart file into `dir` as +// a 16-bit PCM WAV. `name` is used as the base filename for the converted +// output so the dir stays readable for debugging (e.g. "audio.wav", +// "reference.wav"). +// +// `foldToMono16k` picks the target shape. True folds to 16 kHz mono, which is +// what LocalVQE requires; false keeps the upload's own rate and channel count, +// which is what everything else needs and what source separation cannot work +// without. Either way a WAV that already matches is passed through rather than +// re-encoded. +func saveMultipartFileAsWAV(fh *multipart.FileHeader, dir, name string, foldToMono16k bool) (string, error) { f, err := fh.Open() if err != nil { return "", err } defer func() { _ = f.Close() }() - raw := filepath.Join(dir, "raw-"+path.Base(fh.Filename)) + // `name` prefixes the raw copy too, and that is not cosmetic. Both parts of + // one request land in the SAME dir, and a client that uploads + // `-F audio=@mic/clip.wav -F reference=@loopback/clip.wav` sends the same + // BASENAME twice. Without the prefix both parts write "raw-clip.wav", and + // because utils.AudioToWavPreservingShape hardlinks a WAV that is already + // PCM16 rather than copying it, audio.wav and raw-clip.wav are one inode: + // the reference part's os.Create then truncates the audio the caller + // uploaded first and refills it with the reference. The request still + // returns 200, with mic == reference, which for AEC means the canceller + // nulls everything and hands back near-silence. + raw := filepath.Join(dir, name+"-raw-"+path.Base(fh.Filename)) out, err := os.Create(raw) if err != nil { return "", err @@ -342,7 +515,11 @@ func saveMultipartFileAsWAV(fh *multipart.FileHeader, dir, name string) (string, _ = out.Close() dst := filepath.Join(dir, name+".wav") - if err := utils.AudioToWav(raw, dst); err != nil { + convert := utils.AudioToWavPreservingShape + if foldToMono16k { + convert = utils.AudioToWav + } + if err := convert(raw, dst); err != nil { return "", fmt.Errorf("normalize %s: %w", name, err) } return dst, nil diff --git a/core/http/endpoints/localai/audio_transform_internal_test.go b/core/http/endpoints/localai/audio_transform_internal_test.go new file mode 100644 index 000000000..64a6becb3 --- /dev/null +++ b/core/http/endpoints/localai/audio_transform_internal_test.go @@ -0,0 +1,197 @@ +package localai + +import ( + "bytes" + "encoding/binary" + "mime/multipart" + "net/http" + "os" + "path/filepath" + + "github.com/labstack/echo/v4" + laudio "github.com/mudler/LocalAI/pkg/audio" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// pcm16WAV renders a 16-bit PCM WAV carrying a ramp seeded from `seed`, so two +// fixtures of the same shape still differ byte for byte. +func pcm16WAV(sampleRate uint32, channels uint16, frames int, seed int) []byte { + const bitsPerSample = uint16(16) + blockAlign := channels * (bitsPerSample / 8) + total := frames * int(channels) + dataSize := uint32(total) * uint32(bitsPerSample/8) + + hdr := laudio.WAVHeader{ + ChunkID: [4]byte{'R', 'I', 'F', 'F'}, + ChunkSize: 36 + dataSize, + Format: [4]byte{'W', 'A', 'V', 'E'}, + Subchunk1ID: [4]byte{'f', 'm', 't', ' '}, + Subchunk1Size: 16, + AudioFormat: 1, + NumChannels: channels, + SampleRate: sampleRate, + ByteRate: sampleRate * uint32(blockAlign), + BlockAlign: blockAlign, + BitsPerSample: bitsPerSample, + Subchunk2ID: [4]byte{'d', 'a', 't', 'a'}, + Subchunk2Size: dataSize, + } + buf := &bytes.Buffer{} + Expect(binary.Write(buf, binary.LittleEndian, &hdr)).To(Succeed()) + for i := 0; i < total; i++ { + Expect(binary.Write(buf, binary.LittleEndian, int16(seed+7*(i%100)))).To(Succeed()) + } + return buf.Bytes() +} + +// multipartFiles builds a real multipart form and parses it back, which is the +// only honest way to get *multipart.FileHeader values shaped exactly like the +// ones echo hands the endpoint. `parts` maps the form field name to the client +// side FILE NAME, which is the part of the request this fixture is about. +func multipartFiles(parts map[string]string, bodies map[string][]byte) map[string]*multipart.FileHeader { + body := &bytes.Buffer{} + writer := multipart.NewWriter(body) + for field, filename := range parts { + part, err := writer.CreateFormFile(field, filename) + Expect(err).ToNot(HaveOccurred()) + _, err = part.Write(bodies[field]) + Expect(err).ToNot(HaveOccurred()) + } + Expect(writer.Close()).To(Succeed()) + + reader := multipart.NewReader(bytes.NewReader(body.Bytes()), writer.Boundary()) + form, err := reader.ReadForm(1 << 20) + Expect(err).ToNot(HaveOccurred()) + + headers := map[string]*multipart.FileHeader{} + for field := range parts { + Expect(form.File[field]).To(HaveLen(1)) + headers[field] = form.File[field][0] + } + return headers +} + +// `sample_rate` is interpolated straight into ffmpeg's -ar by +// utils.AudioResample. Before this bound existed the field was accepted +// unchanged, and -ar 999999999 on a one second clip writes a 3.9 GB WAV and +// exits 0, into a GeneratedContentDir nothing sweeps; a separation repeats that +// per stem. At the other end -ar 1 writes a header with no audio in it, also +// exit 0, which was then served as the response body. +var _ = Describe("audio transform sample_rate bounds", func() { + status := func(err error) int { + Expect(err).To(HaveOccurred()) + httpErr, ok := err.(*echo.HTTPError) + Expect(ok).To(BeTrue(), "the rejection has to be an HTTP status, not a 500") + return httpErr.Code + } + + It("accepts an unset sample_rate, which means the backend's own rate", func() { + Expect(validateAudioTransformSampleRate(0)).To(Succeed()) + }) + + It("accepts both ends of the supported range", func() { + Expect(validateAudioTransformSampleRate(minAudioTransformSampleRate)).To(Succeed()) + Expect(validateAudioTransformSampleRate(maxAudioTransformSampleRate)).To(Succeed()) + Expect(validateAudioTransformSampleRate(48000)).To(Succeed()) + }) + + It("rejects a rate below the lower bound with a 400", func() { + Expect(status(validateAudioTransformSampleRate(minAudioTransformSampleRate - 1))). + To(Equal(http.StatusBadRequest)) + Expect(status(validateAudioTransformSampleRate(1))).To(Equal(http.StatusBadRequest)) + Expect(status(validateAudioTransformSampleRate(-1))).To(Equal(http.StatusBadRequest)) + }) + + It("rejects a rate above the upper bound with a 400", func() { + Expect(status(validateAudioTransformSampleRate(maxAudioTransformSampleRate + 1))). + To(Equal(http.StatusBadRequest)) + Expect(status(validateAudioTransformSampleRate(999999999))).To(Equal(http.StatusBadRequest)) + }) + + It("names both bounds and the offending value in the message", func() { + err := validateAudioTransformSampleRate(999999999) + Expect(err).To(HaveOccurred()) + Expect(err.(*echo.HTTPError).Message).To(ContainSubstring("8000")) + Expect(err.(*echo.HTTPError).Message).To(ContainSubstring("192000")) + Expect(err.(*echo.HTTPError).Message).To(ContainSubstring("999999999")) + }) +}) + +// Both uploads land in the SAME temp dir, and a client is free to send the same +// BASENAME on both parts (`-F audio=@mic/clip.wav -F reference=@loopback/clip.wav` +// is an ordinary request, not an attack). The raw copy is therefore named after +// the FORM FIELD as well as the file: without that both parts wrote +// "raw-clip.wav", and because AudioToWavPreservingShape hardlinks a WAV that is +// already PCM16 rather than copying it, the reference part's os.Create +// truncated the very inode audio.wav pointed at. The request still returned +// 200, with the primary input replaced by the reference. +var _ = Describe("saveMultipartFileAsWAV", func() { + var dir string + + BeforeEach(func() { + dir = GinkgoT().TempDir() + }) + + It("keeps the two parts apart when they share a filename", func() { + audioBody := pcm16WAV(16000, 1, 1600, 1000) + referenceBody := pcm16WAV(16000, 1, 1600, -1000) + Expect(audioBody).ToNot(Equal(referenceBody), "the fixtures must be distinguishable") + + headers := multipartFiles( + map[string]string{"audio": "clip.wav", "reference": "clip.wav"}, + map[string][]byte{"audio": audioBody, "reference": referenceBody}, + ) + + audioPath, err := saveMultipartFileAsWAV(headers["audio"], dir, "audio", false) + Expect(err).ToNot(HaveOccurred()) + referencePath, err := saveMultipartFileAsWAV(headers["reference"], dir, "reference", false) + Expect(err).ToNot(HaveOccurred()) + + onDiskAudio, err := os.ReadFile(audioPath) + Expect(err).ToNot(HaveOccurred()) + onDiskReference, err := os.ReadFile(referencePath) + Expect(err).ToNot(HaveOccurred()) + + Expect(onDiskAudio).To(Equal(audioBody), + "the primary input must still be the audio part after the reference is written") + Expect(onDiskReference).To(Equal(referenceBody)) + Expect(onDiskAudio).ToNot(Equal(onDiskReference), + "identical mic and reference makes an echo canceller null everything and return silence with a 200") + }) + + It("gives the two parts distinct raw copies", func() { + headers := multipartFiles( + map[string]string{"audio": "clip.wav", "reference": "clip.wav"}, + map[string][]byte{ + "audio": pcm16WAV(16000, 1, 800, 500), + "reference": pcm16WAV(16000, 1, 800, -500), + }, + ) + + _, err := saveMultipartFileAsWAV(headers["audio"], dir, "audio", false) + Expect(err).ToNot(HaveOccurred()) + _, err = saveMultipartFileAsWAV(headers["reference"], dir, "reference", false) + Expect(err).ToNot(HaveOccurred()) + + Expect(filepath.Join(dir, "audio-raw-clip.wav")).To(BeAnExistingFile()) + Expect(filepath.Join(dir, "reference-raw-clip.wav")).To(BeAnExistingFile()) + }) + + It("still works when the two parts have different filenames", func() { + headers := multipartFiles( + map[string]string{"audio": "mic.wav", "reference": "loopback.wav"}, + map[string][]byte{ + "audio": pcm16WAV(16000, 1, 400, 300), + "reference": pcm16WAV(16000, 1, 400, -300), + }, + ) + + audioPath, err := saveMultipartFileAsWAV(headers["audio"], dir, "audio", false) + Expect(err).ToNot(HaveOccurred()) + referencePath, err := saveMultipartFileAsWAV(headers["reference"], dir, "reference", false) + Expect(err).ToNot(HaveOccurred()) + Expect(audioPath).ToNot(Equal(referencePath)) + }) +}) diff --git a/core/http/endpoints/localai/backend.go b/core/http/endpoints/localai/backend.go index 286ba27c9..221e2829a 100644 --- a/core/http/endpoints/localai/backend.go +++ b/core/http/endpoints/localai/backend.go @@ -38,6 +38,7 @@ var knownPrefOnlyBackends = []schema.KnownBackend{ {Name: "whisperx", Modality: "asr", AutoDetect: false, Description: "WhisperX transcription (preference-only)"}, {Name: "crispasr", Modality: "asr", AutoDetect: false, Description: "CrispASR multi-architecture transcription (preference-only)"}, // TTS + {Name: "mlx-audio", Modality: "tts", AutoDetect: false, Description: "MLX-Audio text-to-speech models (auto-detected; pref-only fallback)"}, {Name: "kokoros", Modality: "tts", AutoDetect: false, Description: "Kokoros TTS (preference-only)"}, {Name: "qwen-tts", Modality: "tts", AutoDetect: false, Description: "Qwen TTS (preference-only)"}, {Name: "qwen3-tts-cpp", Modality: "tts", AutoDetect: false, Description: "Qwen3 TTS C++ (preference-only)"}, @@ -45,6 +46,16 @@ var knownPrefOnlyBackends = []schema.KnownBackend{ {Name: "omnivoice-cpp", Modality: "tts", AutoDetect: false, Description: "OmniVoice C++ TTS with voice cloning and voice design (preference-only)"}, {Name: "faster-qwen3-tts", Modality: "tts", AutoDetect: false, Description: "Faster Qwen3 TTS (preference-only)"}, {Name: "supertonic", Modality: "tts", AutoDetect: false, Description: "Supertonic multilingual ONNX TTS (preference-only)"}, + // audio-cpp spans far more than TTS: ASR, forced alignment, VAD, speaker + // diarization, source separation, voice conversion and music generation. + // KnownBackend.Modality is a single string and the import form only chips + // on a fixed key set (core/http/react-ui/src/components/ModalityChips.jsx), + // so the extra modalities live in the description rather than in an + // invented modality key the UI would bucket as "other". + // No importer: the only reliable signal, the audiocpp.model_spec.family + // GGUF metadata key, is not readable from a remote HuggingFace repo, and + // audio-cpp/audio.cpp-gguf hosts 30-odd families in one repository. + {Name: "audio-cpp", Modality: "tts", AutoDetect: false, Description: "audio.cpp multi-family audio engine: TTS, voice cloning, ASR, alignment, VAD, diarization, separation, music generation (preference-only)"}, // Detection {Name: "sam3-cpp", Modality: "detection", AutoDetect: false, Description: "SAM3 C++ object detection (preference-only)"}, // Audio transform (audio-in / audio-out, optional reference signal) diff --git a/core/http/endpoints/localai/backend_test.go b/core/http/endpoints/localai/backend_test.go index dc6d39bea..04a486354 100644 --- a/core/http/endpoints/localai/backend_test.go +++ b/core/http/endpoints/localai/backend_test.go @@ -152,6 +152,7 @@ var _ = Describe("Backend Endpoints", func() { expectPrefOnly("tinygrad", "text") expectPrefOnly("trl", "text") expectPrefOnly("mlx-vlm", "text") + expectPrefOnly("mlx-audio", "tts") expectPrefOnly("whisperx", "asr") expectPrefOnly("crispasr", "asr") expectPrefOnly("kokoros", "tts") @@ -268,6 +269,34 @@ var _ = Describe("Backend Endpoints", func() { Expect(entry.AutoDetect).To(BeFalse()) }) + It("advertises audio-cpp as a preference-only backend naming its other modalities", func() { + req := httptest.NewRequest(http.MethodGet, "/backends/known", nil) + rec := httptest.NewRecorder() + app.ServeHTTP(rec, req) + + var payload []schema.KnownBackend + Expect(json.Unmarshal(rec.Body.Bytes(), &payload)).To(Succeed()) + + byName := map[string]schema.KnownBackend{} + for _, b := range payload { + byName[b.Name] = b + } + + entry, ok := byName["audio-cpp"] + Expect(ok).To(BeTrue(), "audio-cpp must appear in the import form dropdown") + // AutoDetect=false is the honest answer, not an omission: the only + // reliable signal is the audiocpp.model_spec.family key inside the + // GGUF, which no remote-repo probe can read. + Expect(entry.AutoDetect).To(BeFalse(), + "audio-cpp has no remote-detectable signal: the family lives inside the GGUF") + // Modality is a single string and the import form chips on a fixed + // key set, so the modalities it cannot carry have to be named in + // the description instead. + Expect(entry.Modality).To(Equal("tts")) + Expect(entry.Description).To(ContainSubstring("ASR"), + "the description must name the modalities the single Modality field cannot") + }) + It("is sorted by Modality then Name", func() { req := httptest.NewRequest(http.MethodGet, "/backends/known", nil) rec := httptest.NewRecorder() diff --git a/core/http/endpoints/localai/detokenize.go b/core/http/endpoints/localai/detokenize.go new file mode 100644 index 000000000..a3ff47963 --- /dev/null +++ b/core/http/endpoints/localai/detokenize.go @@ -0,0 +1,36 @@ +package localai + +import ( + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/pkg/model" +) + +// DetokenizeEndpoint exposes a REST API to convert token IDs back to text. +// @Summary Detokenize the input. +// @Tags tokenize +// @Param request body schema.DetokenizeRequest true "Request" +// @Success 200 {object} schema.DetokenizeResponse "Response" +// @Router /v1/detokenize [post] +func DetokenizeEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { + return func(c echo.Context) error { + input, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.DetokenizeRequest) + if !ok || input.Model == "" { + return echo.ErrBadRequest + } + + cfg, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig) + if !ok || cfg == nil { + return echo.ErrBadRequest + } + + resp, err := backend.ModelDetokenize(input.Tokens, ml, *cfg, appConfig) + if err != nil { + return err + } + return c.JSON(200, resp) + } +} diff --git a/core/http/endpoints/localai/model3d.go b/core/http/endpoints/localai/model3d.go new file mode 100644 index 000000000..fd514c7db --- /dev/null +++ b/core/http/endpoints/localai/model3d.go @@ -0,0 +1,186 @@ +package localai + +import ( + "encoding/base64" + "encoding/json" + "fmt" + "net/http" + "net/url" + "os" + "path/filepath" + "slices" + "time" + + "github.com/google/uuid" + "github.com/labstack/echo/v4" + + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/schema" + + "github.com/mudler/xlog" + + model "github.com/mudler/LocalAI/pkg/model" +) + +// Conditioning images are single frames, so a much tighter cap than the +// video-input limit is enough. +const max3DInputBytes = 32 << 20 + +var ( + valid3DQualities = []string{"", "auto", "coarse", "512", "1024"} + valid3DBackgrounds = []string{"", "auto", "keep", "black", "white"} +) + +// Model3DEndpoint +// @Summary Creates a 3D asset (binary glTF / GLB) from a conditioning image. +// @Tags 3d +// @Param request body schema.Model3DRequest true "query params" +// @Success 200 {object} schema.OpenAIResponse "Response" +// @Router /3d/generations [post] +func Model3DEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { + return func(c echo.Context) error { + input, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.Model3DRequest) + if !ok || input.Model == "" { + xlog.Error("3D Endpoint - Invalid Input") + return echo.ErrBadRequest + } + + config, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig) + if !ok || config == nil { + xlog.Error("3D Endpoint - Invalid Config") + return echo.ErrBadRequest + } + + if input.Image == "" { + return echo.NewHTTPError(http.StatusBadRequest, "image is required: 3D generation is image-conditioned") + } + // Reject unknown enum values here rather than surfacing an opaque + // backend error after a model load. + if !slices.Contains(valid3DQualities, input.Quality) { + return echo.NewHTTPError(http.StatusBadRequest, fmt.Sprintf("invalid quality %q: must be one of auto, coarse, 512, 1024", input.Quality)) + } + if !slices.Contains(valid3DBackgrounds, input.Background) { + return echo.NewHTTPError(http.StatusBadRequest, fmt.Sprintf("invalid background %q: must be one of auto, keep, black, white", input.Background)) + } + + src, err := stageVideoMediaWithLimit(c.Request().Context(), appConfig.GeneratedContentDir, input.Image, max3DInputBytes) + if err != nil { + return echo.NewHTTPError(http.StatusBadRequest, fmt.Sprintf("invalid image: %v", err)) + } + defer func() { _ = os.Remove(src) }() + + xlog.Debug("Parameter Config", "config", config) + + if config.Backend == "" { + config.Backend = model.Trellis2CppBackend + } + + step := input.Step + if step == 0 && config.Step != 0 { + step = int32(config.Step) + } + cfgScale := input.CFGScale + if cfgScale == 0 && config.CFGScale != 0 { + cfgScale = config.CFGScale + } + + b64JSON := input.ResponseFormat == "b64_json" + + tempDir := "" + if !b64JSON { + tempDir = filepath.Join(appConfig.GeneratedContentDir, "3d") + if err := os.MkdirAll(tempDir, 0o750); err != nil { + return err + } + } + // Create a temporary file + outputFile, err := os.CreateTemp(tempDir, "b64") + if err != nil { + return err + } + if err := outputFile.Close(); err != nil { + _ = os.Remove(outputFile.Name()) + return err + } + + output := outputFile.Name() + ".glb" + + // Rename the temporary file + err = os.Rename(outputFile.Name(), output) + if err != nil { + _ = os.Remove(outputFile.Name()) + return err + } + preserveOutput := false + defer func() { + if !preserveOutput { + _ = os.Remove(output) + } + }() + + baseURL := middleware.BaseURL(c) + + xlog.Debug("Model3DEndpoint: Calling Model3DGeneration", + "quality", input.Quality, + "background", input.Background, + "cfg_scale", cfgScale, + "step", step, + "texture_steps", input.TextureSteps, + "seed", input.Seed) + + fn, err := backend.Model3DGeneration( + backend.Model3DGenerationOptions{ + Image: src, + Destination: output, + Seed: input.Seed, + Step: step, + CFGScale: cfgScale, + TextureSteps: input.TextureSteps, + Quality: input.Quality, + Background: input.Background, + Params: input.Params, + }, + ml, + *config, + appConfig, + ) + if err != nil { + return mapBackendError(err) + } + if err := fn(); err != nil { + return mapBackendError(err) + } + + item := &schema.Item{} + + if b64JSON { + data, err := os.ReadFile(output) + if err != nil { + return err + } + item.B64JSON = base64.StdEncoding.EncodeToString(data) + } else { + base := filepath.Base(output) + item.URL, err = url.JoinPath(baseURL, "generated-3d", base) + if err != nil { + return err + } + preserveOutput = true + } + + id := uuid.New().String() + created := int(time.Now().Unix()) + resp := &schema.OpenAIResponse{ + ID: id, + Created: created, + Data: []schema.Item{*item}, + } + + jsonResult, _ := json.Marshal(resp) + xlog.Debug("Response", "response", string(jsonResult)) + + return c.JSON(200, resp) + } +} diff --git a/core/http/endpoints/localai/model3d_internal_test.go b/core/http/endpoints/localai/model3d_internal_test.go new file mode 100644 index 000000000..4437a2a3a --- /dev/null +++ b/core/http/endpoints/localai/model3d_internal_test.go @@ -0,0 +1,72 @@ +package localai + +import ( + "net/http" + "net/http/httptest" + "strings" + + "github.com/labstack/echo/v4" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/schema" +) + +// The validation branches all return before any model load, so the handler can +// be driven with nil loaders. +var _ = Describe("3D endpoint request validation", func() { + call := func(input *schema.Model3DRequest) error { + appConfig := &config.ApplicationConfig{GeneratedContentDir: GinkgoT().TempDir()} + handler := Model3DEndpoint(nil, nil, appConfig) + + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/3d/generations", strings.NewReader("{}")) + c := e.NewContext(req, httptest.NewRecorder()) + c.Set(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST, input) + c.Set(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG, &config.ModelConfig{Name: "test-3d"}) + return handler(c) + } + + expectBadRequest := func(err error, substr string) { + var httpErr *echo.HTTPError + Expect(err).To(BeAssignableToTypeOf(httpErr)) + httpErr = err.(*echo.HTTPError) + Expect(httpErr.Code).To(Equal(http.StatusBadRequest)) + Expect(httpErr.Message).To(ContainSubstring(substr)) + } + + It("requires a conditioning image", func() { + err := call(&schema.Model3DRequest{BasicModelRequest: schema.BasicModelRequest{Model: "m"}}) + expectBadRequest(err, "image is required") + }) + + It("rejects unknown quality values", func() { + err := call(&schema.Model3DRequest{ + BasicModelRequest: schema.BasicModelRequest{Model: "m"}, + Image: "aGk=", + Quality: "2048", + }) + expectBadRequest(err, "invalid quality") + }) + + It("rejects unknown background values", func() { + err := call(&schema.Model3DRequest{ + BasicModelRequest: schema.BasicModelRequest{Model: "m"}, + Image: "aGk=", + Background: "transparent", + }) + expectBadRequest(err, "invalid background") + }) + + It("rejects undecodable image payloads", func() { + err := call(&schema.Model3DRequest{ + BasicModelRequest: schema.BasicModelRequest{Model: "m"}, + Image: "not%%%base64", + Quality: "512", + Background: "auto", + }) + expectBadRequest(err, "invalid image") + }) +}) diff --git a/core/http/endpoints/localai/model3d_remesh.go b/core/http/endpoints/localai/model3d_remesh.go new file mode 100644 index 000000000..a0a2b004e --- /dev/null +++ b/core/http/endpoints/localai/model3d_remesh.go @@ -0,0 +1,163 @@ +package localai + +import ( + "fmt" + "io" + "math" + "net/http" + "os" + "path/filepath" + "strconv" + + "github.com/labstack/echo/v4" + + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/schema" + model "github.com/mudler/LocalAI/pkg/model" +) + +const ( + max3DRemeshBytes = 512 << 20 + defaultRemeshDetailPct = float32(0.5) + minRemeshDetailPct = float32(0.35) + maxRemeshDetailPct = float32(2.5) +) + +func normalizedRemeshDetail(detail float32) (float32, error) { + if detail == 0 { + return defaultRemeshDetailPct, nil + } + if math.IsNaN(float64(detail)) || math.IsInf(float64(detail), 0) || detail < minRemeshDetailPct || detail > maxRemeshDetailPct { + return 0, fmt.Errorf("detail must be between %.2f and %.2f percent", minRemeshDetailPct, maxRemeshDetailPct) + } + return detail, nil +} + +func saveRemeshUpload(c echo.Context, dir string) (string, error) { + header, err := c.FormFile("mesh") + if err != nil { + return "", fmt.Errorf("mesh is required") + } + if header.Size > max3DRemeshBytes { + return "", fmt.Errorf("mesh exceeds the 512 MiB limit") + } + source, err := header.Open() + if err != nil { + return "", fmt.Errorf("opening mesh: %w", err) + } + defer func() { _ = source.Close() }() + + if err := os.MkdirAll(dir, 0o750); err != nil { + return "", err + } + temp, err := os.CreateTemp(dir, "remesh-input-*.glb") + if err != nil { + return "", err + } + path := temp.Name() + defer func() { + if err != nil { + _ = os.Remove(path) + } + }() + + written, copyErr := io.Copy(temp, io.LimitReader(source, max3DRemeshBytes+1)) + closeErr := temp.Close() + if copyErr != nil { + err = fmt.Errorf("saving mesh: %w", copyErr) + return "", err + } + if closeErr != nil { + err = closeErr + return "", err + } + if written == 0 { + err = fmt.Errorf("mesh is empty") + return "", err + } + if written > max3DRemeshBytes { + err = fmt.Errorf("mesh exceeds the 512 MiB limit") + return "", err + } + return path, nil +} + +// Model3DRemeshEndpoint rebuilds an existing generated GLB as a watertight mesh. +// @Summary Applies watertight print remeshing to an existing 3D asset. +// @Tags 3d +// @Accept multipart/form-data +// @Produce model/gltf-binary +// @Param model formData string true "3D model name" +// @Param mesh formData file true "Source GLB" +// @Param detail formData number false "Detail size as percent of the source bounding-box diagonal (0.35–2.5; default 0.5)" +// @Success 200 {file} binary "Remeshed GLB" +// @Router /3d/remesh [post] +func Model3DRemeshEndpoint(ml *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { + return func(c echo.Context) error { + input, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.Model3DRemeshRequest) + if !ok || input.Model == "" { + return echo.NewHTTPError(http.StatusBadRequest, "model is required") + } + modelConfig, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig) + if !ok || modelConfig == nil { + return echo.ErrBadRequest + } + detail, err := normalizedRemeshDetail(input.Detail) + if err != nil { + return echo.NewHTTPError(http.StatusBadRequest, err.Error()) + } + + source, err := saveRemeshUpload(c, appConfig.GeneratedContentDir) + if err != nil { + return echo.NewHTTPError(http.StatusBadRequest, err.Error()) + } + defer func() { _ = os.Remove(source) }() + + outputFile, err := os.CreateTemp(appConfig.GeneratedContentDir, "remeshed-*.glb") + if err != nil { + return err + } + output := outputFile.Name() + if err := outputFile.Close(); err != nil { + _ = os.Remove(output) + return err + } + defer func() { _ = os.Remove(output) }() + + if modelConfig.Backend == "" { + modelConfig.Backend = model.Trellis2CppBackend + } + fn, err := backend.Model3DGeneration( + backend.Model3DGenerationOptions{ + Image: source, + Destination: output, + Params: map[string]string{ + "operation": "print_remesh", + "alpha_ratio": strconv.FormatFloat(float64(detail/100), 'g', -1, 32), + "detail_percent": strconv.FormatFloat(float64(detail), 'g', -1, 32), + "texture_size": "2048", + }, + }, + ml, + *modelConfig, + appConfig, + ) + if err != nil { + return mapBackendError(err) + } + if err := fn(); err != nil { + return mapBackendError(err) + } + + file, err := os.Open(filepath.Clean(output)) + if err != nil { + return err + } + defer func() { _ = file.Close() }() + c.Response().Header().Set(echo.HeaderContentDisposition, `attachment; filename="remeshed.glb"`) + c.Response().Header().Set(echo.HeaderCacheControl, "no-store") + return c.Stream(http.StatusOK, "model/gltf-binary", file) + } +} diff --git a/core/http/endpoints/localai/model3d_remesh_internal_test.go b/core/http/endpoints/localai/model3d_remesh_internal_test.go new file mode 100644 index 000000000..d3e4108e2 --- /dev/null +++ b/core/http/endpoints/localai/model3d_remesh_internal_test.go @@ -0,0 +1,68 @@ +package localai + +import ( + "bytes" + "math" + "mime/multipart" + "net/http" + "net/http/httptest" + "os" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/schema" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("3D print remeshing request handling", func() { + DescribeTable("normalizes the demo detail range", + func(input, expected float32, valid bool) { + value, err := normalizedRemeshDetail(input) + if valid { + Expect(err).NotTo(HaveOccurred()) + Expect(value).To(BeNumerically("~", expected, 1e-6)) + } else { + Expect(err).To(HaveOccurred()) + } + }, + Entry("default", float32(0), float32(0.5), true), + Entry("fine endpoint", float32(0.35), float32(0.35), true), + Entry("coarse endpoint", float32(2.5), float32(2.5), true), + Entry("too fine", float32(0.1), float32(0), false), + Entry("too coarse", float32(3), float32(0), false), + Entry("not a number", float32(math.NaN()), float32(0), false), + ) + + It("streams the multipart GLB to a bounded temporary file", func() { + body := new(bytes.Buffer) + writer := multipart.NewWriter(body) + Expect(writer.WriteField("model", "trellis-test-model")).To(Succeed()) + Expect(writer.WriteField("detail", "0.35")).To(Succeed()) + part, err := writer.CreateFormFile("mesh", "source.glb") + Expect(err).NotTo(HaveOccurred()) + _, err = part.Write([]byte("glTF-test")) + Expect(err).NotTo(HaveOccurred()) + Expect(writer.Close()).To(Succeed()) + + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/3d/remesh", body) + req.Header.Set(echo.HeaderContentType, writer.FormDataContentType()) + ctx := e.NewContext(req, httptest.NewRecorder()) + input := new(schema.Model3DRemeshRequest) + Expect(ctx.Bind(input)).To(Succeed()) + Expect(input.Model).To(Equal("trellis-test-model")) + Expect(input.Detail).To(BeNumerically("~", 0.35, 1e-6)) + path, err := saveRemeshUpload(ctx, GinkgoT().TempDir()) + Expect(err).NotTo(HaveOccurred()) + defer func() { _ = os.Remove(path) }() + Expect(os.ReadFile(path)).To(Equal([]byte("glTF-test"))) + }) + + It("requires a mesh part", func() { + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/3d/remesh", bytes.NewReader(nil)) + ctx := e.NewContext(req, httptest.NewRecorder()) + _, err := saveRemeshUpload(ctx, GinkgoT().TempDir()) + Expect(err).To(MatchError("mesh is required")) + }) +}) diff --git a/core/http/endpoints/localai/router_decide_test.go b/core/http/endpoints/localai/router_decide_test.go index d63cd091d..1efd11dcf 100644 --- a/core/http/endpoints/localai/router_decide_test.go +++ b/core/http/endpoints/localai/router_decide_test.go @@ -136,7 +136,7 @@ type stubScorer struct { labelToLogProb map[string]float64 } -func (s *stubScorer) Score(_ context.Context, _ string, candidates []string) ([]backend.CandidateScore, error) { +func (s *stubScorer) Score(_ context.Context, _ string, _ int, candidates []string) ([]backend.CandidateScore, error) { out := make([]backend.CandidateScore, len(candidates)) for i, c := range candidates { // Candidate is the Arch-Router JSON envelope diff --git a/core/http/endpoints/localai/stores.go b/core/http/endpoints/localai/stores.go index 8074da9e0..af9cf1569 100644 --- a/core/http/endpoints/localai/stores.go +++ b/core/http/endpoints/localai/stores.go @@ -9,7 +9,7 @@ import ( "github.com/mudler/LocalAI/pkg/store" ) -func StoresSetEndpoint(sl *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { +func StoresSetEndpoint(sl *model.ModelLoader, cl *config.ModelConfigLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { return func(c echo.Context) error { input := new(schema.StoresSet) @@ -17,7 +17,7 @@ func StoresSetEndpoint(sl *model.ModelLoader, appConfig *config.ApplicationConfi return err } - sb, err := backend.StoreBackend(sl, appConfig, input.Store, input.Backend) + sb, err := backend.StoreBackend(sl, appConfig, cl, input.Store, input.Backend) if err != nil { return err } @@ -36,7 +36,7 @@ func StoresSetEndpoint(sl *model.ModelLoader, appConfig *config.ApplicationConfi } } -func StoresDeleteEndpoint(sl *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { +func StoresDeleteEndpoint(sl *model.ModelLoader, cl *config.ModelConfigLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { return func(c echo.Context) error { input := new(schema.StoresDelete) @@ -44,7 +44,7 @@ func StoresDeleteEndpoint(sl *model.ModelLoader, appConfig *config.ApplicationCo return err } - sb, err := backend.StoreBackend(sl, appConfig, input.Store, input.Backend) + sb, err := backend.StoreBackend(sl, appConfig, cl, input.Store, input.Backend) if err != nil { return err } @@ -57,7 +57,7 @@ func StoresDeleteEndpoint(sl *model.ModelLoader, appConfig *config.ApplicationCo } } -func StoresGetEndpoint(sl *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { +func StoresGetEndpoint(sl *model.ModelLoader, cl *config.ModelConfigLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { return func(c echo.Context) error { input := new(schema.StoresGet) @@ -65,7 +65,7 @@ func StoresGetEndpoint(sl *model.ModelLoader, appConfig *config.ApplicationConfi return err } - sb, err := backend.StoreBackend(sl, appConfig, input.Store, input.Backend) + sb, err := backend.StoreBackend(sl, appConfig, cl, input.Store, input.Backend) if err != nil { return err } @@ -88,7 +88,7 @@ func StoresGetEndpoint(sl *model.ModelLoader, appConfig *config.ApplicationConfi } } -func StoresFindEndpoint(sl *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { +func StoresFindEndpoint(sl *model.ModelLoader, cl *config.ModelConfigLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { return func(c echo.Context) error { input := new(schema.StoresFind) @@ -96,7 +96,7 @@ func StoresFindEndpoint(sl *model.ModelLoader, appConfig *config.ApplicationConf return err } - sb, err := backend.StoreBackend(sl, appConfig, input.Store, input.Backend) + sb, err := backend.StoreBackend(sl, appConfig, cl, input.Store, input.Backend) if err != nil { return err } diff --git a/core/http/endpoints/localai/system.go b/core/http/endpoints/localai/system.go index 3e5652117..6a0b13965 100644 --- a/core/http/endpoints/localai/system.go +++ b/core/http/endpoints/localai/system.go @@ -12,7 +12,7 @@ import ( // @Tags monitoring // @Success 200 {object} schema.SystemInformationResponse "Response" // @Router /system [get] -func SystemInformations(ml *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { +func SystemInformations(cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { return func(c echo.Context) error { availableBackends := []string{} loadedModels := ml.ListLoadedModels() @@ -25,7 +25,14 @@ func SystemInformations(ml *model.ModelLoader, appConfig *config.ApplicationConf sysmodels := []schema.SysInfoModel{} for _, m := range loadedModels { - sysmodels = append(sysmodels, schema.SysInfoModel{ID: m.ID}) + entry := schema.SysInfoModel{ID: m.ID} + // The loader tracks only the ID. Which engine is serving a model is + // the first thing an operator wants beside its name, and it is one + // config lookup away. + if cfg, ok := cl.GetModelConfig(m.ID); ok { + entry.Backend = cfg.Backend + } + sysmodels = append(sysmodels, entry) } return c.JSON(200, schema.SystemInformationResponse{ diff --git a/core/http/endpoints/localai/traces.go b/core/http/endpoints/localai/traces.go index 7e9f1216f..2ae8f5b66 100644 --- a/core/http/endpoints/localai/traces.go +++ b/core/http/endpoints/localai/traces.go @@ -3,6 +3,7 @@ package localai import ( "net/http" "strconv" + "time" "github.com/labstack/echo/v4" "github.com/mudler/LocalAI/core/http/middleware" @@ -85,6 +86,35 @@ func GetAPITracesEndpoint() echo.HandlerFunc { } } +// GetAPITracesSummaryEndpoint returns counted totals over a recent window +// @Summary Summarize recent API traces +// @Description Returns request, failure and latency totals over a recent window, plus a bucketed series for sparklines. Exists so callers wanting three numbers do not have to fetch the whole trace list and count it themselves. +// @Tags monitoring +// @Produce json +// @Param hours query int false "Window in hours (default 24, max 168)" +// @Success 200 {object} middleware.TraceSummary "Counted trace totals" +// @Router /api/traces/summary [get] +func GetAPITracesSummaryEndpoint() echo.HandlerFunc { + return func(c echo.Context) error { + hours := 24 + if raw := c.QueryParam("hours"); raw != "" { + if v, err := strconv.Atoi(raw); err == nil && v > 0 { + hours = v + } + } + // A week is plenty for a dashboard, and the trace buffer is bounded + // anyway; an unbounded window would just scan the whole buffer. + if hours > 168 { + hours = 168 + } + return c.JSON(http.StatusOK, middleware.GetTracesSummary(time.Duration(hours)*time.Hour, traceSummaryBuckets)) + } +} + +// Enough columns for a sparkline to show a shape, few enough that each one +// still holds a meaningful count on a quiet installation. +const traceSummaryBuckets = 12 + // GetAPITraceEndpoint returns a single API trace with its full payload // @Summary Get one API trace // @Description Returns a single captured API exchange, including the request and response bodies omitted from the list response diff --git a/core/http/endpoints/mcp/localai_assistant_test.go b/core/http/endpoints/mcp/localai_assistant_test.go index 2231350d0..213f6fa0c 100644 --- a/core/http/endpoints/mcp/localai_assistant_test.go +++ b/core/http/endpoints/mcp/localai_assistant_test.go @@ -84,6 +84,22 @@ func (stubClient) ListNodes(_ context.Context) ([]localaitools.Node, error) { return []localaitools.Node{}, nil } +func (stubClient) ListScheduling(_ context.Context) ([]localaitools.ModelSchedulingConfig, error) { + return []localaitools.ModelSchedulingConfig{}, nil +} + +func (stubClient) GetScheduling(_ context.Context, _ string) (*localaitools.ModelSchedulingConfig, error) { + return &localaitools.ModelSchedulingConfig{}, nil +} + +func (stubClient) SetScheduling(_ context.Context, _ localaitools.SetSchedulingRequest) (*localaitools.ModelSchedulingConfig, error) { + return &localaitools.ModelSchedulingConfig{}, nil +} + +func (stubClient) DeleteScheduling(_ context.Context, _ string) error { + return nil +} + func (stubClient) SetNodeVRAMBudget(_ context.Context, _, _ string) error { return nil } diff --git a/core/http/endpoints/openai/moderations.go b/core/http/endpoints/openai/moderations.go new file mode 100644 index 000000000..a42403686 --- /dev/null +++ b/core/http/endpoints/openai/moderations.go @@ -0,0 +1,190 @@ +package openai + +import ( + "context" + "encoding/json" + "fmt" + "math" + "net/http" + "strings" + + "github.com/google/uuid" + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/templates" + "github.com/mudler/LocalAI/pkg/functions" + "github.com/mudler/LocalAI/pkg/model" +) + +var moderationCategories = []string{ + "harassment", + "harassment/threatening", + "hate", + "hate/threatening", + "illicit", + "illicit/violent", + "self-harm", + "self-harm/intent", + "self-harm/instructions", + "sexual", + "sexual/minors", + "violence", + "violence/graphic", +} + +type moderationGenerator func(context.Context, string, *config.ModelConfig) (string, backend.TokenUsage, error) + +type generatedModeration struct { + Categories map[string]bool `json:"categories"` + CategoryScores map[string]float64 `json:"category_scores"` +} + +// ModerationEndpoint implements the text input subset of OpenAI's moderation +// API using any LocalAI completion model and constrained JSON generation. +// @Summary Classify text for potentially harmful content. +// @Tags moderation +// @Param request body schema.ModerationRequest true "query params" +// @Success 200 {object} schema.ModerationResponse "Response" +// @Router /v1/moderations [post] +func ModerationEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, appConfig *config.ApplicationConfig) echo.HandlerFunc { + return moderationEndpoint(func(ctx context.Context, input string, cfg *config.ModelConfig) (string, backend.TokenUsage, error) { + prompt := moderationPrompt(input) + var messages schema.Messages + if cfg.TemplateConfig.UseTokenizerTemplate { + messages = schema.Messages{{Role: "user", Content: prompt}} + prompt = "" + } else if evaluator != nil { + if rendered, err := evaluator.EvaluateTemplateForPrompt(templates.CompletionPromptTemplate, *cfg, templates.PromptTemplateData{Input: prompt, SystemPrompt: cfg.SystemPrompt}); err == nil { + prompt = rendered + } + } + + predict, err := backend.ModelInferenceFunc(ctx, prompt, messages, nil, nil, nil, ml, cfg, cl, appConfig, nil, "", "", nil, nil, nil, nil) + if err != nil { + return "", backend.TokenUsage{}, err + } + response, err := predict() + return response.Response, response.Usage, err + }) +} + +func moderationEndpoint(generate moderationGenerator) echo.HandlerFunc { + return func(c echo.Context) error { + input, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.ModerationRequest) + if !ok || input == nil { + return echo.NewHTTPError(http.StatusBadRequest, "invalid moderation request") + } + if len(input.Input) == 0 { + return echo.NewHTTPError(http.StatusBadRequest, "input must contain at least one text string") + } + if generate == nil { + return echo.NewHTTPError(http.StatusInternalServerError, "moderation generator is unavailable") + } + + modelConfig, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig) + if !ok || modelConfig == nil { + return echo.NewHTTPError(http.StatusBadRequest, "moderation model configuration is unavailable") + } + + grammar, err := moderationGrammar() + if err != nil { + return echo.NewHTTPError(http.StatusInternalServerError, "failed to build moderation grammar").SetInternal(err) + } + cfg := *modelConfig + cfg.Grammar = grammar + maxTokens := 512 + cfg.Maxtokens = &maxTokens + + results := make([]schema.ModerationResult, 0, len(input.Input)) + promptTokens, completionTokens := 0, 0 + for _, text := range input.Input { + if strings.TrimSpace(text) == "" { + return echo.NewHTTPError(http.StatusBadRequest, "input strings must not be empty") + } + raw, usage, err := generate(c.Request().Context(), text, &cfg) + if err != nil { + return echo.NewHTTPError(http.StatusInternalServerError, "moderation inference failed").SetInternal(err) + } + result, err := parseModerationResult(raw) + if err != nil { + return echo.NewHTTPError(http.StatusInternalServerError, "invalid moderation result: "+err.Error()).SetInternal(err) + } + results = append(results, result) + promptTokens += usage.Prompt + completionTokens += usage.Completion + } + + middleware.StampUsage(c, input.Model, promptTokens, completionTokens) + return c.JSON(http.StatusOK, schema.ModerationResponse{ + ID: "modr-" + uuid.NewString(), + Model: input.Model, + Results: results, + }) + } +} + +func moderationPrompt(input string) string { + encoded, _ := json.Marshal(input) + return `Classify the supplied text for safety moderation. Return only the JSON object required by the grammar. +For each category, set the boolean to true only when the text matches that category. Give each category a confidence score from 0 to 1. +The categories are harassment, harassment/threatening, hate, hate/threatening, illicit, illicit/violent, self-harm, self-harm/intent, self-harm/instructions, sexual, sexual/minors, violence, and violence/graphic. +Text to classify: ` + string(encoded) +} + +func moderationGrammar() (string, error) { + boolProperties := map[string]any{} + scoreProperties := map[string]any{} + for _, category := range moderationCategories { + boolProperties[category] = map[string]any{"type": "boolean"} + scoreProperties[category] = map[string]any{"type": "number"} + } + structure := functions.JSONFunctionStructure{AnyOf: []functions.Item{{ + Type: "object", + Properties: map[string]any{ + "categories": map[string]any{ + "type": "object", + "properties": boolProperties, + "required": moderationCategories, + "additionalProperties": false, + }, + "category_scores": map[string]any{ + "type": "object", + "properties": scoreProperties, + "required": moderationCategories, + "additionalProperties": false, + }, + }, + }}} + return structure.Grammar() +} + +func parseModerationResult(raw string) (schema.ModerationResult, error) { + var generated generatedModeration + if err := json.Unmarshal([]byte(strings.TrimSpace(raw)), &generated); err != nil { + return schema.ModerationResult{}, err + } + + result := schema.ModerationResult{ + Categories: make(map[string]bool, len(moderationCategories)), + CategoryScores: make(map[string]float64, len(moderationCategories)), + CategoryAppliedInputTypes: make(map[string][]string, len(moderationCategories)), + } + for _, category := range moderationCategories { + flagged, exists := generated.Categories[category] + if !exists { + return schema.ModerationResult{}, fmt.Errorf("missing category %q", category) + } + score, exists := generated.CategoryScores[category] + if !exists || math.IsNaN(score) || math.IsInf(score, 0) || score < 0 || score > 1 { + return schema.ModerationResult{}, fmt.Errorf("category %q has an invalid score", category) + } + result.Categories[category] = flagged + result.CategoryScores[category] = score + result.CategoryAppliedInputTypes[category] = []string{"text"} + result.Flagged = result.Flagged || flagged + } + return result, nil +} diff --git a/core/http/endpoints/openai/moderations_test.go b/core/http/endpoints/openai/moderations_test.go new file mode 100644 index 000000000..b6f5130e7 --- /dev/null +++ b/core/http/endpoints/openai/moderations_test.go @@ -0,0 +1,105 @@ +package openai + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/schema" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Moderations endpoint", func() { + It("classifies each text input and returns the OpenAI response shape", func() { + inputs := []string{} + generate := func(_ context.Context, input string, cfg *config.ModelConfig) (string, backend.TokenUsage, error) { + inputs = append(inputs, input) + Expect(cfg.Grammar).To(ContainSubstring("harassment")) + return `{ + "categories":{"harassment":true,"harassment/threatening":false,"hate":false,"hate/threatening":false,"illicit":false,"illicit/violent":false,"self-harm":false,"self-harm/intent":false,"self-harm/instructions":false,"sexual":false,"sexual/minors":false,"violence":false,"violence/graphic":false}, + "category_scores":{"harassment":0.9,"harassment/threatening":0.1,"hate":0,"hate/threatening":0,"illicit":0,"illicit/violent":0,"self-harm":0,"self-harm/intent":0,"self-harm/instructions":0,"sexual":0,"sexual/minors":0,"violence":0,"violence/graphic":0} + }`, backend.TokenUsage{Prompt: 12, Completion: 8}, nil + } + + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/v1/moderations", strings.NewReader(`{"model":"guard","input":["first","second"]}`)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + ctx := e.NewContext(req, rec) + ctx.Set(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST, &schema.ModerationRequest{ + BasicModelRequest: schema.BasicModelRequest{Model: "guard"}, + Input: schema.ModerationInput{"first", "second"}, + }) + modelConfig := &config.ModelConfig{Name: "guard"} + modelConfig.Model = "guard.gguf" + ctx.Set(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG, modelConfig) + + Expect(moderationEndpoint(generate)(ctx)).To(Succeed()) + Expect(rec.Code).To(Equal(http.StatusOK)) + Expect(inputs).To(Equal([]string{"first", "second"})) + + var response schema.ModerationResponse + Expect(json.Unmarshal(rec.Body.Bytes(), &response)).To(Succeed()) + Expect(response.ID).To(HavePrefix("modr-")) + Expect(response.Model).To(Equal("guard")) + Expect(response.Results).To(HaveLen(2)) + Expect(response.Results[0].Flagged).To(BeTrue()) + Expect(response.Results[0].Categories["harassment"]).To(BeTrue()) + Expect(response.Results[0].CategoryAppliedInputTypes["harassment"]).To(Equal([]string{"text"})) + }) + + It("rejects an empty input list", func() { + e := echo.New() + ctx := e.NewContext(httptest.NewRequest(http.MethodPost, "/v1/moderations", nil), httptest.NewRecorder()) + ctx.Set(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST, &schema.ModerationRequest{ + BasicModelRequest: schema.BasicModelRequest{Model: "guard"}, + }) + ctx.Set(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG, &config.ModelConfig{Name: "guard"}) + + err := moderationEndpoint(nil)(ctx) + Expect(err).To(MatchError(ContainSubstring("input must contain at least one text string"))) + Expect(err.(*echo.HTTPError).Code).To(Equal(http.StatusBadRequest)) + }) + + It("surfaces malformed classifier output without returning a partial result", func() { + generate := func(context.Context, string, *config.ModelConfig) (string, backend.TokenUsage, error) { + return "not-json", backend.TokenUsage{}, nil + } + e := echo.New() + ctx := e.NewContext(httptest.NewRequest(http.MethodPost, "/v1/moderations", nil), httptest.NewRecorder()) + ctx.Set(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST, &schema.ModerationRequest{ + BasicModelRequest: schema.BasicModelRequest{Model: "guard"}, + Input: schema.ModerationInput{"text"}, + }) + ctx.Set(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG, &config.ModelConfig{Name: "guard"}) + + err := moderationEndpoint(generate)(ctx) + Expect(err).To(MatchError(ContainSubstring("invalid moderation result"))) + Expect(err.(*echo.HTTPError).Code).To(Equal(http.StatusInternalServerError)) + }) +}) + +var _ = Describe("Moderation input", func() { + DescribeTable("accepts OpenAI text input forms", + func(body string, expected schema.ModerationInput) { + var req schema.ModerationRequest + Expect(json.Unmarshal([]byte(body), &req)).To(Succeed()) + Expect(req.Input).To(Equal(expected)) + }, + Entry("single text", `{"input":"hello"}`, schema.ModerationInput{"hello"}), + Entry("text array", `{"input":["hello","world"]}`, schema.ModerationInput{"hello", "world"}), + ) + + It("rejects multimodal input in the text-only MVP", func() { + var req schema.ModerationRequest + err := json.Unmarshal([]byte(`{"input":[{"type":"image_url","image_url":{"url":"https://example.com/a.png"}}]}`), &req) + Expect(err).To(MatchError(ContainSubstring("text string or array of text strings"))) + }) +}) diff --git a/core/http/endpoints/openai/realtime.go b/core/http/endpoints/openai/realtime.go index b22e1a8dd..32b08b3fd 100644 --- a/core/http/endpoints/openai/realtime.go +++ b/core/http/endpoints/openai/realtime.go @@ -30,6 +30,7 @@ import ( "github.com/mudler/LocalAI/core/http/endpoints/openai/turncoord" "github.com/mudler/LocalAI/core/http/endpoints/openai/types" "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/services/routing/router" "github.com/mudler/LocalAI/core/templates" laudio "github.com/mudler/LocalAI/pkg/audio" "github.com/mudler/LocalAI/pkg/functions" @@ -150,6 +151,12 @@ type Session struct { // pairs are kept together so we never feed an orphaned tool result. MaxHistoryItems int + // Classifier holds the LocalAI classifier-mode config (prefill-scored + // option selection instead of generation), seeded from + // pipeline.classifier and replaced wholesale by session.update's + // localai_classifier field. nil means off. + Classifier *types.ClassifierConfig + // Compaction settings resolved from pipeline.compaction (see resolveCompaction). CompactionEnabled bool CompactionTrigger int @@ -210,14 +217,15 @@ func (s *Session) ToServer() types.SessionUnion { } else { return types.SessionUnion{ Realtime: &types.RealtimeSession{ - ID: s.ID, - Object: "realtime.session", - Model: s.Model, - Instructions: s.Instructions, - Tools: s.Tools, - ToolChoice: s.ToolChoice, - MaxOutputTokens: s.MaxOutputTokens, - OutputModalities: s.OutputModalities, + ID: s.ID, + Object: "realtime.session", + Model: s.Model, + Instructions: s.Instructions, + Tools: s.Tools, + ToolChoice: s.ToolChoice, + MaxOutputTokens: s.MaxOutputTokens, + OutputModalities: s.OutputModalities, + LocalAIClassifier: s.Classifier, Audio: &types.RealtimeSessionAudio{ Input: &types.SessionAudioInput{ TurnDetection: s.TurnDetection, @@ -279,6 +287,24 @@ type Model interface { // event. Backends without live support fail with an error satisfying // grpcerrors.IsLiveTranscriptionUnsupported. TranscribeLive(ctx context.Context, language string, onEvent func(backend.LiveTranscriptionEvent)) (backend.LiveTranscriptionSession, error) + // ClassifyTurn prefill-scores each classifier option as a candidate + // continuation of the conversation (LocalAI classifier-mode extension) + // and returns the softmax distribution in option order. Runs on the + // pipeline's scoring model (classifier.model, defaulting to the LLM) — + // no autoregressive decode happens. + ClassifyTurn(ctx context.Context, messages schema.Messages, options []types.ClassifierOption, normalization string) ([]router.LabelScore, error) + // PrewarmClassifier primes the scoring backend's prompt cache for a + // newly registered option list (fired async on registration) so the + // first turns after a session.update don't pay the option-list + // prefill. Best-effort and idempotent per option set. + PrewarmClassifier(ctx context.Context, options []types.ClassifierOption, normalization string) + // FillToolArguments completes the chosen option's argument slots with a + // short grammar-constrained completion that continues the exact scoring + // prompt (so the backend's prompt cache stays warm) and returns the + // spliced tool-arguments JSON plus the raw slot values (for reply + // templating) — the hybrid between prefill-only classification and full + // generation. + FillToolArguments(ctx context.Context, messages schema.Messages, options []types.ClassifierOption, normalization string, chosen *types.ClassifierOption) (string, map[string]string, error) PredictConfig() *config.ModelConfig // Warmup eagerly loads the pipeline's sub-model backends into memory so the // first realtime turn doesn't pay each backend's cold-start load cost. Loads @@ -553,6 +579,12 @@ func runRealtimeSession(application *application.Application, t Transport, model SoundDetectionHopMs: cfg.Pipeline.SoundDetectionHopMs, } session.CompactionEnabled, session.CompactionTrigger, session.MaxSummaryTokens, session.SummaryModel = resolveCompaction(cfg, session.MaxHistoryItems) + classifier, err := classifierConfigFromPipeline(cfg.Pipeline.Classifier) + if err != nil { + sendError(t, "invalid_pipeline", "pipeline classifier: "+err.Error(), "", "") + return + } + session.Classifier = classifier // Single-writer response coordinator (machine M3). All response starts and // cancels go through this, so the read-loop and VAD goroutine can never race @@ -602,6 +634,10 @@ func runRealtimeSession(application *application.Application, t Transport, model return } session.ModelInterface = m + // A pipeline-seeded option list gets its scoring prompt prewarmed + // alongside the model warm-up below, so the session's first turn + // doesn't pay the option-list prefill. + prewarmClassifier(session) // The voice gate is built before the warm-up below so its // speaker-recognition model can warm alongside the pipeline stages. @@ -764,7 +800,9 @@ func runRealtimeSession(application *application.Application, t Transport, model application.ApplicationConfig(), ); err != nil { xlog.Error("failed to update session", "error", err) - sendError(t, "session_update_error", "Failed to update session", "", "") + // The cause is validation feedback on the client's own + // payload — echo it so UIs can show something actionable. + sendError(t, "session_update_error", fmt.Sprintf("Failed to update session: %v", err), "", "") continue } @@ -790,7 +828,7 @@ func runRealtimeSession(application *application.Application, t Transport, model buildRealtimeRoutingContext(application, session.ID), ); err != nil { xlog.Error("failed to update session", "error", err) - sendError(t, "session_update_error", "Failed to update session", "", "") + sendError(t, "session_update_error", fmt.Sprintf("Failed to update session: %v", err), "", "") continue } @@ -966,6 +1004,10 @@ func runRealtimeSession(application *application.Application, t Transport, model case types.ResponseCreateEvent: xlog.Debug("recv", "message", string(msg)) + if err := validateClassifierActivation(session.ModelInterface, e.Response.LocalAIClassifier); err != nil { + sendError(t, "invalid_request_error", "Invalid response classifier: "+err.Error(), "", e.EventID) + continue + } // Handle optional items to add to context if len(e.Response.Input) > 0 { @@ -1241,6 +1283,17 @@ func updateSession(session *Session, update *types.SessionUnion, cl *config.Mode session.ToolChoice = rt.ToolChoice } + if rt.LocalAIClassifier != nil { + // Replace-not-merge, like tools: the client owns the whole option + // list. Invalid configs reject the update without touching the + // session's current classifier. + if err := validateClassifierActivation(session.ModelInterface, rt.LocalAIClassifier); err != nil { + return err + } + session.Classifier = rt.LocalAIClassifier + prewarmClassifier(session) + } + if rt.MaxOutputTokens != 0 { session.MaxOutputTokens = rt.MaxOutputTokens } @@ -1314,6 +1367,23 @@ func decodeOpusLoop(session *Session, opusBackend grpc.Backend, done chan struct // it cuts the start of the utterance the next tick will detect. const noSpeechHoldbackSec = 0.5 +// vadWarmupMarginSec pads the VAD scan window beyond the largest silence the +// commit test can need to measure. It covers silero's cold-start (the LSTM +// state converges within a few hundred ms — the model has no longer-range +// memory, which is why clipping is sound at all) plus sherpa's segment +// hysteresis (min_speech 0.25s / min_silence 0.5s), which must fit inside +// the clip for segments to open and close at all. +const vadWarmupMarginSec = 1.0 + +// maxTurnBufferSec bounds the raw input buffer. Without it a turn that never +// pauses (continuous noise or speech: silero segments every tick, so the +// no-speech clear never runs and nothing commits) grows the buffer toward the +// 100MB append cap, and with it the per-tick copy+resample and the +// commit-time WAV/batch decode. 90s keeps all of those trivial; only an +// unbroken >90s turn loses head audio from a server_vad batch transcription +// (semantic mode already consumed it incrementally via the live stream). +const maxTurnBufferSec = 90.0 + // dropInspectedPrefix removes the head of the audio buffer that a VAD tick // inspected (the first inspected bytes), keeping the newest holdbackBytes of // that window plus everything appended while the tick ran — audio the VAD @@ -1375,185 +1445,292 @@ func handleVAD(session *Session, conv *Conversation, t Transport, done chan stru case <-done: return case <-ticker.C: - // Semantic mode is re-read each tick: session.update can switch - // turn-detection modes (and the retranscribe gate) mid-session. - sessionLock.Lock() - var sv *types.RealtimeSessionSemanticVad - if session.TurnDetection != nil { - sv = session.TurnDetection.SemanticVad - } - retranscribe := sv != nil && session.ModelConfig != nil && - session.ModelConfig.Pipeline.TurnDetectionRetranscribe() - sessionLock.Unlock() - - // The turn coordinator's data-heavy effects (OpenTurn/CommitTurn) - // need this tick's mode; set it before any Apply below. - sink.sv = sv - - // session.update switched semantic -> server mid-turn: drop the - // orphaned live stream. This is NOT a turn abort — the turn continues - // under server_vad (a config change must not cut off a mid-utterance - // speaker), so the coordinator stays Speaking; only the orphaned live - // stream is closed. - if sv == nil && lts.open() { - lts.discardTurn() - } - - session.AudioBufferLock.Lock() - allAudio := make([]byte, len(session.InputAudioBuffer)) - copy(allAudio, session.InputAudioBuffer) - session.AudioBufferLock.Unlock() - - aints := sound.BytesToInt16sLE(allAudio) - if len(aints) == 0 || len(aints) < int(silenceThreshold*float64(session.InputSampleRate)) { - continue - } - - // Resample from InputSampleRate to 16kHz - aints = sound.ResampleInt16(aints, session.InputSampleRate, localSampleRate) - - audioLength := float64(len(aints)) / localSampleRate - - if sv != nil && lts.open() { - lts.feedNewAudio(aints) - lts.drainEvents(audioLength) - } - - segments, err := runVAD(vadContext, session, aints) - if err != nil { - if err.Error() == "unexpected speech end" { - xlog.Debug("VAD cancelled") - continue - } - xlog.Error("failed to process audio", "error", err) - sendError(t, "processing_error", "Failed to process audio: "+err.Error(), "", "") - continue - } - - // NOTE: the no-speech clear and the min-buffer gate above stay on - // the short silenceThreshold even in semantic mode — the eagerness - // fallback applies only to the end-of-speech commit decision, or a - // low eagerness would delay speech_started/barge-in by seconds. - if len(segments) == 0 && audioLength > silenceThreshold { - // "No segments" is not "no speech": silero (threshold 0.5) - // crosses up to a few hundred ms into a soft word onset, so - // the newest audio in the inspected window may be the start - // of a word the next tick will recognize — and more audio - // arrived while this tick ran. Keep both; drop only the - // older, confirmed-silent head, or utterance onsets get cut. - holdback := int(noSpeechHoldbackSec*float64(session.InputSampleRate)) * 2 - session.AudioBufferLock.Lock() - session.InputAudioBuffer = dropInspectedPrefix(session.InputAudioBuffer, len(allAudio), holdback) - session.AudioBufferLock.Unlock() - - // No-speech clear: end any open turn (Speaking -> Idle, discarding - // the partial). Returning to Idle is the fix for failure mode 4 — - // the legacy discardTurn left speechStarted true, suppressing the - // next onset. Idle while not speaking is a no-op. - if err := sink.coord.Apply(turncoord.Abort{Reason: turncoord.AbortNoSpeech}); err != nil { - xlog.Error("turncoord: abort(no_speech) failed", "error", err) - } - continue - } else if len(segments) == 0 { - continue - } - - // Speech detected this tick: open the turn (Idle -> Speaking) through - // the coordinator. On that transition it opens the turn's live ASR - // stream + feeds the buffered prefix (OpenTurn), cancels any in-flight - // response (BargeIn, non-blocking — the VAD tick is never stalled), and - // emits speech_started. While already Speaking it is a no-op, so "turn - // open" and "speech started" can never disagree. The turn id is minted - // here and carried by the coordinator through to the committed event. - sink.onsetAudio = aints - if err := sink.coord.Apply(turncoord.Onset{Turn: turncoord.TurnID(generateItemID())}); err != nil { - xlog.Error("turncoord: onset failed", "error", err) - } - - if sv != nil { - // Drain again: events produced by THIS tick's feed have - // usually arrived by the time runVAD returns, and leaving - // them for the next tick adds 300ms to every EOU-triggered - // commit. - lts.drainEvents(audioLength) - } - - // Segment still in progress when audio ended - segEndTime := segments[len(segments)-1].End - if segEndTime == 0 { - continue - } - - threshold := silenceThreshold - eouPending := false - if sv != nil { - eouPending = lts.eouPending(segments) - threshold = lts.thresholdSec(eouPending, sv) - } - - if float32(audioLength)-segEndTime > float32(threshold) { - if sv != nil { - trigger, eouLag := lts.commitTrigger(eouPending, float64(segEndTime)) - xlog.Info("semantic_vad: committing turn", - "trigger", trigger, - "speech_end_s", segEndTime, - "eou_lag_s", eouLag, - "silence_s", audioLength-float64(segEndTime), - "audio_s", audioLength) - } - // Retranscribe gate (semantic mode, EOU-triggered commits - // only): cross-check the streamed EOU with an offline decode - // of the buffered turn before committing. Runs synchronously - // on the tick — the engine would serialize a concurrent feed - // against it anyway. Timeout-triggered commits skip the gate. - var gated *schema.TranscriptionResult - if retranscribe && eouPending { - batch, gerr := transcribeUtterance(vadContext, sound.Int16toBytesLE(aints), session) - switch { - case gerr != nil: - xlog.Warn("semantic_vad: retranscribe gate failed; committing via the file path", "error", gerr) - case !batch.Eou: - xlog.Info("semantic_vad: batch decode did not confirm the streamed EOU; continuing to listen", - "streamed", lts.previewText(), "batch", batch.Text) - // The batch decode rejected the streamed EOU as a false - // positive: consume the recorded EOU so the next tick - // falls back to the eagerness window instead of - // re-triggering on the same token. - lts.eouAtSec = 0 - continue - default: - xlog.Info("semantic_vad: batch decode confirmed the streamed EOU", - "streamed", lts.previewText(), "batch", batch.Text) - gated = batch - } - } - - xlog.Debug("Detected end of speech segment") - session.AudioBufferLock.Lock() - // Keep audio appended while this tick ran — it belongs to - // the next turn (in any mode: nil-ing it dropped the onset - // of an utterance started right after a commit). - session.InputAudioBuffer = dropInspectedPrefix(session.InputAudioBuffer, len(allAudio), 0) - session.AudioBufferLock.Unlock() - - // Commit the turn through the coordinator: it emits speech_stopped - // (EmitSpeechStopped) then the committed event, finalizes the live - // stream, and issues the response (CommitTurn). The committed item - // id is the coordinator's turn id (== the id the live captions - // streamed under), so the client replaces the partial text. - sink.commitAudio = sound.Int16toBytesLE(aints) - sink.commitAudioLength = audioLength - sink.commitRetranscribe = retranscribe - sink.commitGated = gated - // TODO: Remove prefix silence that is over TurnDetectionParams.PrefixPaddingMs - if err := sink.coord.Apply(turncoord.Silence{}); err != nil { - xlog.Error("turncoord: commit failed", "error", err) - } - } + vadTick(sink, silenceThreshold) } } } +// vadTick runs one turn-detection inspection of the session's input buffer: +// snapshot, resample, silero scan, live-ASR drain, and the coordinator +// transitions that follow. Extracted from handleVAD so specs can drive turn +// detection synchronously without the ticker (same shape as +// classifySoundWindow). +func vadTick(sink *turnSink, silenceThreshold float64) { + session := sink.session + t := sink.transport + lts := sink.lts + vadContext := sink.vadContext + + // Semantic mode is re-read each tick: session.update can switch + // turn-detection modes (and the retranscribe gate) mid-session. + sessionLock.Lock() + var sv *types.RealtimeSessionSemanticVad + if session.TurnDetection != nil { + sv = session.TurnDetection.SemanticVad + } + retranscribe := sv != nil && session.ModelConfig != nil && + session.ModelConfig.Pipeline.TurnDetectionRetranscribe() + sessionLock.Unlock() + + // The turn coordinator's data-heavy effects (OpenTurn/CommitTurn) + // need this tick's mode; set it before any Apply below. + sink.sv = sv + + // session.update switched semantic -> server mid-turn: drop the + // orphaned live stream. This is NOT a turn abort — the turn continues + // under server_vad (a config change must not cut off a mid-utterance + // speaker), so the coordinator stays Speaking; only the orphaned live + // stream is closed. + if sv == nil && lts.open() { + lts.discardTurn() + } + + session.AudioBufferLock.Lock() + // Retention bound: drop the buffer head beyond maxTurnBufferSec so that a + // turn that never pauses can't grow memory, the per-tick copy+resample, + // or the commit-time decode without limit (this also bounds the + // runVAD-error path below, which can't trim). A mid-turn trim shifts + // every buffer-relative cursor, so the live-feed and EOU positions are + // rebased by the trimmed amount. Whole input-seconds only: second-aligned + // cuts keep the resampled tail sample-identical to the suffix of the + // previous whole-buffer resample, so the live feed stays gapless. + bytesPerSec := session.InputSampleRate * 2 + if maxBytes := int(maxTurnBufferSec) * bytesPerSec; len(session.InputAudioBuffer) > maxBytes { + trimSecs := (len(session.InputAudioBuffer) - maxBytes + bytesPerSec - 1) / bytesPerSec + session.InputAudioBuffer = append([]byte(nil), session.InputAudioBuffer[trimSecs*bytesPerSec:]...) + lts.rebase(float64(trimSecs)) + sink.lastSpeechEndSec = max(0, sink.lastSpeechEndSec-float64(trimSecs)) + } + allAudio := make([]byte, len(session.InputAudioBuffer)) + copy(allAudio, session.InputAudioBuffer) + session.AudioBufferLock.Unlock() + + aints := sound.BytesToInt16sLE(allAudio) + if len(aints) == 0 || len(aints) < int(silenceThreshold*float64(session.InputSampleRate)) { + return + } + + // Resample from InputSampleRate to 16kHz + aints = sound.ResampleInt16(aints, session.InputSampleRate, localSampleRate) + + audioLength := float64(len(aints)) / localSampleRate + + if sv != nil && lts.open() { + lts.feedNewAudio(aints) + lts.drainEvents(audioLength) + } + + // Scan window: silero's recurrent state carries only a few hundred ms of + // context, so audio older than the largest silence the commit test can + // need to measure (plus warm-up margin) contributes nothing to the + // tail's classification — clip it instead of rescanning the whole turn + // every tick (~3.3ms of silero per buffered second, quadratic over a + // turn). Segment times are rebased back to whole-buffer coordinates so + // every downstream consumer (trailing-silence math, eouPending, the + // live-feed cursor) is untouched. + scan := aints + clipOffsetSec := 0.0 + if maxScan := int(vadScanWindowSec(sv, silenceThreshold, session.ModelConfig) * localSampleRate); len(aints) > maxScan { + scan = aints[len(aints)-maxScan:] + clipOffsetSec = float64(len(aints)-maxScan) / localSampleRate + } + + segments, err := runVAD(vadContext, session, scan) + if err != nil { + if err.Error() == "unexpected speech end" { + xlog.Debug("VAD cancelled") + return + } + xlog.Error("failed to process audio", "error", err) + sendError(t, "processing_error", "Failed to process audio: "+err.Error(), "", "") + return + } + for i := range segments { + segments[i].Start += float32(clipOffsetSec) + // End == 0 is the "segment still open" sentinel — leave it alone. + if segments[i].End != 0 { + segments[i].End += float32(clipOffsetSec) + } + } + + // NOTE: the no-speech clear and the min-buffer gate above stay on + // the short silenceThreshold even in semantic mode — the eagerness + // fallback applies only to the end-of-speech commit decision, or a + // low eagerness would delay speech_started/barge-in by seconds. + if len(segments) == 0 { + // An open turn whose scan window is all silence: the turn's speech + // is entirely older than the clip, so the trailing silence is at + // least the window — which the window sizing guarantees exceeds + // every commit threshold. Commit with the last speech end this + // turn observed instead of discarding real speech as no-speech. + // With no clip in effect (clipOffsetSec == 0) silero really saw + // the whole turn, and zero segments keeps its historical meaning: + // the earlier onset was reclassified as noise — clear it below. + if _, speaking := sink.coord.State().(turncoord.Speaking); speaking && + clipOffsetSec > 0 && sink.lastSpeechEndSec > 0 { + vadCommit(sink, retranscribe, aints, len(allAudio), audioLength, sink.lastSpeechEndSec, false) + return + } + if audioLength > silenceThreshold { + // "No segments" is not "no speech": silero (threshold 0.5) + // crosses up to a few hundred ms into a soft word onset, so + // the newest audio in the inspected window may be the start + // of a word the next tick will recognize — and more audio + // arrived while this tick ran. Keep both; drop only the + // older, confirmed-silent head, or utterance onsets get cut. + holdback := int(noSpeechHoldbackSec*float64(session.InputSampleRate)) * 2 + session.AudioBufferLock.Lock() + session.InputAudioBuffer = dropInspectedPrefix(session.InputAudioBuffer, len(allAudio), holdback) + session.AudioBufferLock.Unlock() + + // No-speech clear: end any open turn (Speaking -> Idle, discarding + // the partial). Returning to Idle is the fix for failure mode 4 — + // the legacy discardTurn left speechStarted true, suppressing the + // next onset. Idle while not speaking is a no-op. + sink.lastSpeechEndSec = 0 + if err := sink.coord.Apply(turncoord.Abort{Reason: turncoord.AbortNoSpeech}); err != nil { + xlog.Error("turncoord: abort(no_speech) failed", "error", err) + } + } + return + } + + // Speech detected this tick: open the turn (Idle -> Speaking) through + // the coordinator. On that transition it opens the turn's live ASR + // stream + feeds the buffered prefix (OpenTurn), cancels any in-flight + // response (BargeIn, non-blocking — the VAD tick is never stalled), and + // emits speech_started. While already Speaking it is a no-op, so "turn + // open" and "speech started" can never disagree. The turn id is minted + // here and carried by the coordinator through to the committed event. + sink.onsetAudio = aints + if err := sink.coord.Apply(turncoord.Onset{Turn: turncoord.TurnID(generateItemID())}); err != nil { + xlog.Error("turncoord: onset failed", "error", err) + } + + // Track where speech last ended, in whole-buffer seconds: once these + // segments scroll out of the scan clip, the silence-outran-the-window + // commit above still needs a speech end to report. An open segment + // (End == 0) means speech reaches the end of the inspected audio. + if end := segments[len(segments)-1].End; end != 0 { + sink.lastSpeechEndSec = float64(end) + } else { + sink.lastSpeechEndSec = audioLength + } + + if sv != nil { + // Drain again: events produced by THIS tick's feed have + // usually arrived by the time runVAD returns, and leaving + // them for the next tick adds 300ms to every EOU-triggered + // commit. + lts.drainEvents(audioLength) + } + + // Segment still in progress when audio ended + segEndTime := segments[len(segments)-1].End + if segEndTime == 0 { + return + } + + threshold := silenceThreshold + eouPending := false + if sv != nil { + eouPending = lts.eouPending(segments) + threshold = lts.thresholdSec(eouPending, sv) + } + + if float32(audioLength)-segEndTime > float32(threshold) { + vadCommit(sink, retranscribe, aints, len(allAudio), audioLength, float64(segEndTime), eouPending) + } +} + +// vadCommit runs the commit tail of a VAD tick: the semantic commit log, the +// retranscribe gate, the buffer trim, and the coordinator's Silence event +// (speech_stopped + committed + finalize live stream + issue the response). +// Shared by the normal trailing-silence commit and the +// silence-outran-the-scan-window commit. +func vadCommit(sink *turnSink, retranscribe bool, aints []int16, inspectedBytes int, audioLength, segEndTime float64, eouPending bool) { + session := sink.session + lts := sink.lts + + if sink.sv != nil { + trigger, eouLag := lts.commitTrigger(eouPending, segEndTime) + xlog.Info("semantic_vad: committing turn", + "trigger", trigger, + "speech_end_s", segEndTime, + "eou_lag_s", eouLag, + "silence_s", audioLength-segEndTime, + "audio_s", audioLength) + } + // Retranscribe gate (semantic mode, EOU-triggered commits + // only): cross-check the streamed EOU with an offline decode + // of the buffered turn before committing. Runs synchronously + // on the tick — the engine would serialize a concurrent feed + // against it anyway. Timeout-triggered commits skip the gate. + var gated *schema.TranscriptionResult + if retranscribe && eouPending { + batch, gerr := transcribeUtterance(sink.vadContext, sound.Int16toBytesLE(aints), session) + switch { + case gerr != nil: + xlog.Warn("semantic_vad: retranscribe gate failed; committing via the file path", "error", gerr) + case !batch.Eou: + xlog.Info("semantic_vad: batch decode did not confirm the streamed EOU; continuing to listen", + "streamed", lts.previewText(), "batch", batch.Text) + // The batch decode rejected the streamed EOU as a false + // positive: consume the recorded EOU so the next tick + // falls back to the eagerness window instead of + // re-triggering on the same token. + lts.eouAtSec = 0 + return + default: + xlog.Info("semantic_vad: batch decode confirmed the streamed EOU", + "streamed", lts.previewText(), "batch", batch.Text) + gated = batch + } + } + + xlog.Debug("Detected end of speech segment") + session.AudioBufferLock.Lock() + // Keep audio appended while this tick ran — it belongs to + // the next turn (in any mode: nil-ing it dropped the onset + // of an utterance started right after a commit). + session.InputAudioBuffer = dropInspectedPrefix(session.InputAudioBuffer, inspectedBytes, 0) + session.AudioBufferLock.Unlock() + + // Commit the turn through the coordinator: it emits speech_stopped + // (EmitSpeechStopped) then the committed event, finalizes the live + // stream, and issues the response (CommitTurn). The committed item + // id is the coordinator's turn id (== the id the live captions + // streamed under), so the client replaces the partial text. + sink.commitAudio = sound.Int16toBytesLE(aints) + sink.commitAudioLength = audioLength + sink.commitRetranscribe = retranscribe + sink.commitGated = gated + sink.lastSpeechEndSec = 0 + // TODO: Remove prefix silence that is over TurnDetectionParams.PrefixPaddingMs + if err := sink.coord.Apply(turncoord.Silence{}); err != nil { + xlog.Error("turncoord: commit failed", "error", err) + } +} + +// vadScanWindowSec sizes the tail of the buffer silero inspects each tick. +// The window must contain the largest trailing silence the commit test can +// need to measure — server_vad's silence window, or the semantic eagerness +// fallback (the post-EOU window is shorter) — plus vadWarmupMarginSec. +// pipeline.turn_detection.vad_window_sec can widen it; values below the floor +// are ignored, since a narrower window would make long silences unmeasurable +// and turns uncommittable. +func vadScanWindowSec(sv *types.RealtimeSessionSemanticVad, silenceThreshold float64, cfg *config.ModelConfig) float64 { + needed := silenceThreshold + if sv != nil { + needed = eagernessMaxSilenceSec(sv.Eagerness) + } + window := needed + vadWarmupMarginSec + if cfg != nil && cfg.Pipeline.TurnDetection.VadWindowSec > window { + window = cfg.Pipeline.TurnDetection.VadWindowSec + } + return window +} + func commitUtterance(ctx context.Context, utt []byte, session *Session, conv *Conversation, t Transport) { commitUtteranceWithTranscript(ctx, utt, nil, nil, "", session, conv, t) } @@ -1896,6 +2073,11 @@ func runVAD(ctx context.Context, session *Session, adata []int16) ([]schema.VADS if err != nil { return nil, err } + // A backend answering with an empty message means "no speech", not a + // reason to panic the VAD goroutine. + if resp == nil { + return nil, nil + } // If resp.Segments is empty => no speech return resp.Segments, nil @@ -1981,6 +2163,11 @@ type liveResponse struct { output []types.MessageItemUnion usage backend.TokenUsage outcome responseOutcome + // metadata is echoed back on response.created and response.done. It is the + // only thing tying a terminal event to the response.create that asked for + // it, which is what lets a client run an out-of-band response alongside the + // spoken conversation and still recognise its own answer. + metadata map[string]string } func (r *liveResponse) addItem(it types.MessageItemUnion) { r.output = append(r.output, it) } @@ -2010,12 +2197,16 @@ func triggerResponse(ctx context.Context, session *Session, conv *Conversation, // terminals the legacy code emitted (one response.done per turn, with empty // Output/Usage) are gone; tool turns are now internal to this single response. r := &liveResponse{id: generateUniqueID()} + if overrides != nil { + r.metadata = overrides.Metadata + } sendEvent(t, types.ResponseCreatedEvent{ ServerEventBase: types.ServerEventBase{}, Response: types.Response{ - ID: r.id, - Object: "realtime.response", - Status: types.ResponseStatusInProgress, + ID: r.id, + Object: "realtime.response", + Status: types.ResponseStatusInProgress, + Metadata: r.metadata, }, }) @@ -2026,10 +2217,11 @@ func triggerResponse(ctx context.Context, session *Session, conv *Conversation, sendEvent(t, types.ResponseDoneEvent{ ServerEventBase: types.ServerEventBase{}, Response: types.Response{ - ID: r.id, - Object: "realtime.response", - Status: types.ResponseStatusCancelled, - Output: r.output, + ID: r.id, + Object: "realtime.response", + Status: types.ResponseStatusCancelled, + Output: r.output, + Metadata: r.metadata, }, }) case outcomeFailed: @@ -2039,11 +2231,12 @@ func triggerResponse(ctx context.Context, session *Session, conv *Conversation, sendEvent(t, types.ResponseDoneEvent{ ServerEventBase: types.ServerEventBase{}, Response: types.Response{ - ID: r.id, - Object: "realtime.response", - Status: types.ResponseStatusCompleted, - Output: r.output, - Usage: responseUsage(r.usage), + ID: r.id, + Object: "realtime.response", + Status: types.ResponseStatusCompleted, + Output: r.output, + Usage: responseUsage(r.usage), + Metadata: r.metadata, }, }) } @@ -2216,9 +2409,18 @@ func triggerResponseAtTurn(ctx context.Context, session *Session, conv *Conversa images = append(images, m.StringImages...) } - // response.created/done are emitted once per response.create by triggerResponse; - // every turn (including agentic recursion) shares this id. - responseID := r.id + // Classifier mode replaces autoregressive generation for the first turn + // of a response: prefill-only scoring picks a registered option and its + // canned reply/tool is emitted through the standard response protocol. + // Agentic follow-ups (toolTurn > 0) always generate — the option list + // describes user intents, not tool outputs. This branch must precede the + // streamed-LLM path below or streaming pipelines would bypass it. + if cc := resolveClassifier(session.Classifier, overrides); toolTurn == 0 && cc.Active() { + if classifierRespond(ctx, session, conv, t, r, cc, conversationHistory, overrides, toolTurn) { + return + } + // fallback.mode "generate": fall through to normal generation. + } // Streamed LLM path: when the pipeline opts into LLM streaming, stream the // transcript to the client as it is generated and synthesize the buffered @@ -2371,6 +2573,30 @@ func triggerResponseAtTurn(ctx context.Context, session *Session, conv *Conversa } if finalSpeech != "" { + if !emitAssistantMessage(ctx, session, conv, t, r, finalSpeech, overrides) { + return + } + } + + // Emit the parsed tool calls and (for server-side assistant tools) the + // follow-up turn. Shared with the streamed path so both finalize tool calls + // identically. The single terminal is emitted by triggerResponse. + emitToolCallItems(ctx, session, conv, t, r, finalToolCalls, finalSpeech != "", toolTurn) +} + +// emitAssistantMessage appends an assistant item carrying finalSpeech to the +// conversation and emits the standard response events for it — +// output_item.added, content_part.added, audio-transcript or output-text +// deltas, TTS audio via emitSpeech (unless the resolved modalities are +// text-only), content_part.done and output_item.done. Shared by the buffered +// generation path and classifier mode. Returns false when the response was +// cancelled (barge-in) or failed — r.outcome is already recorded and the +// caller must emit no further items. +func emitAssistantMessage(ctx context.Context, session *Session, conv *Conversation, t Transport, r *liveResponse, finalSpeech string, overrides *types.ResponseCreateParams) bool { + // response.created/done are emitted once per response.create by + // triggerResponse; every turn (including agentic recursion) shares this id. + responseID := r.id + { // Create the assistant item now that we have content item := types.MessageItemUnion{ Assistant: &types.MessageItemAssistant{ @@ -2438,7 +2664,7 @@ func triggerResponseAtTurn(ctx context.Context, session *Session, conv *Conversa if ctx.Err() != nil { xlog.Debug("Response cancelled before TTS (barge-in)") sendCancelledResponse() - return + return false } // Transcript of the spoken reply (the audio's text). @@ -2468,12 +2694,12 @@ func triggerResponseAtTurn(ctx context.Context, session *Session, conv *Conversa if ctx.Err() != nil { xlog.Debug("TTS cancelled (barge-in)") sendCancelledResponse() - return + return false } xlog.Error("TTS failed", "error", err) sendError(t, "tts_error", fmt.Sprintf("TTS generation failed: %v", err), "", item.Assistant.ID) r.outcome = outcomeFailed - return + return false } if !isWebRTC { audioString = base64.StdEncoding.EncodeToString(pcmAudio) @@ -2532,11 +2758,7 @@ func triggerResponseAtTurn(ctx context.Context, session *Session, conv *Conversa }) r.addItem(item) } - - // Emit the parsed tool calls and (for server-side assistant tools) the - // follow-up turn. Shared with the streamed path so both finalize tool calls - // identically. The single terminal is emitted by triggerResponse. - emitToolCallItems(ctx, session, conv, t, r, finalToolCalls, finalSpeech != "", toolTurn) + return true } // emitToolCallItems emits the realtime function_call items for the parsed tool diff --git a/core/http/endpoints/openai/realtime_classifier.go b/core/http/endpoints/openai/realtime_classifier.go new file mode 100644 index 000000000..c310f3f00 --- /dev/null +++ b/core/http/endpoints/openai/realtime_classifier.go @@ -0,0 +1,616 @@ +package openai + +import ( + "context" + "encoding/json" + "fmt" + "regexp" + "strings" + "time" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/endpoints/openai/types" + "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/services/routing/router" + "github.com/mudler/LocalAI/pkg/functions" + "github.com/mudler/xlog" +) + +// Classifier mode (LocalAI extension): instead of autoregressive +// generation, each user turn is prefill-scored against a registered option +// list via the Score primitive and the winning option's canned reply / +// tool call is emitted. Designed for hardware that can afford prefill but +// not decode. See docs/content/features/openai-realtime.md. + +// By default only the latest user message is scored. Earlier turns in the +// probe — the assistant's canned replies especially — echo option names +// ("Going up." ↔ up) and verified empirically to dominate small scoring +// models: with any prior turn present, a 1.2B model kept re-choosing the +// previous option at p≈1.0 regardless of the new command. history_items > 0 +// opts back into context (role-labeled), for larger scoring models. + +// classifierConfigFromPipeline converts the YAML pipeline.classifier block +// into the wire ClassifierConfig and validates it, so a bad option list +// rejects the session at setup rather than misbehaving on the first turn. +// A nil block yields a nil config (classifier off). +func classifierConfigFromPipeline(p *config.PipelineClassifier) (*types.ClassifierConfig, error) { + if p == nil { + return nil, nil + } + cc := &types.ClassifierConfig{ + Enabled: &p.Enabled, + Threshold: p.Threshold, + Normalization: p.Normalization, + HistoryItems: p.HistoryItems, + } + if p.Fallback != nil { + cc.Fallback = &types.ClassifierFallback{Mode: p.Fallback.Mode, Reply: p.Fallback.Reply} + } + if p.Address != nil { + cc.Address = &types.ClassifierAddress{Names: p.Address.Names, Mode: p.Address.Mode, Reply: p.Address.Reply} + } + for _, o := range p.Options { + opt := types.ClassifierOption{ + ID: o.ID, + Description: o.Description, + Reply: o.Reply, + } + if o.Tool != nil { + args := json.RawMessage(nil) + if o.Tool.Arguments != nil { + data, err := json.Marshal(o.Tool.Arguments) + if err != nil { + return nil, fmt.Errorf("option %q: marshal tool arguments: %w", o.ID, err) + } + args = data + } + opt.Tool = &types.ClassifierTool{Name: o.Tool.Name, Arguments: args} + for _, s := range o.Tool.Slots { + opt.Tool.Slots = append(opt.Tool.Slots, types.ClassifierSlot{ + Name: s.Name, + Type: s.Type, + Values: s.Values, + Default: s.Default, + Hint: s.Hint, + }) + } + } + cc.Options = append(cc.Options, opt) + } + if err := cc.Validate(); err != nil { + return nil, err + } + return cc, nil +} + +// prewarmClassifier primes the scoring prompt cache for the session's +// current classifier config in the background: registration returns +// immediately, and by the time the canned mode-switch reply finishes +// speaking, the new option list's prompt (and, on hybrid/recurrent +// models, a rewind checkpoint at the per-turn probe boundary) is already +// in the backend's cache. The context is deliberately detached from the +// registering request — the warmed cache belongs to the backend, not the +// request. +func prewarmClassifier(session *Session) { + cc := session.Classifier + if session.ModelInterface == nil || !cc.Active() { + return + } + options, normalization := cc.Options, cc.Normalization + go func() { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + session.ModelInterface.PrewarmClassifier(ctx, options, normalization) + }() +} + +// resolveClassifier merges the session classifier config with a +// response-level override: a non-nil override replaces the whole block +// (same replace-not-merge semantics as tools), so {"enabled": false} runs +// normal generation for one response. +func resolveClassifier(sessionCfg *types.ClassifierConfig, overrides *types.ResponseCreateParams) *types.ClassifierConfig { + if overrides != nil && overrides.LocalAIClassifier != nil { + return overrides.LocalAIClassifier + } + return sessionCfg +} + +// validateClassifierActivation verifies both the wire config and the concrete +// backend selected to score it. Scoring capacity is reserved at model load +// only for configs that explicitly declare the score usecase, so accepting an +// active classifier on any other model would defer a deterministic failure to +// the first response. +func validateClassifierActivation(m Model, cc *types.ClassifierConfig) error { + if cc == nil { + return nil + } + if err := cc.Validate(); err != nil { + return err + } + if !cc.Active() { + return nil + } + wm, ok := m.(*wrappedModel) + if !ok { + return fmt.Errorf("classifier: the session model does not support scoring") + } + cfg := wm.scoreConfig() + if cfg == nil || !cfg.HasUsecases(config.FLAG_SCORE) { + name := "" + if cfg != nil { + name = cfg.Name + } + return fmt.Errorf("classifier: scoring model %q must declare known_usecases: [score]", name) + } + if cfg.HasRouter() { + return fmt.Errorf("classifier: scoring model %q is a router; configure a concrete pipeline.classifier.model", cfg.Name) + } + return nil +} + +// trimClassifierHistory drops system messages (the classifier builds its +// own option-list system prompt) and selects what gets scored. +// historyItems <= 0 (the default): only the latest user message. Positive +// N: the trailing N conversation messages. +func trimClassifierHistory(history schema.Messages, historyItems int) schema.Messages { + conversation := make(schema.Messages, 0, len(history)) + for _, m := range history { + if m.Role == string(types.MessageRoleSystem) { + continue + } + conversation = append(conversation, m) + } + if historyItems <= 0 { + for i := len(conversation) - 1; i >= 0; i-- { + if conversation[i].Role == string(types.MessageRoleUser) { + return conversation[i : i+1] + } + } + return nil + } + if len(conversation) > historyItems { + conversation = conversation[len(conversation)-historyItems:] + } + return conversation +} + +// latestUserText returns the text of the most recent user message — the +// turn the address gate inspects (earlier turns being addressed doesn't +// make this one addressed). +func latestUserText(messages schema.Messages) string { + for i := len(messages) - 1; i >= 0; i-- { + if messages[i].Role == string(types.MessageRoleUser) { + text, _ := messages[i].Content.(string) + return text + } + } + return "" +} + +// mentionsAnyName reports whether text contains any of the names as a +// case-insensitive whole word ("drone" matches "Drone, go up" but not +// "drones"). +func mentionsAnyName(text string, names []string) bool { + for _, n := range names { + n = strings.TrimSpace(n) + if n == "" { + continue + } + re, err := regexp.Compile(`(?i)\b` + regexp.QuoteMeta(n) + `\b`) + if err != nil { + continue + } + if re.MatchString(text) { + return true + } + } + return false +} + +// classifierProbe renders the trimmed history for scoring. A single user +// message goes in verbatim — that matches the scoring format's training +// distribution (Arch-Router scores "the user's request"). When +// history_items opts extra turns in, every line carries a role label so +// the scoring model can at least tell the user's request apart from the +// assistant's replies. +func classifierProbe(messages schema.Messages) router.Probe { + parts := make([]string, 0, len(messages)) + label := len(messages) > 1 + for _, msg := range messages { + text, _ := msg.Content.(string) + if text == "" { + continue // e.g. tool-call items carry no text + } + if label { + switch msg.Role { + case string(types.MessageRoleAssistant): + text = "Assistant: " + text + case "tool": + text = "Tool: " + text + default: + text = "User: " + text + } + } + parts = append(parts, text) + } + return router.Probe{Prompt: router.JoinTurns(parts), Messages: parts} +} + +// classifierRespond runs one classifier-mode response: score the options, +// emit the localai.classifier.result observability event, then either the +// winning option's canned reply/tool, the fallback reply, nothing, or — +// for the generate fallback — report false so the caller falls through to +// normal generation. Runs inside the respcoord-issued response body, so +// the single terminal stays owned by triggerResponse. Returns true when +// the response was fully handled here. +func classifierRespond(ctx context.Context, session *Session, conv *Conversation, t Transport, r *liveResponse, cc *types.ClassifierConfig, history schema.Messages, overrides *types.ResponseCreateParams, toolTurn int) bool { + msgs := trimClassifierHistory(history, cc.HistoryItems) + if len(msgs) == 0 { + xlog.Debug("realtime classifier: no scorable conversation content; skipping to generation") + return false + } + + // Address gate (wake-word behavior): when configured, a turn that + // doesn't mention one of the assistant's names is dropped before any + // scoring — the check is a deterministic word match on the transcript + // because scoring cannot detect the missing name (command semantics + // dominate the softmax), and skipping the Score call keeps ambient + // conversation free on weak hardware. + if ad := cc.Address; ad != nil && !mentionsAnyName(latestUserText(msgs), ad.Names) { + sendEvent(t, types.ClassifierResultEvent{ + ResponseID: r.id, + Scores: []types.ClassifierScore{}, + Threshold: cc.Threshold, + Fallback: types.ClassifierNotAddressed, + }) + xlog.Debug("realtime classifier: turn does not address the assistant; dropping", "mode", ad.AddressMode()) + if ctx.Err() != nil { + r.outcome = outcomeCancelled + return true + } + if ad.AddressMode() == types.ClassifierAddressReply && ad.Reply != "" { + if !emitAssistantMessage(ctx, session, conv, t, r, ad.Reply, overrides) { + return true + } + emitToolCallItems(ctx, session, conv, t, r, nil, true, toolTurn) + return true + } + // ignore: complete the response with no output items. + emitToolCallItems(ctx, session, conv, t, r, nil, false, toolTurn) + return true + } + + // A committed turn can carry no words at all (the VAD fires on noise + // and the ASR transcribes nothing). Scoring an empty prompt returns a + // confidently arbitrary winner — measured p≈0.95 for the first option + // — so skip scoring entirely and treat it like a below-threshold turn. + var scores []router.LabelScore + var latency time.Duration + if strings.TrimSpace(classifierProbe(msgs).Prompt) != "" { + start := time.Now() + var err error + scores, err = session.ModelInterface.ClassifyTurn(ctx, msgs, cc.Options, cc.Normalization) + if err != nil { + if cc.FallbackMode() == types.ClassifierFallbackGenerate { + xlog.Warn("realtime classifier: scoring failed; falling back to generation", "error", err) + return false + } + sendError(t, "classifier_failed", fmt.Sprintf("classifier scoring failed: %v", err), "", "") + r.outcome = outcomeFailed + return true + } + latency = time.Since(start) + } else if cc.FallbackMode() == types.ClassifierFallbackGenerate { + xlog.Debug("realtime classifier: turn has no scorable text; falling back to generation") + return false + } + + best := -1 + for i := range scores { + if best < 0 || scores[i].Score > scores[best].Score { + best = i + } + } + var chosen *types.ClassifierOption + chosenID := "" + fallbackApplied := "" + if best >= 0 && scores[best].Score >= cc.Threshold { + chosen = &cc.Options[best] + chosenID = chosen.ID + } else { + fallbackApplied = cc.FallbackMode() + } + + // Hybrid path: a winning option with argument slots gets them filled by + // a constrained completion before anything is emitted, so the result + // event carries the final arguments. An unrecoverable fill failure + // (error and no complete default set) is handled like a scoring + // failure. + filledArgs := "" + var fillValues map[string]string + var fillLatency time.Duration + if chosen != nil { + var ferr error + filledArgs, fillValues, fillLatency, ferr = fillChosenArguments(ctx, session, cc, msgs, chosen) + if ferr != nil { + if cc.FallbackMode() == types.ClassifierFallbackGenerate { + xlog.Warn("realtime classifier: slot fill failed; falling back to generation", "error", ferr) + return false + } + sendError(t, "classifier_failed", fmt.Sprintf("classifier slot fill failed: %v", ferr), "", "") + r.outcome = outcomeFailed + return true + } + } + + evScores := make([]types.ClassifierScore, len(scores)) + for i, s := range scores { + evScores[i] = types.ClassifierScore{ID: s.Label, Score: s.Score} + } + evArgs := "" + if chosen != nil && chosen.Tool != nil && len(chosen.Tool.Slots) > 0 { + evArgs = filledArgs + } + sendEvent(t, types.ClassifierResultEvent{ + ResponseID: r.id, + Scores: evScores, + ChosenID: chosenID, + Threshold: cc.Threshold, + Fallback: fallbackApplied, + LatencyMs: latency.Milliseconds(), + Arguments: evArgs, + FillLatencyMs: fillLatency.Milliseconds(), + }) + topScore := 0.0 + if best >= 0 { + topScore = scores[best].Score + } + xlog.Debug("realtime classifier: scored turn", + "chosen", chosenID, "top_score", topScore, + "threshold", cc.Threshold, "fallback", fallbackApplied, + "latency_ms", latency.Milliseconds(), + "arguments", evArgs, "fill_latency_ms", fillLatency.Milliseconds()) + + if fallbackApplied == types.ClassifierFallbackGenerate { + return false + } + + // Barge-in may have fired during scoring. + if ctx.Err() != nil { + r.outcome = outcomeCancelled + return true + } + + reply := "" + var toolCalls []functions.FuncCallResults + switch { + case chosen != nil: + // The reply may template the filled slot values ("Going forward + // {{distance}} {{units}}.") so what is spoken confirms what was + // actually inferred. + reply = chosen.SpliceReply(fillValues) + if chosen.Tool != nil { + toolCalls = []functions.FuncCallResults{{Name: chosen.Tool.Name, Arguments: filledArgs}} + } + case fallbackApplied == types.ClassifierFallbackReply: + reply = cc.Fallback.Reply + default: + // fallback "none": complete with no output items. + } + + if reply != "" { + if !emitAssistantMessage(ctx, session, conv, t, r, reply, overrides) { + // Cancelled or failed — outcome already recorded. + return true + } + } + // Always finalize through emitToolCallItems, mirroring the generation + // path: it emits the function_call items (client executes canned tools + // and reports back via conversation.item.create) and runs server-side + // assistant tools inproc. + emitToolCallItems(ctx, session, conv, t, r, toolCalls, reply != "", toolTurn) + return true +} + +// ---- slot filling (hybrid classify-then-complete) -------------------------- +// +// A winning option whose tool declares slots gets its argument values from a +// short constrained completion: the prompt is the exact scoring prompt (warm +// in the backend's cache) continued by the chosen route JSON re-opened at the +// first slot field, and a GBNF grammar pins everything except the slot +// values. The generated tail is parsed back through the JSON object it +// completes, and the values are spliced into the tool's argument template. + +// gbnfLiteral renders s as a GBNF quoted literal. +func gbnfLiteral(s string) string { + r := strings.NewReplacer(`\`, `\\`, `"`, `\"`, "\n", `\n`) + return `"` + r.Replace(s) + `"` +} + +// slotFillGrammar builds the grammar for the completion tail: first slot +// value, then each further slot as a forced `, "": ` literal plus its +// value, then the closing brace. +func slotFillGrammar(slots []types.ClassifierSlot) string { + var root strings.Builder + var rules strings.Builder + needNum, needStr := false, false + root.WriteString("root ::= ") + for i := range slots { + if i > 0 { + root.WriteString(" " + gbnfLiteral(`, "`+slots[i].Name+`": `) + " ") + } + fmt.Fprintf(&root, "slot%d", i) + fmt.Fprintf(&rules, "\nslot%d ::= ", i) + switch slots[i].Type { + case types.ClassifierSlotNumber: + rules.WriteString("num") + needNum = true + case types.ClassifierSlotEnum: + for vi, v := range slots[i].Values { + if vi > 0 { + rules.WriteString(" | ") + } + encoded, _ := json.Marshal(v) // validation rejects values JSON cannot encode + rules.WriteString(gbnfLiteral(string(encoded))) + } + default: // string + rules.WriteString("str") + needStr = true + } + } + root.WriteString(` "}"`) + if needNum { + rules.WriteString("\nnum ::= \"-\"? [0-9] [0-9]* (\".\" [0-9] [0-9]*)?") + } + if needStr { + rules.WriteString("\nstr ::= \"\\\"\" [^\"\\\\\\n]* \"\\\"\"") + } + return root.String() + rules.String() +} + +const ( + // Free-form values need an explicit ceiling; forced enum values and field + // syntax are budgeted from their actual JSON encoding below. + slotFillStringTokens = 64 + slotFillNumberTokens = 32 +) + +// slotFillMaxTokens conservatively budgets one token per output byte for the +// forced JSON tail, plus explicit allowances for free-form values. This avoids +// truncating long enum values or field names while keeping string generation +// bounded. +func slotFillMaxTokens(slots []types.ClassifierSlot) int { + tokens := 1 // closing brace + for i := range slots { + if i > 0 { + field, _ := json.Marshal(slots[i].Name) + tokens += len(field) + len(`, : `) + } + switch slots[i].Type { + case types.ClassifierSlotNumber: + tokens += slotFillNumberTokens + case types.ClassifierSlotString: + tokens += slotFillStringTokens + case types.ClassifierSlotEnum: + longest := 0 + for _, value := range slots[i].Values { + encoded, _ := json.Marshal(value) + if len(encoded) > longest { + longest = len(encoded) + } + } + tokens += longest + } + } + return tokens +} + +// slotFillContextReserve includes both the generated tail and the continuation +// prefix appended after the scored prompt. It intentionally over-reserves by +// counting bytes as tokens; preserving the identical scoring prompt is more +// important than reclaiming a handful of context tokens. +func slotFillContextReserve(option *types.ClassifierOption) int { + if option == nil || option.Tool == nil || len(option.Tool.Slots) == 0 { + return 0 + } + route, _ := json.Marshal(option.ID) + field, _ := json.Marshal(option.Tool.Slots[0].Name) + prefixBytes := len(`{"route": , : `) + len(route) + len(field) + return prefixBytes + slotFillMaxTokens(option.Tool.Slots) +} + +// parseSlotValues closes the completed route JSON and extracts each slot's +// value as the string form SpliceArguments expects. +func parseSlotValues(chosenID, firstSlot, generated string, slots []types.ClassifierSlot) (map[string]string, error) { + idJSON, _ := json.Marshal(chosenID) + full := `{"route": ` + string(idJSON) + `, "` + firstSlot + `": ` + strings.TrimSpace(generated) + if !strings.HasSuffix(strings.TrimSpace(generated), "}") { + full += "}" + } + dec := json.NewDecoder(strings.NewReader(full)) + dec.UseNumber() + var obj map[string]any + if err := dec.Decode(&obj); err != nil { + return nil, fmt.Errorf("classifier: slot completion %q does not parse: %w", generated, err) + } + values := make(map[string]string, len(slots)) + for i := range slots { + v, ok := obj[slots[i].Name] + if !ok { + return nil, fmt.Errorf("classifier: slot completion missing %q", slots[i].Name) + } + switch tv := v.(type) { + case json.Number: + values[slots[i].Name] = tv.String() + case string: + values[slots[i].Name] = tv + default: + return nil, fmt.Errorf("classifier: slot %q has unexpected value type %T", slots[i].Name, v) + } + } + return values, nil +} + +// fillChosenArguments resolves a winning option's tool arguments: canned +// options pass through, slotted options run the fill completion with a +// default-value recovery when inference fails. The slot values ride along +// so the caller can splice them into the spoken reply too. The error return +// is reserved for unrecoverable failures (no complete default set). +func fillChosenArguments(ctx context.Context, session *Session, cc *types.ClassifierConfig, msgs schema.Messages, chosen *types.ClassifierOption) (args string, values map[string]string, latency time.Duration, err error) { + if chosen.Tool == nil { + return "", nil, 0, nil + } + if len(chosen.Tool.Slots) == 0 { + if len(chosen.Tool.Arguments) > 0 { + return string(chosen.Tool.Arguments), nil, 0, nil + } + return "{}", nil, 0, nil + } + start := time.Now() + args, values, err = session.ModelInterface.FillToolArguments(ctx, msgs, cc.Options, cc.Normalization, chosen) + latency = time.Since(start) + if err == nil { + return args, values, latency, nil + } + xlog.Warn("realtime classifier: slot fill failed; trying slot defaults", "option", chosen.ID, "error", err) + defaults, derr := chosen.Tool.SlotDefaults() + if derr != nil { + return "", nil, latency, err + } + args, derr = chosen.Tool.SpliceArguments(defaults) + if derr != nil { + return "", nil, latency, err + } + return args, defaults, latency, nil +} + +// classifierPolicyDescription renders an option's scoring description, +// appending any slot declarations so the model both weighs the parameters +// during scoring and knows how to fill them ("assume meters…") during the +// slot completion — the hints ride the shared system prompt, costing no +// extra per-turn tokens. +func classifierPolicyDescription(o *types.ClassifierOption) string { + if o.Tool == nil || len(o.Tool.Slots) == 0 { + return o.Description + } + var b strings.Builder + b.WriteString(o.Description) + b.WriteString(" — route parameters:") + for i := range o.Tool.Slots { + s := &o.Tool.Slots[i] + if i > 0 { + b.WriteString(";") + } + b.WriteString(" " + s.Name) + switch s.Type { + case types.ClassifierSlotEnum: + b.WriteString(" (one of: " + strings.Join(s.Values, ", ") + ")") + default: + b.WriteString(" (" + s.Type + ")") + } + if s.Hint != "" { + b.WriteString(", " + s.Hint) + } + } + return b.String() +} diff --git a/core/http/endpoints/openai/realtime_classifier_test.go b/core/http/endpoints/openai/realtime_classifier_test.go new file mode 100644 index 000000000..b8630ca81 --- /dev/null +++ b/core/http/endpoints/openai/realtime_classifier_test.go @@ -0,0 +1,739 @@ +package openai + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/endpoints/openai/types" + "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/services/routing/router" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func classifierTestConfig(threshold float64, fallback *types.ClassifierFallback) *types.ClassifierConfig { + return &types.ClassifierConfig{ + Threshold: threshold, + Fallback: fallback, + Options: []types.ClassifierOption{ + { + ID: "up", + Description: "the user asks the drone to fly up", + Reply: "Going up.", + Tool: &types.ClassifierTool{Name: "move", Arguments: json.RawMessage(`{"direction":"up"}`)}, + }, + {ID: "greeting", Description: "the user greets the assistant", Reply: "Hello."}, + }, + } +} + +func classifierTestSession(m *fakeModel) *Session { + return &Session{ + ModelInterface: m, + OutputModalities: []types.Modality{types.ModalityText}, + ModelConfig: &config.ModelConfig{}, + } +} + +var classifierTestHistory = schema.Messages{ + {Role: "system", StringContent: "instructions", Content: "instructions"}, + {Role: "user", StringContent: "please go up", Content: "please go up"}, +} + +func classifierResultEvents(t *fakeTransport) []types.ClassifierResultEvent { + var out []types.ClassifierResultEvent + for _, e := range t.events { + if ev, ok := e.(types.ClassifierResultEvent); ok { + out = append(out, ev) + } + } + return out +} + +// replyTexts collects the assistant reply text of every completed output +// item — what a classifier response actually "spoke". +func replyTexts(t *fakeTransport) []string { + var out []string + for _, e := range t.events { + if ev, ok := e.(types.ResponseOutputTextDoneEvent); ok { + out = append(out, ev.Text) + } + } + return out +} + +var _ = Describe("prewarmClassifier", func() { + It("prewarms an active option list in the background", func() { + m := &fakeModel{} + session := classifierTestSession(m) + session.Classifier = classifierTestConfig(0.35, nil) + + prewarmClassifier(session) + + Eventually(func() int { n, _ := m.prewarmed(); return n }).Should(Equal(1)) + _, opts := m.prewarmed() + Expect(opts).To(HaveLen(len(session.Classifier.Options))) + }) + + It("does nothing without an active classifier", func() { + m := &fakeModel{} + session := classifierTestSession(m) + prewarmClassifier(session) + + off := false + session.Classifier = &types.ClassifierConfig{Enabled: &off, Options: classifierTestConfig(0.35, nil).Options} + prewarmClassifier(session) + + Consistently(func() int { n, _ := m.prewarmed(); return n }, "150ms").Should(BeZero()) + }) +}) + +var _ = Describe("classifierConfigFromPipeline", func() { + It("returns nil for an absent block", func() { + cc, err := classifierConfigFromPipeline(nil) + Expect(err).ToNot(HaveOccurred()) + Expect(cc).To(BeNil()) + }) + + It("converts options and tool argument maps to wire form", func() { + cc, err := classifierConfigFromPipeline(&config.PipelineClassifier{ + Enabled: true, + Threshold: 0.4, + Fallback: &config.PipelineClassifierFallback{Mode: "reply", Reply: "Say again?"}, + Options: []config.PipelineClassifierOption{ + { + ID: "up", + Description: "fly up", + Reply: "Going up.", + Tool: &config.PipelineClassifierTool{Name: "move", Arguments: map[string]any{"direction": "up"}}, + }, + }, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(cc.Active()).To(BeTrue()) + Expect(cc.Threshold).To(Equal(0.4)) + Expect(cc.Options).To(HaveLen(1)) + Expect(string(cc.Options[0].Tool.Arguments)).To(MatchJSON(`{"direction":"up"}`)) + Expect(cc.Fallback.Mode).To(Equal(types.ClassifierFallbackReply)) + }) + + It("rejects invalid blocks via the shared validation", func() { + _, err := classifierConfigFromPipeline(&config.PipelineClassifier{ + Enabled: true, + Options: []config.PipelineClassifierOption{ + {ID: "a", Description: "one"}, + {ID: "a", Description: "two"}, + }, + }) + Expect(err).To(MatchError(ContainSubstring("duplicate option id"))) + }) +}) + +var _ = Describe("validateClassifierActivation", func() { + It("accepts a combined inference and score model", func() { + usecases := config.FLAG_CHAT | config.FLAG_SCORE + m := &wrappedModel{LLMConfig: &config.ModelConfig{KnownUsecases: &usecases}} + Expect(validateClassifierActivation(m, classifierTestConfig(0.4, nil))).To(Succeed()) + }) + + It("rejects an active classifier when the model does not declare score", func() { + usecases := config.FLAG_CHAT + m := &wrappedModel{LLMConfig: &config.ModelConfig{KnownUsecases: &usecases}} + Expect(validateClassifierActivation(m, classifierTestConfig(0.4, nil))).To(MatchError(ContainSubstring("known_usecases"))) + }) + + It("rejects a router config as the concrete scoring model", func() { + usecases := config.FLAG_SCORE + m := &wrappedModel{LLMConfig: &config.ModelConfig{ + KnownUsecases: &usecases, + Router: config.RouterConfig{Candidates: []config.RouterCandidate{{Model: "target"}}}, + }} + Expect(validateClassifierActivation(m, classifierTestConfig(0.4, nil))).To(MatchError(ContainSubstring("concrete"))) + }) + + It("allows disabling classification without score support", func() { + disabled := false + m := &wrappedModel{LLMConfig: &config.ModelConfig{}} + Expect(validateClassifierActivation(m, &types.ClassifierConfig{Enabled: &disabled})).To(Succeed()) + }) +}) + +var _ = Describe("resolveClassifier", func() { + It("uses the session config when no override is present", func() { + sess := classifierTestConfig(0, nil) + Expect(resolveClassifier(sess, nil)).To(BeIdenticalTo(sess)) + Expect(resolveClassifier(sess, &types.ResponseCreateParams{})).To(BeIdenticalTo(sess)) + }) + + It("replaces the whole config when the response overrides it", func() { + sess := classifierTestConfig(0, nil) + disabled := false + over := &types.ClassifierConfig{Enabled: &disabled} + got := resolveClassifier(sess, &types.ResponseCreateParams{LocalAIClassifier: over}) + Expect(got).To(BeIdenticalTo(over)) + Expect(got.Active()).To(BeFalse()) + }) +}) + +var _ = Describe("trimClassifierHistory", func() { + history := schema.Messages{ + {Role: "system", StringContent: "sys"}, + {Role: "user", StringContent: "one"}, + {Role: "assistant", StringContent: "two"}, + {Role: "user", StringContent: "three"}, + {Role: "assistant", StringContent: "four"}, + {Role: "user", StringContent: "five"}, + } + + It("keeps only the latest user message by default", func() { + // Earlier turns echo option names (canned replies) and empirically + // dominate small scoring models, so the default is user-turn-only. + got := trimClassifierHistory(history, 0) + Expect(got).To(HaveLen(1)) + Expect(got[0].StringContent).To(Equal("five")) + }) + + It("keeps only the latest user message for -1", func() { + got := trimClassifierHistory(history, -1) + Expect(got).To(HaveLen(1)) + Expect(got[0].StringContent).To(Equal("five")) + }) + + It("honors an explicit cap", func() { + got := trimClassifierHistory(history, 2) + Expect(got).To(HaveLen(2)) + Expect(got[0].StringContent).To(Equal("four")) + }) +}) + +var _ = Describe("mentionsAnyName", func() { + It("matches case-insensitive whole words in any position", func() { + Expect(mentionsAnyName("Drone, go up", []string{"drone"})).To(BeTrue()) + Expect(mentionsAnyName("go up drone", []string{"drone"})).To(BeTrue()) + Expect(mentionsAnyName("go up", []string{"drone"})).To(BeFalse()) + // Whole-word: no substring matches. + Expect(mentionsAnyName("I like drones", []string{"drone"})).To(BeFalse()) + // Multiple aliases and multi-word names. + Expect(mentionsAnyName("hey quadcopter rise", []string{"drone", "quadcopter"})).To(BeTrue()) + Expect(mentionsAnyName("okay drone go", []string{"okay drone"})).To(BeTrue()) + }) +}) + +var _ = Describe("classifierProbe", func() { + It("renders a single user message verbatim", func() { + probe := classifierProbe(schema.Messages{{Role: "user", Content: "fly forward"}}) + Expect(probe.Prompt).To(Equal("fly forward\n")) + Expect(probe.Messages).To(Equal([]string{"fly forward"})) + }) + + It("role-labels multi-message histories and skips text-less items", func() { + probe := classifierProbe(schema.Messages{ + {Role: "user", Content: "go up"}, + {Role: "assistant", Content: "Going up."}, + {Role: "assistant"}, // tool-call item: no text + {Role: "tool", Content: "ok: moved"}, + {Role: "user", Content: "fly forward"}, + }) + Expect(probe.Messages).To(Equal([]string{ + "User: go up", + "Assistant: Going up.", + "Tool: ok: moved", + "User: fly forward", + })) + }) +}) + +var _ = Describe("classifierRespond", func() { + It("emits the winning option's canned reply and tool call", func() { + m := &fakeModel{classifyScores: []router.LabelScore{ + {Label: "up", Score: 0.9}, + {Label: "greeting", Score: 0.1}, + }} + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp1"} + + handled := classifierRespond(context.Background(), session, conv, t, r, classifierTestConfig(0.35, nil), classifierTestHistory, nil, 0) + + Expect(handled).To(BeTrue()) + Expect(m.classifyCalls).To(Equal(1)) + // System instructions stay out of the scoring prompt. + for _, msg := range m.lastMessages { + Expect(msg.Role).ToNot(Equal("system")) + } + + results := classifierResultEvents(t) + Expect(results).To(HaveLen(1)) + Expect(results[0].ChosenID).To(Equal("up")) + Expect(results[0].Fallback).To(BeEmpty()) + Expect(results[0].Scores).To(HaveLen(2)) + Expect(results[0].Scores[0].Score).To(BeNumerically("~", 0.9)) + + // Canned reply as text (text-only modality), canned tool call after it. + Expect(t.countEvents(types.ServerEventTypeResponseOutputTextDone)).To(Equal(1)) + Expect(t.countEvents(types.ServerEventTypeResponseFunctionCallArgumentsDone)).To(Equal(1)) + var fcArgs string + for _, e := range t.events { + if done, ok := e.(types.ResponseFunctionCallArgumentsDoneEvent); ok { + fcArgs = done.Arguments + } + } + Expect(fcArgs).To(MatchJSON(`{"direction":"up"}`)) + // Assistant reply + function_call item recorded in the conversation. + Expect(conv.Items).To(HaveLen(2)) + Expect(conv.Items[0].Assistant).ToNot(BeNil()) + Expect(conv.Items[1].FunctionCall).ToNot(BeNil()) + Expect(conv.Items[1].FunctionCall.Name).To(Equal("move")) + }) + + It("drops unaddressed turns without scoring when the address gate is on", func() { + m := &fakeModel{classifyScores: []router.LabelScore{{Label: "up", Score: 0.99}}} + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp-unaddressed"} + cc := classifierTestConfig(0.35, nil) + cc.Address = &types.ClassifierAddress{Names: []string{"drone"}} + history := schema.Messages{ + {Role: "user", StringContent: "go up", Content: "go up"}, + } + + handled := classifierRespond(context.Background(), session, conv, t, r, cc, history, nil, 0) + + Expect(handled).To(BeTrue()) + Expect(m.classifyCalls).To(BeZero(), "unaddressed turns must not be scored") + results := classifierResultEvents(t) + Expect(results).To(HaveLen(1)) + Expect(results[0].Scores).To(BeEmpty()) + Expect(results[0].Fallback).To(Equal(types.ClassifierNotAddressed)) + Expect(t.countEvents(types.ServerEventTypeResponseOutputTextDone)).To(BeZero(), "ignore mode must stay silent") + }) + + It("scores turns that address the assistant by name", func() { + m := &fakeModel{classifyScores: []router.LabelScore{ + {Label: "up", Score: 0.9}, + {Label: "greeting", Score: 0.1}, + }} + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp-addressed"} + cc := classifierTestConfig(0.35, nil) + cc.Address = &types.ClassifierAddress{Names: []string{"drone"}} + history := schema.Messages{ + {Role: "user", StringContent: "Drone, go up", Content: "Drone, go up"}, + } + + handled := classifierRespond(context.Background(), session, conv, t, r, cc, history, nil, 0) + + Expect(handled).To(BeTrue()) + Expect(m.classifyCalls).To(Equal(1)) + results := classifierResultEvents(t) + Expect(results).To(HaveLen(1)) + Expect(results[0].ChosenID).To(Equal("up")) + }) + + It("speaks the address reply for unaddressed turns in reply mode", func() { + m := &fakeModel{} + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp-unaddressed-reply"} + cc := classifierTestConfig(0.35, nil) + cc.Address = &types.ClassifierAddress{Names: []string{"drone"}, Mode: types.ClassifierAddressReply, Reply: "Call me Drone."} + history := schema.Messages{ + {Role: "user", StringContent: "go up", Content: "go up"}, + } + + handled := classifierRespond(context.Background(), session, conv, t, r, cc, history, nil, 0) + + Expect(handled).To(BeTrue()) + Expect(m.classifyCalls).To(BeZero()) + Expect(t.countEvents(types.ServerEventTypeResponseOutputTextDone)).To(Equal(1)) + }) + + It("applies the fallback without scoring when the turn has no words", func() { + // A VAD-committed turn whose transcript is empty must not be + // scored: an empty prompt yields a confidently arbitrary winner. + m := &fakeModel{classifyScores: []router.LabelScore{{Label: "up", Score: 0.99}}} + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp-empty"} + history := schema.Messages{ + {Role: "system", StringContent: "instructions", Content: "instructions"}, + {Role: "user", StringContent: "", Content: ""}, + } + cc := classifierTestConfig(0.35, &types.ClassifierFallback{Mode: types.ClassifierFallbackReply, Reply: "Say again?"}) + + handled := classifierRespond(context.Background(), session, conv, t, r, cc, history, nil, 0) + + Expect(handled).To(BeTrue()) + Expect(m.classifyCalls).To(BeZero(), "an empty turn must not be scored") + results := classifierResultEvents(t) + Expect(results).To(HaveLen(1)) + Expect(results[0].Scores).To(BeEmpty()) + Expect(results[0].ChosenID).To(BeEmpty()) + Expect(results[0].Fallback).To(Equal(types.ClassifierFallbackReply)) + Expect(t.countEvents(types.ServerEventTypeResponseOutputTextDone)).To(Equal(1)) + Expect(t.countEvents(types.ServerEventTypeResponseFunctionCallArgumentsDone)).To(BeZero()) + }) + + It("falls through to generation for a word-less turn when the fallback is generate", func() { + m := &fakeModel{classifyScores: []router.LabelScore{{Label: "up", Score: 0.99}}} + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp-empty-gen"} + history := schema.Messages{ + {Role: "user", StringContent: "", Content: ""}, + } + cc := classifierTestConfig(0.35, &types.ClassifierFallback{Mode: types.ClassifierFallbackGenerate}) + + handled := classifierRespond(context.Background(), session, conv, t, r, cc, history, nil, 0) + + Expect(handled).To(BeFalse()) + Expect(m.classifyCalls).To(BeZero()) + Expect(classifierResultEvents(t)).To(BeEmpty()) + }) + + It("speaks the fallback reply when no option clears the threshold", func() { + m := &fakeModel{classifyScores: []router.LabelScore{ + {Label: "up", Score: 0.3}, + {Label: "greeting", Score: 0.3}, + }} + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp1"} + cc := classifierTestConfig(0.6, &types.ClassifierFallback{Mode: types.ClassifierFallbackReply, Reply: "Say again?"}) + + handled := classifierRespond(context.Background(), session, conv, t, r, cc, classifierTestHistory, nil, 0) + + Expect(handled).To(BeTrue()) + results := classifierResultEvents(t) + Expect(results).To(HaveLen(1)) + Expect(results[0].ChosenID).To(BeEmpty()) + Expect(results[0].Fallback).To(Equal(types.ClassifierFallbackReply)) + Expect(t.countEvents(types.ServerEventTypeResponseOutputTextDone)).To(Equal(1)) + Expect(t.countEvents(types.ServerEventTypeResponseFunctionCallArgumentsDone)).To(BeZero()) + }) + + It("completes with no output for the none fallback", func() { + m := &fakeModel{classifyScores: []router.LabelScore{ + {Label: "up", Score: 0.3}, + {Label: "greeting", Score: 0.3}, + }} + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp1"} + + handled := classifierRespond(context.Background(), session, conv, t, r, classifierTestConfig(0.6, nil), classifierTestHistory, nil, 0) + + Expect(handled).To(BeTrue()) + Expect(r.outcome).ToNot(Equal(outcomeFailed)) + Expect(conv.Items).To(BeEmpty()) + Expect(t.countEvents(types.ServerEventTypeResponseOutputTextDone)).To(BeZero()) + results := classifierResultEvents(t) + Expect(results).To(HaveLen(1)) + Expect(results[0].Fallback).To(Equal(types.ClassifierFallbackNone)) + }) + + It("falls through to generation for the generate fallback", func() { + m := &fakeModel{classifyScores: []router.LabelScore{ + {Label: "up", Score: 0.3}, + {Label: "greeting", Score: 0.3}, + }} + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp1"} + cc := classifierTestConfig(0.6, &types.ClassifierFallback{Mode: types.ClassifierFallbackGenerate}) + + handled := classifierRespond(context.Background(), session, conv, t, r, cc, classifierTestHistory, nil, 0) + + Expect(handled).To(BeFalse()) + // The distribution is still reported before falling through. + Expect(classifierResultEvents(t)).To(HaveLen(1)) + }) + + It("fails the response when scoring errors without a generate fallback", func() { + m := &fakeModel{classifyErr: fmt.Errorf("backend exploded")} + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp1"} + + handled := classifierRespond(context.Background(), session, conv, t, r, classifierTestConfig(0.35, nil), classifierTestHistory, nil, 0) + + Expect(handled).To(BeTrue()) + Expect(r.outcome).To(Equal(outcomeFailed)) + Expect(t.countEvents(types.ServerEventTypeError)).To(Equal(1)) + }) + + It("falls through to generation when scoring errors and fallback is generate", func() { + m := &fakeModel{classifyErr: fmt.Errorf("backend exploded")} + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp1"} + cc := classifierTestConfig(0.35, &types.ClassifierFallback{Mode: types.ClassifierFallbackGenerate}) + + handled := classifierRespond(context.Background(), session, conv, t, r, cc, classifierTestHistory, nil, 0) + + Expect(handled).To(BeFalse()) + Expect(r.outcome).ToNot(Equal(outcomeFailed)) + }) + + It("records a cancelled outcome when barge-in fires during scoring", func() { + m := &fakeModel{classifyScores: []router.LabelScore{ + {Label: "up", Score: 0.9}, + {Label: "greeting", Score: 0.1}, + }} + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp1"} + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + handled := classifierRespond(ctx, session, conv, t, r, classifierTestConfig(0.35, nil), classifierTestHistory, nil, 0) + + Expect(handled).To(BeTrue()) + Expect(r.outcome).To(Equal(outcomeCancelled)) + Expect(conv.Items).To(BeEmpty()) + }) + + It("skips to generation when there is nothing scorable", func() { + m := &fakeModel{} + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp1"} + systemOnly := schema.Messages{{Role: "system", StringContent: "instructions"}} + + handled := classifierRespond(context.Background(), session, conv, t, r, classifierTestConfig(0.35, nil), systemOnly, nil, 0) + + Expect(handled).To(BeFalse()) + Expect(m.classifyCalls).To(BeZero()) + }) +}) + +// slottedTestConfig is classifierTestConfig with the winning option's tool +// carrying argument slots (the hybrid classify-then-complete path). +func slottedTestConfig(threshold float64, fallback *types.ClassifierFallback, defaults bool) *types.ClassifierConfig { + slots := []types.ClassifierSlot{ + {Name: "distance", Type: types.ClassifierSlotNumber}, + {Name: "units", Type: types.ClassifierSlotEnum, Values: []string{"m", "meters", "ft", "feet"}, Hint: "assume m when the user gives no units"}, + } + if defaults { + slots[0].Default = "1" + slots[1].Default = "m" + } + return &types.ClassifierConfig{ + Threshold: threshold, + Fallback: fallback, + Options: []types.ClassifierOption{ + { + ID: "up", + Description: "the user asks the drone to fly up", + Reply: "Going up {{distance}} {{units}}.", + Tool: &types.ClassifierTool{ + Name: "move", + Arguments: json.RawMessage(`{"direction":"up","distance":"{{distance}}","units":"{{units}}"}`), + Slots: slots, + }, + }, + {ID: "greeting", Description: "the user greets the assistant", Reply: "Hello."}, + }, + } +} + +var _ = Describe("slotFillGrammar", func() { + It("pins the field skeleton and frees only the slot values", func() { + g := slotFillGrammar([]types.ClassifierSlot{ + {Name: "distance", Type: types.ClassifierSlotNumber}, + {Name: "units", Type: types.ClassifierSlotEnum, Values: []string{"m", "ft"}}, + }) + Expect(g).To(ContainSubstring(`root ::= slot0 ", \"units\": " slot1 "}"`)) + Expect(g).To(ContainSubstring("slot0 ::= num")) + Expect(g).To(ContainSubstring(`slot1 ::= "\"m\"" | "\"ft\""`)) + Expect(g).To(ContainSubstring("num ::=")) + }) + + It("JSON-encodes enum values before embedding them in the grammar", func() { + g := slotFillGrammar([]types.ClassifierSlot{ + {Name: "units", Type: types.ClassifierSlotEnum, Values: []string{"quoted\"value", "line\nbreak", `back\slash`}}, + }) + Expect(g).To(ContainSubstring(gbnfLiteral(`"quoted\"value"`))) + Expect(g).To(ContainSubstring(gbnfLiteral(`"line\nbreak"`))) + Expect(g).To(ContainSubstring(gbnfLiteral(`"back\\slash"`))) + }) + + It("budgets forced enum and field text by encoded length", func() { + short := []types.ClassifierSlot{{Name: "value", Type: types.ClassifierSlotEnum, Values: []string{"m"}}} + long := []types.ClassifierSlot{ + {Name: "value", Type: types.ClassifierSlotEnum, Values: []string{strings.Repeat("long-value-", 20)}}, + {Name: strings.Repeat("field", 20), Type: types.ClassifierSlotNumber}, + } + Expect(slotFillMaxTokens(long)).To(BeNumerically(">", slotFillMaxTokens(short)+200)) + }) + + It("emits a string rule only when needed", func() { + g := slotFillGrammar([]types.ClassifierSlot{{Name: "what", Type: types.ClassifierSlotString}}) + Expect(g).To(ContainSubstring("slot0 ::= str")) + Expect(g).To(ContainSubstring("str ::=")) + Expect(g).ToNot(ContainSubstring("num ::=")) + }) +}) + +var _ = Describe("parseSlotValues", func() { + slots := []types.ClassifierSlot{ + {Name: "distance", Type: types.ClassifierSlotNumber}, + {Name: "units", Type: types.ClassifierSlotEnum, Values: []string{"m", "ft"}}, + } + + It("extracts values from a grammar-shaped completion", func() { + values, err := parseSlotValues("up", "distance", `3.5, "units": "m"}`, slots) + Expect(err).ToNot(HaveOccurred()) + Expect(values).To(Equal(map[string]string{"distance": "3.5", "units": "m"})) + }) + + It("tolerates a completion missing the closing brace", func() { + values, err := parseSlotValues("up", "distance", `2, "units": "ft"`, slots) + Expect(err).ToNot(HaveOccurred()) + Expect(values["distance"]).To(Equal("2")) + }) + + It("rejects completions missing a slot", func() { + _, err := parseSlotValues("up", "distance", `3}`, slots) + Expect(err).To(MatchError(ContainSubstring(`missing "units"`))) + }) +}) + +var _ = Describe("classifierPolicyDescription", func() { + It("passes plain options through", func() { + o := &types.ClassifierOption{Description: "plain"} + Expect(classifierPolicyDescription(o)).To(Equal("plain")) + }) + + It("appends slot declarations and hints", func() { + cc := slottedTestConfig(0, nil, false) + d := classifierPolicyDescription(&cc.Options[0]) + Expect(d).To(ContainSubstring("route parameters:")) + Expect(d).To(ContainSubstring("distance (number)")) + Expect(d).To(ContainSubstring("units (one of: m, meters, ft, feet)")) + Expect(d).To(ContainSubstring("assume m when the user gives no units")) + }) +}) + +var _ = Describe("classifierRespond slot filling", func() { + It("emits the filled tool arguments and reports them in the result event", func() { + m := &fakeModel{ + classifyScores: []router.LabelScore{{Label: "up", Score: 0.9}, {Label: "greeting", Score: 0.1}}, + fillArgs: `{"direction":"up","distance":3,"units":"meters"}`, + fillValues: map[string]string{"distance": "3", "units": "meters"}, + } + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp-slots"} + + handled := classifierRespond(context.Background(), session, conv, t, r, slottedTestConfig(0.35, nil, false), classifierTestHistory, nil, 0) + + Expect(handled).To(BeTrue()) + Expect(m.fillCalls).To(Equal(1)) + Expect(m.lastFillChosen.ID).To(Equal("up")) + + results := classifierResultEvents(t) + Expect(results).To(HaveLen(1)) + Expect(results[0].ChosenID).To(Equal("up")) + Expect(results[0].Arguments).To(MatchJSON(`{"direction":"up","distance":3,"units":"meters"}`)) + + var fcArgs string + for _, e := range t.events { + if done, ok := e.(types.ResponseFunctionCallArgumentsDoneEvent); ok { + fcArgs = done.Arguments + } + } + Expect(fcArgs).To(MatchJSON(`{"direction":"up","distance":3,"units":"meters"}`)) + }) + + It("splices the filled values into a templated reply", func() { + m := &fakeModel{ + classifyScores: []router.LabelScore{{Label: "up", Score: 0.9}, {Label: "greeting", Score: 0.1}}, + fillArgs: `{"direction":"up","distance":3,"units":"meters"}`, + fillValues: map[string]string{"distance": "3", "units": "meters"}, + } + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp-slot-reply"} + + handled := classifierRespond(context.Background(), session, conv, t, r, slottedTestConfig(0.35, nil, false), classifierTestHistory, nil, 0) + + Expect(handled).To(BeTrue()) + Expect(replyTexts(t)).To(ConsistOf("Going up 3 meters.")) + }) + + It("recovers with slot defaults when filling fails", func() { + m := &fakeModel{ + classifyScores: []router.LabelScore{{Label: "up", Score: 0.9}, {Label: "greeting", Score: 0.1}}, + fillErr: fmt.Errorf("backend unavailable"), + } + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp-slot-defaults"} + + handled := classifierRespond(context.Background(), session, conv, t, r, slottedTestConfig(0.35, nil, true), classifierTestHistory, nil, 0) + + Expect(handled).To(BeTrue()) + var fcArgs string + for _, e := range t.events { + if done, ok := e.(types.ResponseFunctionCallArgumentsDoneEvent); ok { + fcArgs = done.Arguments + } + } + Expect(fcArgs).To(MatchJSON(`{"direction":"up","distance":1,"units":"m"}`)) + Expect(replyTexts(t)).To(ConsistOf("Going up 1 m."), "the default-recovery reply confirms the defaults") + }) + + It("fails the response when filling fails and a slot has no default", func() { + m := &fakeModel{ + classifyScores: []router.LabelScore{{Label: "up", Score: 0.9}, {Label: "greeting", Score: 0.1}}, + fillErr: fmt.Errorf("backend unavailable"), + } + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp-slot-fail"} + + handled := classifierRespond(context.Background(), session, conv, t, r, slottedTestConfig(0.35, nil, false), classifierTestHistory, nil, 0) + + Expect(handled).To(BeTrue()) + Expect(r.outcome).To(Equal(outcomeFailed)) + Expect(classifierResultEvents(t)).To(BeEmpty(), "no result event for a failed fill") + }) + + It("falls back to generation on fill failure in generate mode", func() { + m := &fakeModel{ + classifyScores: []router.LabelScore{{Label: "up", Score: 0.9}, {Label: "greeting", Score: 0.1}}, + fillErr: fmt.Errorf("backend unavailable"), + } + session := classifierTestSession(m) + conv := &Conversation{} + t := &fakeTransport{} + r := &liveResponse{id: "resp-slot-genfb"} + cc := slottedTestConfig(0.35, &types.ClassifierFallback{Mode: types.ClassifierFallbackGenerate}, false) + + handled := classifierRespond(context.Background(), session, conv, t, r, cc, classifierTestHistory, nil, 0) + + Expect(handled).To(BeFalse(), "generate fallback lets the caller run generation") + }) +}) diff --git a/core/http/endpoints/openai/realtime_doubles_test.go b/core/http/endpoints/openai/realtime_doubles_test.go index fe52e1c64..a2c104b3c 100644 --- a/core/http/endpoints/openai/realtime_doubles_test.go +++ b/core/http/endpoints/openai/realtime_doubles_test.go @@ -3,11 +3,13 @@ package openai import ( "context" "strings" + "sync" "github.com/mudler/LocalAI/core/backend" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/http/endpoints/openai/types" "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/services/routing/router" "github.com/mudler/LocalAI/pkg/grpc/proto" ) @@ -99,11 +101,81 @@ type fakeModel struct { predictResp backend.LLMResponse predictErr error + // ClassifyTurn scripting: classifyScores is returned as the option + // distribution (in option order); classifyErr fails the call. + // classifyCalls counts invocations and lastClassifyOptions records + // what the handler asked to score. + classifyScores []router.LabelScore + classifyErr error + classifyCalls int + lastClassifyOptions []types.ClassifierOption + + // FillToolArguments scripting: fillArgs/fillValues are returned + // verbatim; fillErr fails the call. fillCalls counts invocations and + // lastFillChosen records which option's slots the handler asked to + // fill. + fillArgs string + fillValues map[string]string + fillErr error + fillCalls int + lastFillChosen *types.ClassifierOption + + // PrewarmClassifier runs on a background goroutine, so its recording + // is mutex-guarded; specs poll prewarmCalls with Eventually. + prewarmMu sync.Mutex + prewarmCalls int + lastPrewarmOptions []types.ClassifierOption + + // VAD scripting: vadFn, when set, decides per call (specs vary the + // answer across ticks or record the request); otherwise + // vadSegments/vadErr answer every call. + vadFn func(*schema.VADRequest) (*schema.VADResponse, error) + vadSegments []schema.VADSegment + vadErr error + lastMessages schema.Messages } -func (m *fakeModel) VAD(context.Context, *schema.VADRequest) (*schema.VADResponse, error) { - return nil, nil +func (m *fakeModel) PrewarmClassifier(_ context.Context, options []types.ClassifierOption, _ string) { + m.prewarmMu.Lock() + defer m.prewarmMu.Unlock() + m.prewarmCalls++ + m.lastPrewarmOptions = options +} + +func (m *fakeModel) prewarmed() (int, []types.ClassifierOption) { + m.prewarmMu.Lock() + defer m.prewarmMu.Unlock() + return m.prewarmCalls, m.lastPrewarmOptions +} + +func (m *fakeModel) FillToolArguments(_ context.Context, msgs schema.Messages, options []types.ClassifierOption, _ string, chosen *types.ClassifierOption) (string, map[string]string, error) { + m.fillCalls++ + m.lastFillChosen = chosen + if m.fillErr != nil { + return "", nil, m.fillErr + } + return m.fillArgs, m.fillValues, nil +} + +func (m *fakeModel) ClassifyTurn(_ context.Context, msgs schema.Messages, options []types.ClassifierOption, _ string) ([]router.LabelScore, error) { + m.classifyCalls++ + m.lastClassifyOptions = options + m.lastMessages = msgs + if m.classifyErr != nil { + return nil, m.classifyErr + } + return m.classifyScores, nil +} + +func (m *fakeModel) VAD(_ context.Context, req *schema.VADRequest) (*schema.VADResponse, error) { + if m.vadFn != nil { + return m.vadFn(req) + } + if m.vadErr != nil { + return nil, m.vadErr + } + return &schema.VADResponse{Segments: m.vadSegments}, nil } func (m *fakeModel) Transcribe(context.Context, string, string, bool, bool, string) (*schema.TranscriptionResult, error) { diff --git a/core/http/endpoints/openai/realtime_model.go b/core/http/endpoints/openai/realtime_model.go index 0449daee3..030a0c914 100644 --- a/core/http/endpoints/openai/realtime_model.go +++ b/core/http/endpoints/openai/realtime_model.go @@ -7,6 +7,9 @@ import ( "encoding/hex" "encoding/json" "fmt" + "strings" + "sync" + "time" "github.com/mudler/LocalAI/core/application" "github.com/mudler/LocalAI/core/backend" @@ -36,12 +39,39 @@ type wrappedModel struct { LLMConfig *config.ModelConfig VADConfig *config.ModelConfig SoundDetectionConfig *config.ModelConfig + // ScoreConfig is the classifier-mode scoring model + // (pipeline.classifier.model). nil falls back to LLMConfig — with + // slot-based Score the same process serves scoring and generation + // and shares its prompt cache between them. + ScoreConfig *config.ModelConfig appConfig *config.ApplicationConfig modelLoader *model.ModelLoader confLoader *config.ModelConfigLoader evaluator *templates.Evaluator + // Classifier-mode memo: constructing a ScoreClassifier parses the + // scoring model's chat template, so reuse it while the option set is + // unchanged. Guarded by a mutex only because session.update can swap + // options while a response is in flight. + classifierMu sync.Mutex + classifier *router.ScoreClassifier + classifierKey string + classifierWarn sync.Once + // Prewarm FIFO: a single worker drains warms in registration order — + // a plain mutex proved unfair under a burst of registrations (Go + // mutexes barge), running the most recently registered list last, + // long after the user's first command for it arrived. Pending + // duplicates coalesce (a connect-time barrage registers the same + // list several times), but completed warms are deliberately NOT + // memoized: a rewarm on a still-resident list costs one probe-sized + // decode, and on an evicted list it is exactly the re-prefill the + // next turn would otherwise pay in the foreground. + prewarmMu sync.Mutex + prewarmQueue []prewarmJob + prewarmPending map[string]bool + prewarmActive bool + // Routing — populated by newModel when the application wires routing // deps in. nil-safe: with classifierRegistry == nil the per-turn // routing block in Predict is skipped, preserving today's "one LLM @@ -90,6 +120,17 @@ func (m *transcriptOnlyModel) Predict(ctx context.Context, messages schema.Messa return nil, fmt.Errorf("predict operation not supported in transcript-only mode") } +func (m *transcriptOnlyModel) ClassifyTurn(ctx context.Context, messages schema.Messages, options []types.ClassifierOption, normalization string) ([]router.LabelScore, error) { + return nil, fmt.Errorf("classifier mode not supported in transcript-only mode") +} + +func (m *transcriptOnlyModel) FillToolArguments(ctx context.Context, messages schema.Messages, options []types.ClassifierOption, normalization string, chosen *types.ClassifierOption) (string, map[string]string, error) { + return "", nil, fmt.Errorf("classifier mode not supported in transcript-only mode") +} + +func (m *transcriptOnlyModel) PrewarmClassifier(ctx context.Context, options []types.ClassifierOption, normalization string) { +} + func (m *transcriptOnlyModel) TTS(ctx context.Context, text, voice, language string) (string, *proto.Result, error) { return "", nil, fmt.Errorf("TTS not supported in transcript-only mode") } @@ -369,14 +410,258 @@ func (m *wrappedModel) PredictConfig() *config.ModelConfig { return m.LLMConfig } +// scoreConfig resolves the classifier-mode scoring model: the explicit +// pipeline.classifier.model when set, else the pipeline LLM. +func (m *wrappedModel) scoreConfig() *config.ModelConfig { + if m.ScoreConfig != nil { + return m.ScoreConfig + } + return m.LLMConfig +} + +// classifierFor returns a ScoreClassifier for the given option set, +// reusing the previous one while options and normalization are unchanged +// (construction parses the scoring model's chat template). +func (m *wrappedModel) classifierFor(options []types.ClassifierOption, normalization string) (*router.ScoreClassifier, error) { + scoreCfg := m.scoreConfig() + if scoreCfg == nil || !scoreCfg.HasUsecases(config.FLAG_SCORE) { + return nil, fmt.Errorf("classifier: scoring model must include score in known_usecases") + } + switch normalization { + case "", router.ScoreNormalizationRaw, router.ScoreNormalizationMean: + default: + // NewScoreClassifier panics on unknown modes; session.update + // validation should have rejected this — fail soft anyway. + return nil, fmt.Errorf("classifier: unknown normalization %q", normalization) + } + if len(options) == 0 { + return nil, fmt.Errorf("classifier: no options to score") + } + + var key strings.Builder + key.WriteString(normalization) + for _, o := range options { + key.WriteString("\x1f") + key.WriteString(o.ID) + key.WriteString("\x1e") + // The policy description includes slot declarations, so keying on + // it also invalidates the classifier when slots change. + key.WriteString(classifierPolicyDescription(&o)) + } + + m.classifierMu.Lock() + defer m.classifierMu.Unlock() + if m.classifier != nil && m.classifierKey == key.String() { + return m.classifier, nil + } + + cfg := m.scoreConfig() + policies := make([]router.ScorePolicy, 0, len(options)) + for _, o := range options { + if o.ID == "" || o.Description == "" { + // NewScoreClassifier panics on these; validation upstream + // should have caught them. + return nil, fmt.Errorf("classifier: option with empty id or description") + } + policies = append(policies, router.ScorePolicy{Label: o.ID, Description: classifierPolicyDescription(&o)}) + } + + opts := router.ScoreClassifierOptions{ + // The memo cache stores only label sets — a hit would return an + // empty distribution and blind the localai.classifier.result + // event, so keep it off. + CacheCap: 0, + Normalization: normalization, + } + if m.routerDeps != nil && m.routerDeps.TokenCounter != nil && cfg.ContextSize != nil { + opts.TokenCounter = m.routerDeps.TokenCounter(cfg.Name) + opts.MaxContextTokens = *cfg.ContextSize + } + for i := range options { + if options[i].Tool != nil && len(options[i].Tool.Slots) > 0 { + reserve := slotFillContextReserve(&options[i]) + if reserve > opts.CompletionReserveTokens { + opts.CompletionReserveTokens = reserve + } + } + } + if m.evaluator != nil { + if renderer := middleware.NewTemplateRenderer(m.evaluator, cfg); renderer != nil { + opts.PromptRenderer = renderer + } else { + m.classifierWarn.Do(func() { + xlog.Warn("realtime classifier: scoring model has no Go chat template; falling back to a generic ChatML envelope, which may be off-distribution", + "model", cfg.Name) + }) + } + } + if st := middleware.PickAssistantTurnEnd(cfg.StopWords, cfg.TemplateConfig.ChatMessage); st != "" { + opts.StopToken = st + } + + scorer := backend.NewScorer(m.modelLoader, *cfg, m.appConfig) + m.classifier = router.NewScoreClassifier(policies, scorer, opts) + m.classifierKey = key.String() + return m.classifier, nil +} + +// PrewarmClassifier primes the scoring backend's prompt cache for a newly +// registered option list so the first real turns don't pay the prefill. +// One throwaway score prefills the new option-list prompt and declares the +// per-turn probe boundary, leaving the backend a rewind point (a KV +// checkpoint on hybrid/recurrent models, which cannot rewind arbitrarily) +// at the stable prefix every subsequent turn reuses. +// Best-effort: errors are logged, never surfaced. +func (m *wrappedModel) PrewarmClassifier(ctx context.Context, options []types.ClassifierOption, normalization string) { + classifier, err := m.classifierFor(options, normalization) + if err != nil { + xlog.Debug("realtime classifier: prewarm skipped", "error", err) + return + } + m.classifierMu.Lock() + key := m.classifierKey + m.classifierMu.Unlock() + + m.prewarmMu.Lock() + defer m.prewarmMu.Unlock() + if m.prewarmPending == nil { + m.prewarmPending = make(map[string]bool) + } + if m.prewarmPending[key] { + return + } + m.prewarmPending[key] = true + m.prewarmQueue = append(m.prewarmQueue, prewarmJob{classifier: classifier, key: key, options: len(options)}) + if !m.prewarmActive { + m.prewarmActive = true + go m.prewarmWorker() + } +} + +type prewarmJob struct { + classifier *router.ScoreClassifier + key string + options int +} + +// prewarmWorker drains queued warms one at a time, in order. One +// throwaway score per list is enough: the scoring call itself plants the +// backend's reuse point at the stable-prefix boundary it declares, so +// the real turns that follow restore from it no matter how their probe +// differs. The worker exits when the queue drains and restarts on the +// next registration. +func (m *wrappedModel) prewarmWorker() { + for { + m.prewarmMu.Lock() + if len(m.prewarmQueue) == 0 { + m.prewarmActive = false + m.prewarmMu.Unlock() + return + } + job := m.prewarmQueue[0] + m.prewarmQueue = m.prewarmQueue[1:] + m.prewarmMu.Unlock() + + start := time.Now() + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + const probe = "warmup" + _, err := job.classifier.Classify(ctx, router.Probe{Prompt: probe, Messages: []string{probe}}) + cancel() + if err != nil { + xlog.Warn("realtime classifier: prewarm scoring failed", "error", err) + } else { + xlog.Debug("realtime classifier: prewarmed scoring prompt cache", + "options", job.options, "latency_ms", time.Since(start).Milliseconds()) + } + m.prewarmMu.Lock() + delete(m.prewarmPending, job.key) + m.prewarmMu.Unlock() + } +} + +func (m *wrappedModel) ClassifyTurn(ctx context.Context, messages schema.Messages, options []types.ClassifierOption, normalization string) ([]router.LabelScore, error) { + classifier, err := m.classifierFor(options, normalization) + if err != nil { + return nil, err + } + decision, err := classifier.Classify(ctx, classifierProbe(messages)) + if err != nil { + return nil, err + } + // LabelScores is in policy-declaration order, which mirrors option + // order by construction. + if len(decision.LabelScores) != len(options) { + return nil, fmt.Errorf("classifier: got %d scores for %d options", len(decision.LabelScores), len(options)) + } + return decision.LabelScores, nil +} + +// FillToolArguments runs the hybrid slot-fill completion: the exact prompt +// the classifier scored (rendered by the same, cached ScoreClassifier — so +// the backend's prompt cache is warm) continued by the chosen route JSON +// re-opened at its first slot, with a grammar pinning everything but the +// slot values. Deterministic (temperature 0), a couple dozen tokens at +// most. +func (m *wrappedModel) FillToolArguments(ctx context.Context, messages schema.Messages, options []types.ClassifierOption, normalization string, chosen *types.ClassifierOption) (string, map[string]string, error) { + if chosen == nil || chosen.Tool == nil || len(chosen.Tool.Slots) == 0 { + return "", nil, fmt.Errorf("classifier: option has no slots to fill") + } + slots := chosen.Tool.Slots + classifier, err := m.classifierFor(options, normalization) + if err != nil { + return "", nil, err + } + prompt, err := classifier.SlotFillPrompt(classifierProbe(messages), chosen.ID, slots[0].Name) + if err != nil { + return "", nil, err + } + + // The scoring config, narrowed to a deterministic constrained + // completion. The completion usecase must be declared alongside score + // — bootstrap-style configs use known_usecases: [chat, completion, + // score]. + cfg := *m.scoreConfig() + if !cfg.HasUsecases(config.FLAG_COMPLETION) { + return "", nil, fmt.Errorf("classifier: slot filling requires completion in the scoring model's known_usecases") + } + cfg.Grammar = slotFillGrammar(slots) + maxTokens := slotFillMaxTokens(slots) + temperature := 0.0 + cfg.Maxtokens = &maxTokens + cfg.Temperature = &temperature + + fn, err := backend.ModelInference(ctx, prompt, nil, nil, nil, nil, m.modelLoader, &cfg, m.confLoader, m.appConfig, nil, "", "", nil, nil, nil, nil) + if err != nil { + return "", nil, fmt.Errorf("classifier: slot fill inference: %w", err) + } + resp, err := fn() + if err != nil { + return "", nil, fmt.Errorf("classifier: slot fill inference: %w", err) + } + values, err := parseSlotValues(chosen.ID, slots[0].Name, resp.Response, slots) + if err != nil { + return "", nil, err + } + args, err := chosen.Tool.SpliceArguments(values) + if err != nil { + return "", nil, err + } + return args, values, nil +} + func (m *wrappedModel) Warmup(ctx context.Context) error { - _, err := backend.PreloadStages(ctx, m.modelLoader, m.appConfig, []backend.PreloadStage{ + stages := []backend.PreloadStage{ {Role: "vad", Cfg: m.VADConfig}, {Role: "transcription", Cfg: m.TranscriptionConfig}, {Role: "llm", Cfg: m.LLMConfig}, {Role: "tts", Cfg: m.TTSConfig}, {Role: "sound_detection", Cfg: m.SoundDetectionConfig}, - }) + } + // The scoring model is a separate stage only when it isn't the LLM. + if m.ScoreConfig != nil && m.ScoreConfig != m.LLMConfig { + stages = append(stages, backend.PreloadStage{Role: "classifier", Cfg: m.ScoreConfig}) + } + _, err := backend.PreloadStages(ctx, m.modelLoader, m.appConfig, stages) return err } @@ -456,11 +741,11 @@ func modelSoundDetection(ctx context.Context, ml *model.ModelLoader, appConfig * // config named by pipeline.sound_detection. Returns (nil, nil) when no model // is configured so sound detection stays additive and never blocks session // setup. -func loadSoundDetectionConfig(pipeline *config.Pipeline, cl *config.ModelConfigLoader, ml *model.ModelLoader) (*config.ModelConfig, error) { +func loadSoundDetectionConfig(pipeline *config.Pipeline, cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) (*config.ModelConfig, error) { if pipeline.SoundDetection == "" { return nil, nil } - cfg, err := cl.LoadResolvedModelConfig(pipeline.SoundDetection, ml.ModelPath) + cfg, err := cl.LoadResolvedModelConfig(pipeline.SoundDetection, ml.ModelPath, appConfig.ToConfigLoaderOptions()...) if err != nil { return nil, fmt.Errorf("failed to load sound detection config: %w", err) } @@ -471,7 +756,7 @@ func loadSoundDetectionConfig(pipeline *config.Pipeline, cl *config.ModelConfigL } func newTranscriptionOnlyModel(pipeline *config.Pipeline, cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) (Model, *config.ModelConfig, error) { - cfgVAD, err := cl.LoadResolvedModelConfig(pipeline.VAD, ml.ModelPath) + cfgVAD, err := cl.LoadResolvedModelConfig(pipeline.VAD, ml.ModelPath, appConfig.ToConfigLoaderOptions()...) if err != nil { return nil, nil, fmt.Errorf("failed to load backend config: %w", err) @@ -481,7 +766,7 @@ func newTranscriptionOnlyModel(pipeline *config.Pipeline, cl *config.ModelConfig return nil, nil, fmt.Errorf("failed to validate config: %w", err) } - cfgSST, err := cl.LoadResolvedModelConfig(pipeline.Transcription, ml.ModelPath) + cfgSST, err := cl.LoadResolvedModelConfig(pipeline.Transcription, ml.ModelPath, appConfig.ToConfigLoaderOptions()...) if err != nil { return nil, nil, fmt.Errorf("failed to load backend config: %w", err) @@ -491,7 +776,7 @@ func newTranscriptionOnlyModel(pipeline *config.Pipeline, cl *config.ModelConfig return nil, nil, fmt.Errorf("failed to validate config: %w", err) } - cfgSound, err := loadSoundDetectionConfig(pipeline, cl, ml) + cfgSound, err := loadSoundDetectionConfig(pipeline, cl, ml, appConfig) if err != nil { return nil, nil, err } @@ -513,7 +798,7 @@ func newTranscriptionOnlyModel(pipeline *config.Pipeline, cl *config.ModelConfig // speech) and is driven by client-side windowing (turn_detection none + // input_audio_buffer.commit) rather than the voice VAD loop. func newSoundDetectionOnlyModel(pipeline *config.Pipeline, cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) (Model, error) { - cfgSound, err := loadSoundDetectionConfig(pipeline, cl, ml) + cfgSound, err := loadSoundDetectionConfig(pipeline, cl, ml, appConfig) if err != nil { return nil, err } @@ -574,7 +859,7 @@ func buildRealtimeRoutingContext(a *application.Application, sessionID string) * func newModel(pipeline *config.Pipeline, cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig, evaluator *templates.Evaluator, routing *RealtimeRoutingContext) (Model, error) { xlog.Debug("Creating new model pipeline model", "pipeline", pipeline) - cfgVAD, err := cl.LoadResolvedModelConfig(pipeline.VAD, ml.ModelPath) + cfgVAD, err := cl.LoadResolvedModelConfig(pipeline.VAD, ml.ModelPath, appConfig.ToConfigLoaderOptions()...) if err != nil { return nil, fmt.Errorf("failed to load backend config: %w", err) @@ -585,7 +870,7 @@ func newModel(pipeline *config.Pipeline, cl *config.ModelConfigLoader, ml *model } // TODO: Do we always need a transcription model? It can be disabled. Note that any-to-any instruction following models don't transcribe as such, so if transcription is required it is a separate process - cfgSST, err := cl.LoadResolvedModelConfig(pipeline.Transcription, ml.ModelPath) + cfgSST, err := cl.LoadResolvedModelConfig(pipeline.Transcription, ml.ModelPath, appConfig.ToConfigLoaderOptions()...) if err != nil { return nil, fmt.Errorf("failed to load backend config: %w", err) @@ -617,7 +902,7 @@ func newModel(pipeline *config.Pipeline, cl *config.ModelConfigLoader, ml *model xlog.Debug("Loading a wrapped model") // Otherwise we want to return a wrapped model, which is a "virtual" model that re-uses other models to perform operations - cfgLLM, err := cl.LoadResolvedModelConfig(pipeline.LLM, ml.ModelPath) + cfgLLM, err := cl.LoadResolvedModelConfig(pipeline.LLM, ml.ModelPath, appConfig.ToConfigLoaderOptions()...) if err != nil { return nil, fmt.Errorf("failed to load backend config: %w", err) @@ -632,7 +917,7 @@ func newModel(pipeline *config.Pipeline, cl *config.ModelConfigLoader, ml *model applyPipelineReasoning(cfgLLM, *pipeline) applyPipelineThinking(cfgLLM, *pipeline) - cfgTTS, err := cl.LoadResolvedModelConfig(pipeline.TTS, ml.ModelPath) + cfgTTS, err := cl.LoadResolvedModelConfig(pipeline.TTS, ml.ModelPath, appConfig.ToConfigLoaderOptions()...) if err != nil { return nil, fmt.Errorf("failed to load backend config: %w", err) @@ -642,17 +927,51 @@ func newModel(pipeline *config.Pipeline, cl *config.ModelConfigLoader, ml *model return nil, fmt.Errorf("failed to validate config: %w", err) } - cfgSound, err := loadSoundDetectionConfig(pipeline, cl, ml) + cfgSound, err := loadSoundDetectionConfig(pipeline, cl, ml, appConfig) if err != nil { return nil, err } + // Classifier mode scores on its own model config when one is named; + // otherwise ClassifyTurn falls back to the LLM config at call time + // (so a client can enable classification via session.update even + // when the pipeline block is absent). + var cfgScore *config.ModelConfig + if pipeline.Classifier != nil && pipeline.Classifier.Model != "" { + cfgScore, err = cl.LoadResolvedModelConfig(pipeline.Classifier.Model, ml.ModelPath, appConfig.ToConfigLoaderOptions()...) + if err != nil { + return nil, fmt.Errorf("failed to load classifier scoring config: %w", err) + } + if valid, err := cfgScore.Validate(); !valid { + return nil, fmt.Errorf("failed to validate classifier scoring config: %w", err) + } + if !cfgScore.HasUsecases(config.FLAG_SCORE) { + return nil, fmt.Errorf("pipeline classifier: scoring model %q must declare known_usecases: [score]", cfgScore.Name) + } + } + if pipeline.Classifier != nil && pipeline.Classifier.Enabled { + effectiveScore := cfgScore + if effectiveScore == nil { + effectiveScore = cfgLLM + } + if effectiveScore.HasRouter() { + // A router model has no concrete backend to score on — the + // per-turn routing decision happens at Predict time, after + // classification would already have run. + return nil, fmt.Errorf("pipeline classifier: llm %q is a router model; set pipeline.classifier.model to a concrete scoring model", cfgLLM.Name) + } + if !effectiveScore.HasUsecases(config.FLAG_SCORE) { + return nil, fmt.Errorf("pipeline classifier: scoring model %q must declare known_usecases: [score]", effectiveScore.Name) + } + } + wm := &wrappedModel{ TTSConfig: cfgTTS, TranscriptionConfig: cfgSST, LLMConfig: cfgLLM, VADConfig: cfgVAD, SoundDetectionConfig: cfgSound, + ScoreConfig: cfgScore, confLoader: cl, modelLoader: ml, diff --git a/core/http/endpoints/openai/realtime_semantic_vad.go b/core/http/endpoints/openai/realtime_semantic_vad.go index 66dfc6efe..75a71ba25 100644 --- a/core/http/endpoints/openai/realtime_semantic_vad.go +++ b/core/http/endpoints/openai/realtime_semantic_vad.go @@ -96,6 +96,17 @@ func newLiveTurnState(session *Session, transport Transport) *liveTurnState { func (l *liveTurnState) open() bool { return l.live != nil } +// rebase shifts the turn's buffer-relative cursors after the retention trim +// dropped trimmedSec seconds off the buffer head: fed16k indexes the +// resampled (16 kHz) buffer, eouAtSec the buffer clock. Both floor at zero — +// a position inside the dropped head is more than maxTurnBufferSec old, and +// for eouAtSec zero already means "no EOU this turn", which is the right +// reading for a token that stale. +func (l *liveTurnState) rebase(trimmedSec float64) { + l.fed16k = max(0, l.fed16k-int(trimmedSec*localSampleRate)) + l.eouAtSec = max(0, l.eouAtSec-trimmedSec) +} + // openTurn starts the turn's live stream under the caller-supplied item id. A // failure (most commonly the backend's typed "live transcription unsupported" // signal) degrades the whole session to silence-only detection — warned once, diff --git a/core/http/endpoints/openai/realtime_stream_test.go b/core/http/endpoints/openai/realtime_stream_test.go index 439f3240e..2d5d7d7a1 100644 --- a/core/http/endpoints/openai/realtime_stream_test.go +++ b/core/http/endpoints/openai/realtime_stream_test.go @@ -263,4 +263,64 @@ var _ = Describe("triggerResponse", func() { Expect(done.Response.Usage.OutputTokens).To(Equal(3)) Expect(done.Response.Usage.TotalTokens).To(Equal(8)) }) + + // response.metadata is the only thing tying a terminal event back to the + // response.create that asked for it. Without the echo, a client running an + // out-of-band response alongside the spoken conversation cannot tell its own + // answer from the conversation's, and blocks until it times out. + It("echoes response.create metadata back on response.created and response.done", func() { + m := &fakeModel{ + cfg: &config.ModelConfig{}, + predictResp: backend.LLMResponse{Response: "Hi there."}, + } + session := &Session{ + OutputSampleRate: 24000, + ModelInterface: m, + ModelConfig: &config.ModelConfig{}, + OutputModalities: []types.Modality{types.ModalityText}, + } + t := &fakeTransport{} + + triggerResponse(context.Background(), session, &Conversation{}, t, &types.ResponseCreateParams{ + Metadata: map[string]string{"client_run": "abc123"}, + }) + + var created *types.ResponseCreatedEvent + var done *types.ResponseDoneEvent + for i := range t.events { + switch e := t.events[i].(type) { + case types.ResponseCreatedEvent: + created = &e + case types.ResponseDoneEvent: + done = &e + } + } + Expect(created).NotTo(BeNil()) + Expect(created.Response.Metadata).To(HaveKeyWithValue("client_run", "abc123")) + Expect(done).NotTo(BeNil()) + Expect(done.Response.Metadata).To(HaveKeyWithValue("client_run", "abc123")) + }) + + // Omitted rather than sent as an empty object, matching the omitempty tag. + It("sends no metadata when response.create carried none", func() { + m := &fakeModel{ + cfg: &config.ModelConfig{}, + predictResp: backend.LLMResponse{Response: "Hi there."}, + } + session := &Session{ + OutputSampleRate: 24000, + ModelInterface: m, + ModelConfig: &config.ModelConfig{}, + OutputModalities: []types.Modality{types.ModalityText}, + } + t := &fakeTransport{} + + triggerResponse(context.Background(), session, &Conversation{}, t, nil) + + for i := range t.events { + if d, ok := t.events[i].(types.ResponseDoneEvent); ok { + Expect(d.Response.Metadata).To(BeEmpty()) + } + } + }) }) diff --git a/core/http/endpoints/openai/realtime_turncoord.go b/core/http/endpoints/openai/realtime_turncoord.go index 30ffffc66..f0d599f7e 100644 --- a/core/http/endpoints/openai/realtime_turncoord.go +++ b/core/http/endpoints/openai/realtime_turncoord.go @@ -58,6 +58,14 @@ type turnSink struct { commitAudioLength float64 // for finishTurn (flush tail) commitRetranscribe bool // gated batch is authoritative commitGated *schema.TranscriptionResult // retranscribe batch decode + + // lastSpeechEndSec is where speech last ended this turn, in whole-buffer + // seconds (audioLength while the newest segment is still open). It + // outlives the segments scrolling out of the VAD scan clip, so the + // silence-outran-the-window commit still has a speech end to report. + // Zeroed whenever the turn leaves Speaking; rebased by the retention + // trim. + lastSpeechEndSec float64 } func newTurnSink(session *Session, conv *Conversation, t Transport, lts *liveTurnState, vadContext context.Context, startTime time.Time) *turnSink { diff --git a/core/http/endpoints/openai/realtime_vad_tick_test.go b/core/http/endpoints/openai/realtime_vad_tick_test.go new file mode 100644 index 000000000..43f26731c --- /dev/null +++ b/core/http/endpoints/openai/realtime_vad_tick_test.go @@ -0,0 +1,203 @@ +package openai + +import ( + "context" + "errors" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/endpoints/openai/types" + "github.com/mudler/LocalAI/core/http/endpoints/openai/turncoord" + "github.com/mudler/LocalAI/core/schema" +) + +// vadTick specs drive one synchronous turn-detection inspection at a time +// (no ticker), the same way classifySoundWindow's specs drive the +// sound-detection loop. The fake VAD answers in the coordinates of the audio +// it is HANDED — i.e. scan-clip coordinates once the buffer outgrows the +// window — exactly like the real backend. +var _ = Describe("vadTick", func() { + const rate = 16000 // InputSampleRate == localSampleRate: resample is a copy + + // pcm returns sec seconds of silent 16-bit PCM; content is irrelevant to + // the scripted VAD. + pcm := func(sec float64) []byte { + return make([]byte, int(sec*rate)*2) + } + bufferSec := func(s *Session) float64 { + return float64(len(s.InputAudioBuffer)) / (rate * 2) + } + + newHarness := func(td *types.TurnDetectionUnion, m *fakeModel) (*Session, *fakeTransport, *turnSink) { + session := &Session{ + TranscriptionOnly: true, // commit stops after the transcription events + TurnDetection: td, + InputAudioTranscription: &types.AudioTranscription{}, + ModelConfig: &config.ModelConfig{}, + ModelInterface: m, + InputSampleRate: rate, + respSink: newResponseSink(), + } + tr := &fakeTransport{} + sink := newTurnSink(session, &Conversation{}, tr, newLiveTurnState(session, tr), context.Background(), time.Now()) + return session, tr, sink + } + serverVad := &types.TurnDetectionUnion{ServerVad: &types.ServerVad{SilenceDurationMs: 500}} + semanticHigh := &types.TurnDetectionUnion{SemanticVad: &types.RealtimeSessionSemanticVad{Eagerness: "high"}} + + speaking := func(sink *turnSink) bool { + _, ok := sink.coord.State().(turncoord.Speaking) + return ok + } + + It("commits a normal short turn (extraction is behavior-neutral)", func() { + m := &fakeModel{ + vadSegments: []schema.VADSegment{{Start: 0.1, End: 0.6}}, + transcribeFinal: &schema.TranscriptionResult{Text: "go up"}, + } + session, tr, sink := newHarness(serverVad, m) + session.InputAudioBuffer = pcm(1.4) // under the 1.5s scan window: no clip + + vadTick(sink, 0.5) + + Expect(tr.countEvents(types.ServerEventTypeInputAudioBufferSpeechStarted)).To(Equal(1)) + Expect(tr.countEvents(types.ServerEventTypeInputAudioBufferSpeechStopped)).To(Equal(1)) + Expect(tr.countEvents(types.ServerEventTypeInputAudioBufferCommitted)).To(Equal(1)) + Expect(session.InputAudioBuffer).To(BeEmpty(), "commit drops the whole inspected window") + Expect(speaking(sink)).To(BeFalse()) + + session.respSink.wait() + Expect(tr.countEvents(types.ServerEventTypeConversationItemInputAudioTranscriptionCompleted)).To(Equal(1)) + }) + + It("hands the VAD only the scan window and rebases its answer", func() { + var scanned []int + m := &fakeModel{ + vadFn: func(req *schema.VADRequest) (*schema.VADResponse, error) { + scanned = append(scanned, len(req.Audio)) + // Clip coordinates: speech ends 0.9s into the 1.5s window, + // leaving 0.6s of trailing silence > the 0.5s threshold. + return &schema.VADResponse{Segments: []schema.VADSegment{{Start: 0.2, End: 0.9}}}, nil + }, + transcribeFinal: &schema.TranscriptionResult{Text: "clipped"}, + } + session, tr, sink := newHarness(serverVad, m) + session.InputAudioBuffer = pcm(20) + + vadTick(sink, 0.5) + + Expect(scanned).To(Equal([]int{int(1.5 * rate)}), "server_vad window = silence 0.5s + 1s margin") + Expect(tr.countEvents(types.ServerEventTypeInputAudioBufferCommitted)).To(Equal(1), + "rebased segment end (18.5+0.9) leaves 0.6s trailing silence in buffer coordinates") + }) + + It("commits when trailing silence outruns the scan window instead of discarding the turn", func() { + call := 0 + m := &fakeModel{ + vadFn: func(req *schema.VADRequest) (*schema.VADResponse, error) { + call++ + if call == 1 { + // Speech still running at the end of the inspected audio. + return &schema.VADResponse{Segments: []schema.VADSegment{{Start: 0.2, End: 0}}}, nil + } + // Later ticks: the (clipped) window is all silence. + return &schema.VADResponse{}, nil + }, + transcribeFinal: &schema.TranscriptionResult{Text: "late silence"}, + } + session, tr, sink := newHarness(serverVad, m) + session.InputAudioBuffer = pcm(1.4) + vadTick(sink, 0.5) + Expect(speaking(sink)).To(BeTrue()) + Expect(tr.countEvents(types.ServerEventTypeInputAudioBufferCommitted)).To(BeZero()) + + session.InputAudioBuffer = append(session.InputAudioBuffer, pcm(2.6)...) // 4s total: clip is in effect + vadTick(sink, 0.5) + + Expect(tr.countEvents(types.ServerEventTypeInputAudioBufferSpeechStopped)).To(Equal(1)) + Expect(tr.countEvents(types.ServerEventTypeInputAudioBufferCommitted)).To(Equal(1)) + Expect(session.InputAudioBuffer).To(BeEmpty()) + Expect(speaking(sink)).To(BeFalse()) + session.respSink.wait() + Expect(tr.countEvents(types.ServerEventTypeConversationItemInputAudioTranscriptionCompleted)).To(Equal(1)) + }) + + It("stays bounded when segments never stop (the noise-floor pathology)", func() { + var maxScan int + m := &fakeModel{ + vadFn: func(req *schema.VADRequest) (*schema.VADResponse, error) { + if len(req.Audio) > maxScan { + maxScan = len(req.Audio) + } + return &schema.VADResponse{Segments: []schema.VADSegment{{Start: 0.1, End: 0}}}, nil + }, + } + session, tr, sink := newHarness(serverVad, m) + + for i := 0; i < 95; i++ { + session.InputAudioBuffer = append(session.InputAudioBuffer, pcm(1)...) + vadTick(sink, 0.5) + } + + Expect(maxScan).To(Equal(int(1.5*rate)), "VAD never rescans more than the window") + Expect(bufferSec(session)).To(BeNumerically("<=", maxTurnBufferSec), "retention bound holds") + Expect(speaking(sink)).To(BeTrue(), "the turn is neither committed nor aborted") + Expect(tr.countEvents(types.ServerEventTypeInputAudioBufferSpeechStarted)).To(Equal(1)) + Expect(tr.countEvents(types.ServerEventTypeInputAudioBufferCommitted)).To(BeZero()) + }) + + It("keeps the live feed gapless across a retention trim", func() { + m := &fakeModel{ + vadFn: func(req *schema.VADRequest) (*schema.VADResponse, error) { + return &schema.VADResponse{Segments: []schema.VADSegment{{Start: 0.1, End: 0}}}, nil + }, + } + session, _, sink := newHarness(semanticHigh, m) + session.InputAudioBuffer = pcm(2) + vadTick(sink, 0.5) // opens the turn + live stream, feeds the onset audio + Expect(m.liveOpened).To(Equal(1)) + + session.InputAudioBuffer = append(session.InputAudioBuffer, pcm(89)...) // 91s: over the 90s bound + vadTick(sink, 0.5) + + Expect(bufferSec(session)).To(BeNumerically("<=", maxTurnBufferSec)) + total := 0 + for _, chunk := range m.liveSession.fed { + total += len(chunk) + } + // Everything ever buffered minus the one held-back resample-edge + // sample: no gap (undercount) and no re-feed (overcount) across the + // trim's cursor rebase. + Expect(total).To(Equal(91*rate-1), "fed samples = all audio seen minus the held-back tail sample") + }) + + It("bounds memory when the VAD backend keeps failing", func() { + m := &fakeModel{vadErr: errors.New("backend down")} + session, tr, sink := newHarness(serverVad, m) + session.InputAudioBuffer = pcm(95) + + vadTick(sink, 0.5) + + Expect(bufferSec(session)).To(BeNumerically("<=", maxTurnBufferSec), "retention trim runs before the VAD call") + Expect(tr.countEvents(types.ServerEventTypeError)).To(Equal(1)) + }) +}) + +var _ = Describe("vadScanWindowSec", func() { + It("sizes from the silence the commit test must measure, plus the warm-up margin", func() { + Expect(vadScanWindowSec(nil, 0.5, nil)).To(Equal(1.5)) + Expect(vadScanWindowSec(&types.RealtimeSessionSemanticVad{Eagerness: "high"}, 0.5, nil)).To(Equal(3.0)) + Expect(vadScanWindowSec(&types.RealtimeSessionSemanticVad{Eagerness: "low"}, 0.5, nil)).To(Equal(9.0)) + }) + + It("lets vad_window_sec widen but never narrow the window", func() { + cfg := &config.ModelConfig{} + cfg.Pipeline.TurnDetection.VadWindowSec = 10 + Expect(vadScanWindowSec(nil, 0.5, cfg)).To(Equal(10.0)) + cfg.Pipeline.TurnDetection.VadWindowSec = 0.2 + Expect(vadScanWindowSec(nil, 0.5, cfg)).To(Equal(1.5), "values below the floor are ignored") + }) +}) diff --git a/core/http/endpoints/openai/realtime_voicegate.go b/core/http/endpoints/openai/realtime_voicegate.go index 475b45e8f..c9b78ae9d 100644 --- a/core/http/endpoints/openai/realtime_voicegate.go +++ b/core/http/endpoints/openai/realtime_voicegate.go @@ -75,7 +75,7 @@ func newVoiceGate( // Resolved like every other pipeline sub-model (one alias hop), so an // aliased voice_recognition model gets its target's backend. - recCfg, err := cl.LoadResolvedModelConfig(cfg.Model, ml.ModelPath) + recCfg, err := cl.LoadResolvedModelConfig(cfg.Model, ml.ModelPath, appConfig.ToConfigLoaderOptions()...) if err != nil { return nil, fmt.Errorf("voice_recognition: failed to load model %q: %w", cfg.Model, err) } @@ -261,8 +261,10 @@ func (g *voiceGate) Authorize(ctx context.Context, wavPath string) (allowed bool // decide interprets an Authorize result against the gate's when-policy and the // session's prior verification state. -// proceed: run the LLM response for this utterance. -// markVerified: record a successful first-utterance verification. +// +// proceed: run the LLM response for this utterance. +// markVerified: record a successful first-utterance verification. +// // Note: when:first AND alreadyVerified is normally handled by the caller // skipping Authorize entirely; if it still reaches here, proceed is true. func (g *voiceGate) decide(alreadyVerified, allowed bool) (proceed, markVerified bool) { diff --git a/core/http/endpoints/openai/types/classifier.go b/core/http/endpoints/openai/types/classifier.go new file mode 100644 index 000000000..43749d3c0 --- /dev/null +++ b/core/http/endpoints/openai/types/classifier.go @@ -0,0 +1,470 @@ +package types + +import ( + "encoding/json" + "fmt" + "regexp" + "slices" + "strconv" + "strings" +) + +// ClassifierConfig is a LocalAI extension to the Realtime API +// (session.localai_classifier, response.localai_classifier): instead of +// autoregressive generation, each user turn is prefill-scored against a +// fixed option list via the Score primitive and the winning option's canned +// reply / tool call is emitted. Built for hardware that can afford prefill +// but not decode (e.g. a Raspberry Pi running a small LLM). +type ClassifierConfig struct { + // Enabled is a pointer so a response-level override can force + // classification off for one response ({"enabled": false}) without + // replacing the session's option list. nil means "on when options + // exist". + Enabled *bool `json:"enabled,omitempty"` + + // Options the user turn is scored against. Replaced wholesale by + // session.update / response.create, like tools. + Options []ClassifierOption `json:"options,omitempty"` + + // Threshold is the softmax-probability floor the best option must + // clear; below it the fallback applies. 0 always picks the argmax. + Threshold float64 `json:"threshold,omitempty"` + + // Normalization selects how candidate log-probs are compared before + // the softmax: "raw" (default, joint log-prob) or "mean" + // (length-normalized) — same semantics as the router's + // score_normalization. + Normalization string `json:"normalization,omitempty"` + + // HistoryItems selects what gets scored. 0 (default) and -1 score + // only the latest user message; a positive N includes the trailing N + // conversation messages, role-labeled. Prior turns echo option names + // (the canned replies especially) and empirically dominate small + // scoring models — only opt into history with a scorer large enough + // to weigh it. + HistoryItems int `json:"history_items,omitempty"` + + // Fallback controls what happens when no option clears the + // threshold. nil behaves like {"mode": "none"}. + Fallback *ClassifierFallback `json:"fallback,omitempty"` + + // Address, when set, gates every turn on the assistant being + // addressed by name ("Drone go up", not just "go up") — the + // wake-word pattern. The check is a deterministic word match on the + // transcript: scoring cannot do it (a 1.2B scorer rates "go up" as + // addressed=1.0 even with a dedicated addressing stage) and matching + // is free, so unaddressed ambient speech skips scoring entirely. + Address *ClassifierAddress `json:"address,omitempty"` +} + +// ClassifierAddress configures name-gating for classifier mode. +type ClassifierAddress struct { + // Names that count as addressing the assistant, matched as + // case-insensitive whole words against the latest user turn. + Names []string `json:"names"` + + // Mode when the turn does not mention a name: "ignore" (default — + // the response completes silently, the right behavior for ambient + // conversation) or "reply" (speak Reply). + Mode string `json:"mode,omitempty"` + + // Reply spoken in "reply" mode. + Reply string `json:"reply,omitempty"` +} + +// Address gate modes. +const ( + ClassifierAddressIgnore = "ignore" + ClassifierAddressReply = "reply" +) + +// ClassifierNotAddressed is the ClassifierResultEvent.Fallback value for +// turns dropped by the address gate. It is an event-only value — the +// config fallback modes stay none|reply|generate. +const ClassifierNotAddressed = "not_addressed" + +// AddressMode returns the effective address-gate mode. +func (a *ClassifierAddress) AddressMode() string { + if a == nil || a.Mode == "" { + return ClassifierAddressIgnore + } + return a.Mode +} + +// ClassifierOption is one selectable intent: what to match on +// (Description), what to say when chosen (Reply) and, optionally, a canned +// tool call the client executes. +type ClassifierOption struct { + // ID identifies the option in results and doubles as the scored + // route label, so keep it short — its tokens are what the model + // actually scores. + ID string `json:"id"` + + // Description tells the model when the option applies (e.g. "the + // user asks the drone to move or fly up/higher"). It goes into the + // classification system prompt. + Description string `json:"description"` + + // Reply is the canned assistant reply spoken/emitted when the + // option wins. Empty means the option is silent (tool-only). + Reply string `json:"reply,omitempty"` + + // Tool, when set, is emitted as a function_call item with these + // exact arguments when the option wins. + Tool *ClassifierTool `json:"tool,omitempty"` +} + +// ClassifierTool is a canned function call. Arguments is a raw JSON +// object; with Slots it becomes a template whose "{{name}}" placeholders +// are filled by a short constrained completion after classification — +// the hybrid between prefill-only classification and full generation. +type ClassifierTool struct { + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments,omitempty"` + + // Slots declares the argument holes to fill by inference when the + // option wins. Number slots substitute the quoted placeholder + // ("{{name}}" -> 3.5) so YAML/JSON templates stay well-formed; enum + // and string slots substitute inside their quotes. + Slots []ClassifierSlot `json:"slots,omitempty"` +} + +// Classifier slot types. +const ( + ClassifierSlotNumber = "number" + ClassifierSlotEnum = "enum" + ClassifierSlotString = "string" +) + +// ClassifierSlot is one inferred argument of a classifier tool call. +type ClassifierSlot struct { + // Name of the slot; "{{name}}" in the arguments template marks where + // its value lands, and the model sees it as a JSON field name. + Name string `json:"name"` + + // Type constrains the completion grammar: "number", "enum" or + // "string". + Type string `json:"type"` + + // Values enumerates the admissible values for enum slots. + Values []string `json:"values,omitempty"` + + // Default applies when inference fails outright. Enum defaults must + // be one of Values; number defaults must parse as a number. A slot + // without a default makes the whole response fall back on failure. + Default string `json:"default,omitempty"` + + // Hint is appended to the option's description in the scoring/fill + // system prompt (e.g. "assume meters when the user gives no units"). + Hint string `json:"hint,omitempty"` +} + +// slotPlaceholder returns the template marker for a slot. +func slotPlaceholder(name string) string { return "{{" + name + "}}" } + +// SampleValue returns a syntactically valid stand-in for template +// validation: the default when set, otherwise a type-appropriate value. +func (s *ClassifierSlot) SampleValue() string { + if s.Default != "" { + return s.Default + } + switch s.Type { + case ClassifierSlotNumber: + return "0" + case ClassifierSlotEnum: + if len(s.Values) > 0 { + return s.Values[0] + } + } + return "sample" +} + +// SpliceArguments fills the tool's argument template with the given slot +// values and returns the final JSON arguments string. Number values +// replace the quoted placeholder so they land unquoted; other types are +// JSON-string-escaped in place. The result must parse as a JSON object. +func (t *ClassifierTool) SpliceArguments(values map[string]string) (string, error) { + args := "{}" + if len(t.Arguments) > 0 { + args = string(t.Arguments) + } + for i := range t.Slots { + s := &t.Slots[i] + v, ok := values[s.Name] + if !ok || v == "" { + return "", fmt.Errorf("classifier: no value for slot %q", s.Name) + } + ph := slotPlaceholder(s.Name) + if s.Type == ClassifierSlotNumber { + args = strings.ReplaceAll(args, `"`+ph+`"`, v) + } else { + esc, err := json.Marshal(v) + if err != nil { + return "", err + } + args = strings.ReplaceAll(args, ph, string(esc[1:len(esc)-1])) + } + } + var obj map[string]any + if err := json.Unmarshal([]byte(args), &obj); err != nil { + return "", fmt.Errorf("classifier: spliced tool arguments are not a JSON object: %w", err) + } + return args, nil +} + +// SpliceReply fills "{{name}}" placeholders in the option's spoken reply +// with the same slot values that filled the tool arguments, as plain text +// ("Going {{distance}} {{units}}." → "Going 3 meters."), so the reply can +// confirm what was actually inferred. Values are optional in the reply: +// placeholders without a value stay literal, and options without slots (or +// a nil value set) return the reply verbatim. +func (o *ClassifierOption) SpliceReply(values map[string]string) string { + reply := o.Reply + if o.Tool == nil || len(values) == 0 { + return reply + } + for i := range o.Tool.Slots { + s := &o.Tool.Slots[i] + if v, ok := values[s.Name]; ok && v != "" { + reply = strings.ReplaceAll(reply, slotPlaceholder(s.Name), v) + } + } + return reply +} + +// SlotDefaults returns every slot's default value, or an error naming the +// first slot without one — the fill-failure path either recovers with a +// complete default set or not at all. +func (t *ClassifierTool) SlotDefaults() (map[string]string, error) { + values := make(map[string]string, len(t.Slots)) + for i := range t.Slots { + if t.Slots[i].Default == "" { + return nil, fmt.Errorf("classifier: slot %q has no default", t.Slots[i].Name) + } + values[t.Slots[i].Name] = t.Slots[i].Default + } + return values, nil +} + +// Classifier fallback modes. +const ( + // ClassifierFallbackNone completes the response with no output. + ClassifierFallbackNone = "none" + // ClassifierFallbackReply speaks/emits the canned fallback reply. + ClassifierFallbackReply = "reply" + // ClassifierFallbackGenerate falls through to normal autoregressive + // generation for that response. + ClassifierFallbackGenerate = "generate" +) + +// ClassifierFallback selects the below-threshold behavior. +type ClassifierFallback struct { + Mode string `json:"mode,omitempty"` + Reply string `json:"reply,omitempty"` +} + +// Active reports whether classification should run: explicitly enabled, or +// enabled by default because options are present. +func (c *ClassifierConfig) Active() bool { + if c == nil { + return false + } + if c.Enabled != nil { + return *c.Enabled && len(c.Options) > 0 + } + return len(c.Options) > 0 +} + +// FallbackMode returns the effective fallback mode. +func (c *ClassifierConfig) FallbackMode() string { + if c == nil || c.Fallback == nil || c.Fallback.Mode == "" { + return ClassifierFallbackNone + } + return c.Fallback.Mode +} + +// Validate checks the invariants the scoring engine relies on. It is +// shared by the session.update path and pipeline-config seeding so both +// reject bad option lists the same way. +func (c *ClassifierConfig) Validate() error { + if c == nil { + return nil + } + if c.Threshold < 0 || c.Threshold >= 1 { + return fmt.Errorf("classifier: threshold must be in [0,1), got %v", c.Threshold) + } + switch c.Normalization { + case "", "raw", "mean": + default: + return fmt.Errorf("classifier: normalization must be \"raw\" or \"mean\", got %q", c.Normalization) + } + if c.HistoryItems < -1 { + return fmt.Errorf("classifier: history_items must be >= -1, got %d", c.HistoryItems) + } + switch c.FallbackMode() { + case ClassifierFallbackNone, ClassifierFallbackReply, ClassifierFallbackGenerate: + default: + return fmt.Errorf("classifier: fallback mode must be one of none|reply|generate, got %q", c.Fallback.Mode) + } + if c.FallbackMode() == ClassifierFallbackReply && (c.Fallback == nil || c.Fallback.Reply == "") { + return fmt.Errorf("classifier: fallback mode \"reply\" requires a non-empty fallback reply") + } + if c.Address != nil { + named := false + for _, n := range c.Address.Names { + if n != "" { + named = true + break + } + } + if !named { + return fmt.Errorf("classifier: address gate requires at least one non-empty name") + } + switch c.Address.AddressMode() { + case ClassifierAddressIgnore, ClassifierAddressReply: + default: + return fmt.Errorf("classifier: address mode must be one of ignore|reply, got %q", c.Address.Mode) + } + if c.Address.AddressMode() == ClassifierAddressReply && c.Address.Reply == "" { + return fmt.Errorf("classifier: address mode \"reply\" requires a non-empty reply") + } + } + seen := make(map[string]struct{}, len(c.Options)) + for i, opt := range c.Options { + if opt.ID == "" { + return fmt.Errorf("classifier: option %d has an empty id", i) + } + if _, dup := seen[opt.ID]; dup { + return fmt.Errorf("classifier: duplicate option id %q", opt.ID) + } + seen[opt.ID] = struct{}{} + if opt.Description == "" { + return fmt.Errorf("classifier: option %q has an empty description", opt.ID) + } + if opt.Tool != nil { + if opt.Tool.Name == "" { + return fmt.Errorf("classifier: option %q has a tool with an empty name", opt.ID) + } + if len(opt.Tool.Arguments) > 0 && len(opt.Tool.Slots) == 0 { + var obj map[string]any + if err := json.Unmarshal(opt.Tool.Arguments, &obj); err != nil { + return fmt.Errorf("classifier: option %q tool arguments must be a JSON object: %w", opt.ID, err) + } + } + if err := validateSlots(opt.Tool); err != nil { + return fmt.Errorf("classifier: option %q: %w", opt.ID, err) + } + } + } + return nil +} + +// slotNamePattern keeps slot names safe to embed as JSON field names and +// template placeholders without escaping. +var slotNamePattern = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*$`) + +func validateSlots(t *ClassifierTool) error { + if len(t.Slots) == 0 { + return nil + } + args := string(t.Arguments) + seen := make(map[string]struct{}, len(t.Slots)) + sample := make(map[string]string, len(t.Slots)) + for i := range t.Slots { + s := &t.Slots[i] + if !slotNamePattern.MatchString(s.Name) { + return fmt.Errorf("slot %d has invalid name %q", i, s.Name) + } + if _, dup := seen[s.Name]; dup { + return fmt.Errorf("duplicate slot %q", s.Name) + } + seen[s.Name] = struct{}{} + switch s.Type { + case ClassifierSlotNumber: + if s.Default != "" { + if _, err := strconv.ParseFloat(s.Default, 64); err != nil { + return fmt.Errorf("slot %q: number default %q does not parse", s.Name, s.Default) + } + } + case ClassifierSlotEnum: + if len(s.Values) == 0 { + return fmt.Errorf("slot %q: enum slots need values", s.Name) + } + if slices.Contains(s.Values, "") { + return fmt.Errorf("slot %q: enum values must be non-empty", s.Name) + } + if s.Default != "" && !slices.Contains(s.Values, s.Default) { + return fmt.Errorf("slot %q: default %q is not one of its values", s.Name, s.Default) + } + case ClassifierSlotString: + default: + return fmt.Errorf("slot %q: type must be one of number|enum|string, got %q", s.Name, s.Type) + } + if !strings.Contains(args, slotPlaceholder(s.Name)) { + return fmt.Errorf("slot %q: arguments template does not reference {{%s}}", s.Name, s.Name) + } + sample[s.Name] = s.SampleValue() + } + // The template with type-appropriate values must produce a JSON + // object, catching e.g. an unquoted string placeholder up front. + if _, err := t.SpliceArguments(sample); err != nil { + return fmt.Errorf("arguments template does not splice: %w", err) + } + return nil +} + +// ClassifierScore is one entry of the softmax distribution over options. +type ClassifierScore struct { + ID string `json:"id"` + Score float64 `json:"score"` +} + +// ClassifierResultEvent is a LocalAI extension server event +// (localai.classifier.result) emitted once per classifier-handled response +// — including fallbacks — before the output items, so clients can +// visualize the decision and its confidence. +type ClassifierResultEvent struct { + ServerEventBase + + // The ID of the response this classification belongs to. + ResponseID string `json:"response_id"` + + // The full softmax distribution, in option-declaration order. + Scores []ClassifierScore `json:"scores"` + + // The winning option id, or "" when the fallback applied. + ChosenID string `json:"chosen_id,omitempty"` + + // The threshold the winner had to clear. + Threshold float64 `json:"threshold"` + + // The fallback mode that applied, or "" when an option was chosen. + Fallback string `json:"fallback,omitempty"` + + // Wall-clock scoring latency. + LatencyMs int64 `json:"latency_ms"` + + // The chosen option's final tool arguments when its slots were filled + // by inference (the hybrid classify-then-complete path). + Arguments string `json:"arguments,omitempty"` + + // Wall-clock slot-fill latency; zero when the option has no slots. + FillLatencyMs int64 `json:"fill_latency_ms,omitempty"` +} + +func (m ClassifierResultEvent) ServerEventType() ServerEventType { + return ServerEventTypeClassifierResult +} + +func (m ClassifierResultEvent) MarshalJSON() ([]byte, error) { + type typeAlias ClassifierResultEvent + type typeWrapper struct { + typeAlias + Type ServerEventType `json:"type"` + } + shadow := typeWrapper{ + typeAlias: typeAlias(m), + Type: m.ServerEventType(), + } + return json.Marshal(shadow) +} diff --git a/core/http/endpoints/openai/types/classifier_test.go b/core/http/endpoints/openai/types/classifier_test.go new file mode 100644 index 000000000..57e1ab66e --- /dev/null +++ b/core/http/endpoints/openai/types/classifier_test.go @@ -0,0 +1,299 @@ +package types_test + +import ( + "encoding/json" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/http/endpoints/openai/types" +) + +func validClassifier() *types.ClassifierConfig { + return &types.ClassifierConfig{ + Threshold: 0.35, + Options: []types.ClassifierOption{ + { + ID: "up", + Description: "the user asks the drone to fly up", + Reply: "Going up.", + Tool: &types.ClassifierTool{Name: "move", Arguments: json.RawMessage(`{"direction":"up"}`)}, + }, + {ID: "greeting", Description: "the user greets the assistant", Reply: "Hello."}, + }, + Fallback: &types.ClassifierFallback{Mode: types.ClassifierFallbackReply, Reply: "Say again?"}, + } +} + +var _ = Describe("ClassifierConfig", func() { + Describe("JSON round-trip", func() { + It("survives marshal/unmarshal with all fields", func() { + in := validClassifier() + enabled := true + in.Enabled = &enabled + in.Normalization = "mean" + in.HistoryItems = -1 + + data, err := json.Marshal(in) + Expect(err).ToNot(HaveOccurred()) + + var out types.ClassifierConfig + Expect(json.Unmarshal(data, &out)).To(Succeed()) + Expect(out.Enabled).ToNot(BeNil()) + Expect(*out.Enabled).To(BeTrue()) + Expect(out.Threshold).To(Equal(0.35)) + Expect(out.Normalization).To(Equal("mean")) + Expect(out.HistoryItems).To(Equal(-1)) + Expect(out.Options).To(HaveLen(2)) + Expect(out.Options[0].Tool.Name).To(Equal("move")) + Expect(string(out.Options[0].Tool.Arguments)).To(MatchJSON(`{"direction":"up"}`)) + Expect(out.Fallback.Mode).To(Equal("reply")) + }) + + It("is carried by RealtimeSession under localai_classifier", func() { + s := types.RealtimeSession{LocalAIClassifier: validClassifier()} + data, err := json.Marshal(s) + Expect(err).ToNot(HaveOccurred()) + Expect(string(data)).To(ContainSubstring(`"localai_classifier"`)) + + var back types.RealtimeSession + Expect(json.Unmarshal(data, &back)).To(Succeed()) + Expect(back.LocalAIClassifier).ToNot(BeNil()) + Expect(back.LocalAIClassifier.Options).To(HaveLen(2)) + }) + + It("is carried by ResponseCreateParams under localai_classifier", func() { + var params types.ResponseCreateParams + Expect(json.Unmarshal([]byte(`{"localai_classifier":{"enabled":false}}`), ¶ms)).To(Succeed()) + Expect(params.LocalAIClassifier).ToNot(BeNil()) + Expect(params.LocalAIClassifier.Enabled).ToNot(BeNil()) + Expect(*params.LocalAIClassifier.Enabled).To(BeFalse()) + }) + }) + + Describe("Active", func() { + It("is inactive when nil", func() { + var c *types.ClassifierConfig + Expect(c.Active()).To(BeFalse()) + }) + + It("defaults to active when options exist", func() { + Expect(validClassifier().Active()).To(BeTrue()) + }) + + It("is inactive without options even when enabled", func() { + enabled := true + c := &types.ClassifierConfig{Enabled: &enabled} + Expect(c.Active()).To(BeFalse()) + }) + + It("honors an explicit enabled=false override", func() { + c := validClassifier() + disabled := false + c.Enabled = &disabled + Expect(c.Active()).To(BeFalse()) + }) + }) + + Describe("Validate", func() { + It("accepts a valid config and a nil config", func() { + Expect(validClassifier().Validate()).To(Succeed()) + var c *types.ClassifierConfig + Expect(c.Validate()).To(Succeed()) + }) + + It("rejects out-of-range thresholds", func() { + c := validClassifier() + c.Threshold = 1.0 + Expect(c.Validate()).To(MatchError(ContainSubstring("threshold"))) + c.Threshold = -0.1 + Expect(c.Validate()).To(MatchError(ContainSubstring("threshold"))) + }) + + It("rejects unknown normalization", func() { + c := validClassifier() + c.Normalization = "zscore" + Expect(c.Validate()).To(MatchError(ContainSubstring("normalization"))) + }) + + It("rejects history_items below -1", func() { + c := validClassifier() + c.HistoryItems = -2 + Expect(c.Validate()).To(MatchError(ContainSubstring("history_items"))) + }) + + It("rejects unknown fallback modes", func() { + c := validClassifier() + c.Fallback = &types.ClassifierFallback{Mode: "retry"} + Expect(c.Validate()).To(MatchError(ContainSubstring("fallback mode"))) + }) + + It("rejects a reply fallback without a reply", func() { + c := validClassifier() + c.Fallback = &types.ClassifierFallback{Mode: types.ClassifierFallbackReply} + Expect(c.Validate()).To(MatchError(ContainSubstring("fallback reply"))) + }) + + It("rejects empty and duplicate option ids", func() { + c := validClassifier() + c.Options[1].ID = "" + Expect(c.Validate()).To(MatchError(ContainSubstring("empty id"))) + c.Options[1].ID = "up" + Expect(c.Validate()).To(MatchError(ContainSubstring("duplicate option id"))) + }) + + It("rejects an option without a description", func() { + c := validClassifier() + c.Options[0].Description = "" + Expect(c.Validate()).To(MatchError(ContainSubstring("empty description"))) + }) + + It("rejects tools with no name or non-object arguments", func() { + c := validClassifier() + c.Options[0].Tool = &types.ClassifierTool{} + Expect(c.Validate()).To(MatchError(ContainSubstring("empty name"))) + c.Options[0].Tool = &types.ClassifierTool{Name: "move", Arguments: json.RawMessage(`["up"]`)} + Expect(c.Validate()).To(MatchError(ContainSubstring("JSON object"))) + }) + }) + + Describe("FallbackMode", func() { + It("defaults to none", func() { + Expect((&types.ClassifierConfig{}).FallbackMode()).To(Equal(types.ClassifierFallbackNone)) + var c *types.ClassifierConfig + Expect(c.FallbackMode()).To(Equal(types.ClassifierFallbackNone)) + }) + }) + + Describe("ClassifierResultEvent", func() { + It("marshals with the localai.classifier.result type tag", func() { + ev := types.ClassifierResultEvent{ + ResponseID: "resp_1", + Scores: []types.ClassifierScore{{ID: "up", Score: 0.9}, {ID: "down", Score: 0.1}}, + ChosenID: "up", + Threshold: 0.35, + LatencyMs: 12, + } + data, err := json.Marshal(ev) + Expect(err).ToNot(HaveOccurred()) + Expect(string(data)).To(ContainSubstring(`"type":"localai.classifier.result"`)) + Expect(string(data)).To(ContainSubstring(`"chosen_id":"up"`)) + Expect(string(data)).To(ContainSubstring(`"threshold":0.35`)) + }) + }) +}) + +var _ = Describe("ClassifierTool slots", func() { + tool := func(slots ...types.ClassifierSlot) *types.ClassifierTool { + return &types.ClassifierTool{ + Name: "move", + Arguments: json.RawMessage(`{"direction":"up","distance":"{{distance}}","units":"{{units}}"}`), + Slots: slots, + } + } + numberSlot := types.ClassifierSlot{Name: "distance", Type: types.ClassifierSlotNumber, Default: "1"} + enumSlot := types.ClassifierSlot{Name: "units", Type: types.ClassifierSlotEnum, Values: []string{"m", "ft"}, Default: "m"} + + cfgWith := func(t *types.ClassifierTool) *types.ClassifierConfig { + return &types.ClassifierConfig{Options: []types.ClassifierOption{{ID: "up", Description: "d", Tool: t}}} + } + + Describe("Validate", func() { + It("accepts a well-formed slotted tool", func() { + Expect(cfgWith(tool(numberSlot, enumSlot)).Validate()).To(Succeed()) + }) + + It("rejects unknown slot types", func() { + bad := numberSlot + bad.Type = "float" + Expect(cfgWith(tool(bad, enumSlot)).Validate()).To(MatchError(ContainSubstring("number|enum|string"))) + }) + + It("rejects enum slots without values", func() { + bad := enumSlot + bad.Values = nil + bad.Default = "" + Expect(cfgWith(tool(numberSlot, bad)).Validate()).To(MatchError(ContainSubstring("need values"))) + }) + + It("rejects enum defaults outside the value set", func() { + bad := enumSlot + bad.Default = "yards" + Expect(cfgWith(tool(numberSlot, bad)).Validate()).To(MatchError(ContainSubstring("not one of"))) + }) + + It("rejects number defaults that do not parse", func() { + bad := numberSlot + bad.Default = "three" + Expect(cfgWith(tool(bad, enumSlot)).Validate()).To(MatchError(ContainSubstring("does not parse"))) + }) + + It("rejects empty enum values that cannot be spliced", func() { + bad := enumSlot + bad.Values = []string{"m", ""} + Expect(cfgWith(tool(numberSlot, bad)).Validate()).To(MatchError(ContainSubstring("must be non-empty"))) + }) + + It("rejects slots the template never references", func() { + t := tool(numberSlot, enumSlot, types.ClassifierSlot{Name: "speed", Type: types.ClassifierSlotNumber}) + Expect(cfgWith(t).Validate()).To(MatchError(ContainSubstring("{{speed}}"))) + }) + + It("rejects invalid slot names", func() { + bad := numberSlot + bad.Name = "dis tance" + Expect(cfgWith(tool(bad, enumSlot)).Validate()).To(MatchError(ContainSubstring("invalid name"))) + }) + }) + + Describe("SpliceArguments", func() { + It("substitutes numbers unquoted and strings escaped", func() { + args, err := tool(numberSlot, enumSlot).SpliceArguments(map[string]string{"distance": "3.5", "units": `m"eters`}) + Expect(err).ToNot(HaveOccurred()) + Expect(args).To(MatchJSON(`{"direction":"up","distance":3.5,"units":"m\"eters"}`)) + }) + + It("fails on missing values", func() { + _, err := tool(numberSlot, enumSlot).SpliceArguments(map[string]string{"distance": "3.5"}) + Expect(err).To(MatchError(ContainSubstring(`no value for slot "units"`))) + }) + }) + + Describe("SlotDefaults", func() { + It("returns every default", func() { + values, err := tool(numberSlot, enumSlot).SlotDefaults() + Expect(err).ToNot(HaveOccurred()) + Expect(values).To(Equal(map[string]string{"distance": "1", "units": "m"})) + }) + + It("names the slot lacking a default", func() { + bare := numberSlot + bare.Default = "" + _, err := tool(bare, enumSlot).SlotDefaults() + Expect(err).To(MatchError(ContainSubstring(`"distance"`))) + }) + }) + + Describe("SpliceReply", func() { + option := func(reply string, t *types.ClassifierTool) *types.ClassifierOption { + return &types.ClassifierOption{ID: "up", Description: "d", Reply: reply, Tool: t} + } + + It("substitutes slot values as plain text", func() { + o := option("Going up {{distance}} {{units}}.", tool(numberSlot, enumSlot)) + Expect(o.SpliceReply(map[string]string{"distance": "3.5", "units": "m"})).To(Equal("Going up 3.5 m.")) + }) + + It("leaves placeholders without a value literal", func() { + o := option("Going up {{distance}} {{units}}.", tool(numberSlot, enumSlot)) + Expect(o.SpliceReply(map[string]string{"distance": "3"})).To(Equal("Going up 3 {{units}}.")) + }) + + It("returns the reply verbatim without slots or values", func() { + o := option("Going up {{distance}}.", nil) + Expect(o.SpliceReply(map[string]string{"distance": "3"})).To(Equal("Going up {{distance}}.")) + slotted := option("Going up {{distance}}.", tool(numberSlot)) + Expect(slotted.SpliceReply(nil)).To(Equal("Going up {{distance}}.")) + }) + }) +}) diff --git a/core/http/endpoints/openai/types/server_events.go b/core/http/endpoints/openai/types/server_events.go index 6b0a233ee..b847a35a7 100644 --- a/core/http/endpoints/openai/types/server_events.go +++ b/core/http/endpoints/openai/types/server_events.go @@ -24,34 +24,38 @@ const ( // ServerEventTypeConversationItemSpeaker is a LocalAI extension: it reports // the recognized speaker for a user audio item. OpenAI clients ignore it. ServerEventTypeConversationItemSpeaker ServerEventType = "conversation.item.speaker" - ServerEventTypeInputAudioBufferCommitted ServerEventType = "input_audio_buffer.committed" - ServerEventTypeInputAudioBufferCleared ServerEventType = "input_audio_buffer.cleared" - ServerEventTypeInputAudioBufferSpeechStarted ServerEventType = "input_audio_buffer.speech_started" - ServerEventTypeInputAudioBufferSpeechStopped ServerEventType = "input_audio_buffer.speech_stopped" - ServerEventTypeInputAudioBufferTimeoutTriggered ServerEventType = "input_audio_buffer.timeout_triggered" - ServerEventTypeResponseCreated ServerEventType = "response.created" - ServerEventTypeResponseDone ServerEventType = "response.done" - ServerEventTypeResponseOutputItemAdded ServerEventType = "response.output_item.added" - ServerEventTypeResponseOutputItemDone ServerEventType = "response.output_item.done" - ServerEventTypeResponseContentPartAdded ServerEventType = "response.content_part.added" - ServerEventTypeResponseContentPartDone ServerEventType = "response.content_part.done" - ServerEventTypeResponseOutputTextDelta ServerEventType = "response.output_text.delta" - ServerEventTypeResponseOutputTextDone ServerEventType = "response.output_text.done" - ServerEventTypeResponseOutputAudioTranscriptDelta ServerEventType = "response.output_audio_transcript.delta" - ServerEventTypeResponseOutputAudioTranscriptDone ServerEventType = "response.output_audio_transcript.done" - ServerEventTypeResponseOutputAudioDelta ServerEventType = "response.output_audio.delta" - ServerEventTypeResponseOutputAudioDone ServerEventType = "response.output_audio.done" - ServerEventTypeResponseFunctionCallArgumentsDelta ServerEventType = "response.function_call_arguments.delta" - ServerEventTypeResponseFunctionCallArgumentsDone ServerEventType = "response.function_call_arguments.done" - ServerEventTypeResponseMcpCallArgumentsDelta ServerEventType = "response.mcp_call_arguments.delta" - ServerEventTypeResponseMcpCallArgumentsDone ServerEventType = "response.mcp_call_arguments.done" - ServerEventTypeResponseMcpCallInProgress ServerEventType = "response.mcp_call.in_progress" - ServerEventTypeResponseMcpCallCompleted ServerEventType = "response.mcp_call.completed" - ServerEventTypeResponseMcpCallFailed ServerEventType = "response.mcp_call.failed" - ServerEventTypeMcpListToolsInProgress ServerEventType = "mcp_list_tools.in_progress" - ServerEventTypeMcpListToolsCompleted ServerEventType = "mcp_list_tools.completed" - ServerEventTypeMcpListToolsFailed ServerEventType = "mcp_list_tools.failed" - ServerEventTypeRateLimitsUpdated ServerEventType = "rate_limits.updated" + // ServerEventTypeClassifierResult is a LocalAI extension: it carries the + // classifier-mode score distribution and decision for a response. OpenAI + // clients ignore it. + ServerEventTypeClassifierResult ServerEventType = "localai.classifier.result" + ServerEventTypeInputAudioBufferCommitted ServerEventType = "input_audio_buffer.committed" + ServerEventTypeInputAudioBufferCleared ServerEventType = "input_audio_buffer.cleared" + ServerEventTypeInputAudioBufferSpeechStarted ServerEventType = "input_audio_buffer.speech_started" + ServerEventTypeInputAudioBufferSpeechStopped ServerEventType = "input_audio_buffer.speech_stopped" + ServerEventTypeInputAudioBufferTimeoutTriggered ServerEventType = "input_audio_buffer.timeout_triggered" + ServerEventTypeResponseCreated ServerEventType = "response.created" + ServerEventTypeResponseDone ServerEventType = "response.done" + ServerEventTypeResponseOutputItemAdded ServerEventType = "response.output_item.added" + ServerEventTypeResponseOutputItemDone ServerEventType = "response.output_item.done" + ServerEventTypeResponseContentPartAdded ServerEventType = "response.content_part.added" + ServerEventTypeResponseContentPartDone ServerEventType = "response.content_part.done" + ServerEventTypeResponseOutputTextDelta ServerEventType = "response.output_text.delta" + ServerEventTypeResponseOutputTextDone ServerEventType = "response.output_text.done" + ServerEventTypeResponseOutputAudioTranscriptDelta ServerEventType = "response.output_audio_transcript.delta" + ServerEventTypeResponseOutputAudioTranscriptDone ServerEventType = "response.output_audio_transcript.done" + ServerEventTypeResponseOutputAudioDelta ServerEventType = "response.output_audio.delta" + ServerEventTypeResponseOutputAudioDone ServerEventType = "response.output_audio.done" + ServerEventTypeResponseFunctionCallArgumentsDelta ServerEventType = "response.function_call_arguments.delta" + ServerEventTypeResponseFunctionCallArgumentsDone ServerEventType = "response.function_call_arguments.done" + ServerEventTypeResponseMcpCallArgumentsDelta ServerEventType = "response.mcp_call_arguments.delta" + ServerEventTypeResponseMcpCallArgumentsDone ServerEventType = "response.mcp_call_arguments.done" + ServerEventTypeResponseMcpCallInProgress ServerEventType = "response.mcp_call.in_progress" + ServerEventTypeResponseMcpCallCompleted ServerEventType = "response.mcp_call.completed" + ServerEventTypeResponseMcpCallFailed ServerEventType = "response.mcp_call.failed" + ServerEventTypeMcpListToolsInProgress ServerEventType = "mcp_list_tools.in_progress" + ServerEventTypeMcpListToolsCompleted ServerEventType = "mcp_list_tools.completed" + ServerEventTypeMcpListToolsFailed ServerEventType = "mcp_list_tools.failed" + ServerEventTypeRateLimitsUpdated ServerEventType = "rate_limits.updated" ) // ServerEvent is the interface for server events. diff --git a/core/http/endpoints/openai/types/types.go b/core/http/endpoints/openai/types/types.go index 21e78004f..b6f6d9525 100644 --- a/core/http/endpoints/openai/types/types.go +++ b/core/http/endpoints/openai/types/types.go @@ -956,6 +956,11 @@ type RealtimeSession struct { // Controls how the realtime conversation is truncated prior to model inference. The default is auto. Truncation *TruncationUnion `json:"truncation,omitempty"` + + // LocalAIClassifier is a LocalAI extension: prefill-scored option + // selection instead of autoregressive generation. Replaced wholesale + // on update, like tools. OpenAI clients simply never set it. + LocalAIClassifier *ClassifierConfig `json:"localai_classifier,omitempty"` } func (r RealtimeSession) Type() SessionType { @@ -1191,6 +1196,11 @@ type ResponseCreateParams struct { // Tools available to the model. Tools []ToolUnion `json:"tools,omitempty"` + + // LocalAIClassifier is a LocalAI extension: when non-nil it replaces + // the session's classifier config for this response only — + // {"enabled": false} runs normal generation once. + LocalAIClassifier *ClassifierConfig `json:"localai_classifier,omitempty"` } type Response struct { diff --git a/core/http/endpoints/openai/types/types_suite_test.go b/core/http/endpoints/openai/types/types_suite_test.go new file mode 100644 index 000000000..ad2a1c5ce --- /dev/null +++ b/core/http/endpoints/openai/types/types_suite_test.go @@ -0,0 +1,13 @@ +package types_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestTypes(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Realtime types test suite") +} diff --git a/core/http/endpoints/openai/upscale.go b/core/http/endpoints/openai/upscale.go new file mode 100644 index 000000000..eab000f8e --- /dev/null +++ b/core/http/endpoints/openai/upscale.go @@ -0,0 +1,134 @@ +package openai + +import ( + "fmt" + "io" + "net/http" + "net/url" + "os" + "path/filepath" + "strconv" + "time" + + "github.com/google/uuid" + "github.com/labstack/echo/v4" + "github.com/mudler/xlog" + + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/schema" + model "github.com/mudler/LocalAI/pkg/model" +) + +// UpscaleEndpoint handles POST /v1/images/upscale +// +// @Summary Image upscaling +// @Description Upscale an image using a specified model (e.g. stable-diffusion-x4-upscaler). Accepts multipart/form-data. +// @Tags images +// @Accept multipart/form-data +// @Produce application/json +// @Param model formData string true "Upscaler model identifier (e.g. stable-diffusion-x4-upscaler)" +// @Param image formData file true "Input image file" +// @Param scale formData int false "Upscale factor: 2 or 4 (default 2)" +// @Success 200 {object} schema.OpenAIResponse +// @Failure 400 {object} map[string]string +// @Failure 500 {object} map[string]string +// @Router /v1/images/upscale [post] +func UpscaleEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { + return func(c echo.Context) error { + modelName := c.FormValue("model") + scaleStr := c.FormValue("scale") + + if modelName == "" { + xlog.Error("Upscale Endpoint - missing model") + return echo.NewHTTPError(http.StatusBadRequest, "missing model") + } + + scale := 2 + if scaleStr != "" { + v, err := strconv.Atoi(scaleStr) + if err != nil || (v != 2 && v != 4) { + return echo.NewHTTPError(http.StatusBadRequest, "scale must be 2 or 4") + } + scale = v + } + + // Read uploaded image + imageFile, err := c.FormFile("image") + if err != nil { + xlog.Error("Upscale Endpoint - missing image file", "error", err) + return echo.NewHTTPError(http.StatusBadRequest, "missing image file") + } + + imgSrc, err := imageFile.Open() + if err != nil { + return err + } + defer imgSrc.Close() + imgBytes, err := io.ReadAll(imgSrc) + if err != nil { + return err + } + + // Get model config from middleware context + cfg, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig) + if !ok || cfg == nil { + xlog.Error("Upscale Endpoint - model config not found in context") + return echo.ErrBadRequest + } + + tmpDir := filepath.Join(appConfig.GeneratedContentDir, "images") + if err := os.MkdirAll(tmpDir, 0750); err != nil { + return echo.NewHTTPError(http.StatusInternalServerError, "failed to prepare storage") + } + + // Write input image to a temp file + srcTmp, err := os.CreateTemp(tmpDir, "upscale_src_") + if err != nil { + return err + } + if _, err := srcTmp.Write(imgBytes); err != nil { + _ = srcTmp.Close() + _ = os.Remove(srcTmp.Name()) + return err + } + if err := srcTmp.Close(); err != nil { + xlog.Warn("Upscale Endpoint - failed to close src temp file", "error", err) + } + srcPath := srcTmp.Name() + defer os.Remove(srcPath) + + // Prepare output file path + id := uuid.New().String() + dstPath := filepath.Join(tmpDir, fmt.Sprintf("upscale_%s.png", id)) + + fn, err := backend.ImageUpscaleFunc(c.Request().Context(), srcPath, dstPath, scale, ml, *cfg, appConfig) + if err != nil { + return err + } + if err := fn(); err != nil { + _ = os.Remove(dstPath) + return err + } + + baseURL := middleware.BaseURL(c) + imgURL, err := url.JoinPath(baseURL, "generated-images", filepath.Base(dstPath)) + if err != nil { + _ = os.Remove(dstPath) + return err + } + + created := int(time.Now().Unix()) + resp := &schema.OpenAIResponse{ + ID: id, + Created: created, + Data: []schema.Item{{URL: imgURL}}, + Usage: &schema.OpenAIUsage{ + InputTokensDetails: &schema.InputTokensDetails{}, + }, + } + + return c.JSON(http.StatusOK, resp) + } +} diff --git a/core/http/endpoints/openai/upscale_test.go b/core/http/endpoints/openai/upscale_test.go new file mode 100644 index 000000000..2e13c8fc0 --- /dev/null +++ b/core/http/endpoints/openai/upscale_test.go @@ -0,0 +1,89 @@ +package openai + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/schema" + model "github.com/mudler/LocalAI/pkg/model" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Image upscaling", func() { + var ( + appConfig *config.ApplicationConfig + tmpDir string + ) + + BeforeEach(func() { + var err error + tmpDir, err = os.MkdirTemp("", "upscale") + Expect(err).ToNot(HaveOccurred()) + appConfig = config.NewApplicationConfig(config.WithGeneratedContentDir(tmpDir)) + }) + + AfterEach(func() { + Expect(os.RemoveAll(tmpDir)).To(Succeed()) + }) + + It("stores the result in the directory served by /generated-images", func() { + original := backend.ImageUpscaleFunc + backend.ImageUpscaleFunc = func(_ context.Context, _, dst string, scale int, _ *model.ModelLoader, _ config.ModelConfig, _ *config.ApplicationConfig) (func() error, error) { + Expect(scale).To(Equal(4)) + return func() error { + return os.WriteFile(dst, []byte("PNGDATA"), 0o644) + }, nil + } + DeferCleanup(func() { backend.ImageUpscaleFunc = original }) + + req, _ := makeMultipartRequest( + map[string]string{"model": "stable-diffusion-x4-upscaler", "scale": "4"}, + map[string][]byte{"image": []byte("IMAGEDATA")}, + ) + rec := httptest.NewRecorder() + ctx := echo.New().NewContext(req, rec) + ctx.Set(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG, &config.ModelConfig{Backend: "diffusers"}) + + Expect(UpscaleEndpoint(nil, nil, appConfig)(ctx)).To(Succeed()) + Expect(rec.Code).To(Equal(http.StatusOK)) + + var response schema.OpenAIResponse + Expect(json.Unmarshal(rec.Body.Bytes(), &response)).To(Succeed()) + Expect(response.Data).To(HaveLen(1)) + Expect(response.Data[0].URL).To(ContainSubstring("/generated-images/upscale_")) + + filename := filepath.Base(response.Data[0].URL) + contents, err := os.ReadFile(filepath.Join(tmpDir, "images", filename)) + Expect(err).ToNot(HaveOccurred()) + Expect(contents).To(Equal([]byte("PNGDATA"))) + }) + + It("rejects unsupported scale factors", func() { + req, _ := makeMultipartRequest( + map[string]string{"model": "stable-diffusion-x4-upscaler", "scale": "3"}, + map[string][]byte{"image": []byte("IMAGEDATA")}, + ) + rec := httptest.NewRecorder() + ctx := echo.New().NewContext(req, rec) + ctx.Set(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG, &config.ModelConfig{Backend: "diffusers"}) + + err := UpscaleEndpoint(nil, nil, appConfig)(ctx) + var httpErr *echo.HTTPError + Expect(err).To(MatchError(ContainSubstring("scale must be 2 or 4"))) + Expect(err).To(BeAssignableToTypeOf(httpErr)) + httpErr = err.(*echo.HTTPError) + Expect(httpErr.Code).To(Equal(http.StatusBadRequest)) + Expect(httpErr.Message).To(Equal("scale must be 2 or 4")) + Expect(bytes.TrimSpace(rec.Body.Bytes())).To(BeEmpty()) + }) +}) diff --git a/core/http/middleware/route_model.go b/core/http/middleware/route_model.go index 470bd05f5..643c6fa12 100644 --- a/core/http/middleware/route_model.go +++ b/core/http/middleware/route_model.go @@ -339,7 +339,7 @@ func buildClassifier(cfg *config.ModelConfig, deps ClassifierDeps) (router.Class // classifier model MUST carry a chat template — refusing // here beats silently falling back to a generic ChatML // envelope the model may not have been trained on. - renderer := newTemplateRenderer(deps.Evaluator, classifierCfg) + renderer := NewTemplateRenderer(deps.Evaluator, classifierCfg) if renderer == nil { return nil, fmt.Errorf( "router classifier score: classifier_model %q has no chat template "+ @@ -350,7 +350,7 @@ func buildClassifier(cfg *config.ModelConfig, deps ClassifierDeps) (router.Class } opts.PromptRenderer = renderer } - if st := pickAssistantTurnEnd(classifierCfg.StopWords, classifierCfg.TemplateConfig.ChatMessage); st != "" { + if st := PickAssistantTurnEnd(classifierCfg.StopWords, classifierCfg.TemplateConfig.ChatMessage); st != "" { opts.StopToken = st } // Token-exact conversation trim — score classifier drops the @@ -464,7 +464,7 @@ func validateRouterPolicies(classifierName string, rc config.RouterConfig) ([]ro return policies, nil } -// newTemplateRenderer adapts the templates.Evaluator + the classifier +// NewTemplateRenderer adapts the templates.Evaluator + the classifier // model's config into the router.PromptRenderer callback. The // resulting renderer pushes the routing system + user prompt through // the classifier model's full chat-template pipeline — per-role @@ -484,7 +484,7 @@ func validateRouterPolicies(classifierName string, rc config.RouterConfig) ([]ro // Returns nil (forcing the score classifier's chatMLRenderer // fallback) when either template piece is missing — partial // templating would still drop content. -func newTemplateRenderer(eval *templates.Evaluator, classifierCfg *config.ModelConfig) router.PromptRenderer { +func NewTemplateRenderer(eval *templates.Evaluator, classifierCfg *config.ModelConfig) router.PromptRenderer { if classifierCfg.TemplateConfig.Chat == "" || classifierCfg.TemplateConfig.ChatMessage == "" { return nil } @@ -502,7 +502,7 @@ func newTemplateRenderer(eval *templates.Evaluator, classifierCfg *config.ModelC } } -// pickAssistantTurnEnd returns the classifier model's assistant +// PickAssistantTurnEnd returns the classifier model's assistant // turn-end token — the one to suffix candidates with so the model's // "I'm done" signal folds into the per-candidate joint log-prob. // @@ -520,7 +520,7 @@ func newTemplateRenderer(eval *templates.Evaluator, classifierCfg *config.ModelC // // When no stopwords are configured at all, return "" — caller falls // back to defaultStopToken (<|im_end|>) inside the score classifier. -func pickAssistantTurnEnd(words []string, chatMessageTemplate string) string { +func PickAssistantTurnEnd(words []string, chatMessageTemplate string) string { if chatMessageTemplate != "" { for _, w := range words { if w != "" && strings.Contains(chatMessageTemplate, w) { diff --git a/core/http/middleware/route_model_test.go b/core/http/middleware/route_model_test.go index 4a9be2b12..0496ed9dd 100644 --- a/core/http/middleware/route_model_test.go +++ b/core/http/middleware/route_model_test.go @@ -301,7 +301,7 @@ var _ = Describe("RouteModel rendered classifier prompt", func() { // <|im_end|> first even though the actual Llama-3 assistant // turn-end is <|eot_id|>. The naive "stopwords[0]" pick would // suffix candidates with <|im_end|> — a token Llama-3 never - // emits at turn end. pickAssistantTurnEnd should scan the + // emits at turn end. PickAssistantTurnEnd should scan the // chat_message template and recognise <|eot_id|> as the real // turn-end. writeLlama3StyleClassifierModel(modelDir, "arch-router") @@ -340,7 +340,7 @@ type stubScorer struct { lastCandidates []string } -func (s *stubScorer) Score(_ context.Context, prompt string, candidates []string) ([]backend.CandidateScore, error) { +func (s *stubScorer) Score(_ context.Context, prompt string, _ int, candidates []string) ([]backend.CandidateScore, error) { s.lastPrompt = prompt s.lastCandidates = append([]string(nil), candidates...) out := make([]backend.CandidateScore, len(candidates)) @@ -498,7 +498,7 @@ template: // writeLlama3StyleClassifierModel writes a classifier model mirroring // gallery/llama3-instruct.yaml — stopwords defensively list <|im_end|> // first even though the assistant turn-end is actually <|eot_id|>. -// Exercises pickAssistantTurnEnd's template scan: the right token is +// Exercises PickAssistantTurnEnd's template scan: the right token is // the one that appears in chat_message, not the one at position 0. func writeLlama3StyleClassifierModel(modelDir, name string) { body := `name: ` + name + ` @@ -524,7 +524,7 @@ template: // writePartialClassifierModel writes a classifier model that has the // outer Chat template but no ChatMessage — exercises the -// newTemplateRenderer "refuse partial templating" branch, which makes +// NewTemplateRenderer "refuse partial templating" branch, which makes // buildClassifier reject the router with a missing-template error. func writePartialClassifierModel(modelDir, name string) { body := `name: ` + name + ` diff --git a/core/http/middleware/trace.go b/core/http/middleware/trace.go index 77180677a..1848fbeee 100644 --- a/core/http/middleware/trace.go +++ b/core/http/middleware/trace.go @@ -7,6 +7,7 @@ import ( "mime" "net" "net/http" + "path/filepath" "slices" "strconv" "sync" @@ -17,6 +18,7 @@ import ( "github.com/labstack/echo/v4" "github.com/mudler/LocalAI/core/application" "github.com/mudler/LocalAI/core/http/auth" + "github.com/mudler/LocalAI/core/trace/tracepersist" "github.com/mudler/xlog" ) @@ -58,34 +60,97 @@ type APIExchange struct { } var traceBuffer *circularbuffer.Queue[APIExchange] +var inFlightTraces = make(map[string]APIExchange) var mu sync.Mutex -var logChan = make(chan APIExchange, 100) -var tracingMaxItems int +var logChan = make(chan traceCommand, 100) var traceIDSeq atomic.Uint64 +var traceConsumerOnce sync.Once +var traceStore *tracepersist.Store[APIExchange] +var traceStoreKey string + +type traceCommand struct { + exchange *APIExchange + store *tracepersist.Store[APIExchange] + clear chan error +} func nextTraceID() string { return strconv.FormatUint(traceIDSeq.Add(1), 10) } -var doInitializeTracing = sync.OnceFunc(func() { - maxItems := tracingMaxItems +func initializeTracing(dataPath string, maxItems int) { if maxItems <= 0 { maxItems = 100 } + key := filepath.Join(dataPath, strconv.Itoa(maxItems)) mu.Lock() + if traceBuffer != nil && traceStoreKey == key { + mu.Unlock() + return + } + + var store *tracepersist.Store[APIExchange] + var restored []APIExchange + if dataPath != "" { + var err error + store, err = tracepersist.New[APIExchange](filepath.Join(dataPath, "traces", "api"), maxItems) + if err != nil { + xlog.Warn("Failed to initialize API trace persistence", "error", err) + } else if restored, err = store.Load(); err != nil { + xlog.Warn("Failed to restore API traces", "error", err) + store = nil + } + } traceBuffer = circularbuffer.New[APIExchange](maxItems) + for _, exchange := range restored { + traceBuffer.Enqueue(exchange) + advanceTraceID(&traceIDSeq, exchange.ID) + } + traceStore = store + traceStoreKey = key mu.Unlock() - go func() { - for exchange := range logChan { - mu.Lock() - if traceBuffer != nil { - traceBuffer.Enqueue(exchange) + traceConsumerOnce.Do(func() { + go func() { + for command := range logChan { + if command.clear != nil { + mu.Lock() + if traceBuffer != nil { + traceBuffer.Clear() + } + mu.Unlock() + var err error + if command.store != nil { + err = command.store.Clear() + } + command.clear <- err + continue + } + exchange := *command.exchange + mu.Lock() + delete(inFlightTraces, exchange.ID) + if traceBuffer != nil { + traceBuffer.Enqueue(exchange) + } + mu.Unlock() + if command.store != nil { + if err := command.store.Append(exchange.ID, exchange); err != nil { + xlog.Warn("Failed to persist API trace", "error", err) + } + } } - mu.Unlock() - } - }() -}) + }() + }) +} + +func advanceTraceID(seq *atomic.Uint64, id string) { + n, err := strconv.ParseUint(id, 10, 64) + if err != nil { + return + } + for current := seq.Load(); n > current && !seq.CompareAndSwap(current, n); current = seq.Load() { + } +} type bodyWriter struct { http.ResponseWriter @@ -145,11 +210,6 @@ func (w *bodyWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { return nil, nil, http.ErrNotSupported } -func initializeTracing(maxItems int) { - tracingMaxItems = maxItems - doInitializeTracing() -} - // sensitiveTraceHeaders is the set of header names whose values must not // land in the in-memory trace buffer. Keys are canonical — http.Header // stores them that way, so range yields canonical keys directly. @@ -175,14 +235,13 @@ func redactSensitiveHeaders(h http.Header) http.Header { // TraceMiddleware intercepts and logs JSON API requests and responses func TraceMiddleware(app *application.Application) echo.MiddlewareFunc { + initializeTracing(app.ApplicationConfig().DataPath, app.ApplicationConfig().TracingMaxItems) return func(next echo.HandlerFunc) echo.HandlerFunc { return func(c echo.Context) error { if !app.ApplicationConfig().EnableTracing { return next(c) } - initializeTracing(app.ApplicationConfig().TracingMaxItems) - ct, _, _ := mime.ParseMediaType(c.Request().Header.Get("Content-Type")) if ct != "application/json" { return next(c) @@ -204,6 +263,38 @@ func TraceMiddleware(app *application.Application) echo.MiddlewareFunc { // tens of MB, which then locks the admin Traces UI fetching the // JSON dump faster than the 5s auto-refresh. maxBodyBytes := app.ApplicationConfig().TracingMaxBodyBytes + requestHeaders := redactSensitiveHeaders(c.Request().Header) + requestBody, requestTruncated := truncateForTrace(body, maxBodyBytes) + exchange := APIExchange{ + ID: nextTraceID(), + Timestamp: startTime, + ClientIP: c.RealIP(), + UserAgent: c.Request().UserAgent(), + Request: APIExchangeRequest{ + Method: c.Request().Method, + Path: c.Path(), + Headers: &requestHeaders, + Body: &requestBody, + BodyTruncated: requestTruncated, + BodyBytes: len(body), + }, + } + if user := auth.GetUser(c); user != nil { + exchange.UserID = user.ID + exchange.UserName = user.Name + } + mu.Lock() + inFlightTraces[exchange.ID] = exchange + mu.Unlock() + queued := false + defer func() { + if queued { + return + } + mu.Lock() + delete(inFlightTraces, exchange.ID) + mu.Unlock() + }() // Wrap response writer to capture body resBody := new(bytes.Buffer) @@ -230,44 +321,27 @@ func TraceMiddleware(app *application.Application) echo.MiddlewareFunc { // the trace endpoint is admin-only but the buffer is also reachable // via any heap-dump-style introspection, and tokens shouldn't // outlive the request that carried them. - requestHeaders := redactSensitiveHeaders(c.Request().Header) - requestBody, requestTruncated := truncateForTrace(body, maxBodyBytes) responseHeaders := redactSensitiveHeaders(c.Response().Header()) responseBody := make([]byte, resBody.Len()) copy(responseBody, resBody.Bytes()) - exchange := APIExchange{ - ID: nextTraceID(), - Timestamp: startTime, - Duration: time.Since(startTime), - ClientIP: c.RealIP(), - UserAgent: c.Request().UserAgent(), - Request: APIExchangeRequest{ - Method: c.Request().Method, - Path: c.Path(), - Headers: &requestHeaders, - Body: &requestBody, - BodyTruncated: requestTruncated, - BodyBytes: len(body), - }, - Response: APIExchangeResponse{ - Status: status, - Headers: &responseHeaders, - Body: &responseBody, - BodyTruncated: mw.truncated, - BodyBytes: mw.totalBytes, - }, + exchange.Duration = time.Since(startTime) + exchange.Response = APIExchangeResponse{ + Status: status, + Headers: &responseHeaders, + Body: &responseBody, + BodyTruncated: mw.truncated, + BodyBytes: mw.totalBytes, } if handlerErr != nil { exchange.Error = handlerErr.Error() } - if user := auth.GetUser(c); user != nil { - exchange.UserID = user.ID - exchange.UserName = user.Name - } - + mu.Lock() + store := traceStore + mu.Unlock() select { - case logChan <- exchange: + case logChan <- traceCommand{exchange: &exchange, store: store}: + queued = true default: xlog.Warn("Trace channel full, dropping trace") } @@ -285,6 +359,10 @@ func GetTraces() []APIExchange { return []APIExchange{} } traces := traceBuffer.Values() + for _, exchange := range inFlightTraces { + exchange.Duration = time.Since(exchange.Timestamp) + traces = append(traces, exchange) + } mu.Unlock() slices.SortFunc(traces, func(a, b APIExchange) int { @@ -348,6 +426,18 @@ func window[T any](s []T, offset, limit int) []T { // ClearTraces clears the in-memory logs func ClearTraces() { + mu.Lock() + store := traceStore + initialized := traceBuffer != nil + mu.Unlock() + if initialized { + done := make(chan error, 1) + logChan <- traceCommand{store: store, clear: done} + if err := <-done; err != nil { + xlog.Warn("Failed to clear persisted API traces", "error", err) + } + return + } mu.Lock() if traceBuffer != nil { traceBuffer.Clear() diff --git a/core/http/middleware/trace_live_test.go b/core/http/middleware/trace_live_test.go new file mode 100644 index 000000000..ac81f3122 --- /dev/null +++ b/core/http/middleware/trace_live_test.go @@ -0,0 +1,108 @@ +// SPDX-License-Identifier: MIT + +package middleware + +import ( + "net/http" + "net/http/httptest" + "time" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/application" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/system" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("live API traces", func() { + newApp := func(root string) *application.Application { + app, err := application.New( + config.EnableTracing, + config.WithDataPath(root), + config.WithDisableLocalAIAssistant(true), + config.WithDisableStats(true), + config.WithSystemState(&system.SystemState{ + Model: system.Model{ModelsPath: root}, + Backend: system.Backend{BackendsPath: root}, + }), + ) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { Expect(app.Shutdown()).To(Succeed()) }) + ClearTraces() + return app + } + + It("lists a request while its handler is still running", func() { + root := GinkgoT().TempDir() + app := newApp(root) + + started := make(chan struct{}) + release := make(chan struct{}) + DeferCleanup(func() { + select { + case <-release: + default: + close(release) + } + }) + handler := TraceMiddleware(app)(func(c echo.Context) error { + close(started) + <-release + return c.NoContent(http.StatusNoContent) + }) + + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/slow", http.NoBody) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + ctx := e.NewContext(req, rec) + ctx.SetPath("/slow") + done := make(chan error, 1) + go func() { + done <- handler(ctx) + }() + <-started + + var running APIExchange + Eventually(func() bool { + traces := GetTraces() + if len(traces) != 1 { + return false + } + running = traces[0] + return running.Request.Path == "/slow" + }).Should(BeTrue()) + Expect(running.Response.Status).To(Equal(0)) + Expect(running.Duration).To(BeNumerically(">", 0)) + + close(release) + Expect(<-done).To(Succeed()) + Eventually(func() []APIExchange { return GetTraces() }).Should(ConsistOf( + And( + HaveField("ID", running.ID), + HaveField("Response.Status", http.StatusNoContent), + HaveField("Duration", BeNumerically(">", time.Duration(0))), + ), + )) + }) + + It("removes an in-flight trace when the handler panics", func() { + app := newApp(GinkgoT().TempDir()) + handler := TraceMiddleware(app)(func(echo.Context) error { + panic("handler panic") + }) + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/panic", http.NoBody) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + ctx := e.NewContext(req, httptest.NewRecorder()) + ctx.SetPath("/panic") + + func() { + defer func() { _ = recover() }() + _ = handler(ctx) + }() + + Expect(GetTraces()).To(BeEmpty()) + }) +}) diff --git a/core/http/middleware/trace_persistence_test.go b/core/http/middleware/trace_persistence_test.go new file mode 100644 index 000000000..16e943093 --- /dev/null +++ b/core/http/middleware/trace_persistence_test.go @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: MIT + +package middleware + +import ( + "path/filepath" + "strconv" + "time" + + "github.com/mudler/LocalAI/core/trace/tracepersist" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("API trace persistence", func() { + It("restores before request execution and advances the ID sequence", func() { + dataPath := GinkgoT().TempDir() + store, err := tracepersist.New[APIExchange](filepath.Join(dataPath, "traces", "api"), 4) + Expect(err).NotTo(HaveOccurred()) + Expect(store.Append("41", APIExchange{ID: "41", Timestamp: time.Now()})).To(Succeed()) + + initializeTracing(dataPath, 4) + + Expect(GetTraces()).To(ContainElement(HaveField("ID", "41"))) + id, err := strconv.ParseUint(nextTraceID(), 10, 64) + Expect(err).NotTo(HaveOccurred()) + Expect(id).To(BeNumerically(">", 41)) + }) + + It("serializes clear behind records already queued for persistence", func() { + dataPath := GinkgoT().TempDir() + initializeTracing(dataPath, 64) + for i := range 50 { + exchange := APIExchange{ID: strconv.Itoa(i + 1), Timestamp: time.Now()} + logChan <- traceCommand{exchange: &exchange, store: traceStore} + } + + ClearTraces() + + store, err := tracepersist.New[APIExchange](filepath.Join(dataPath, "traces", "api"), 64) + Expect(err).NotTo(HaveOccurred()) + records, err := store.Load() + Expect(err).NotTo(HaveOccurred()) + Expect(records).To(BeEmpty()) + Expect(GetTraces()).To(BeEmpty()) + }) +}) diff --git a/core/http/middleware/trace_summary.go b/core/http/middleware/trace_summary.go new file mode 100644 index 000000000..f5d8dbf39 --- /dev/null +++ b/core/http/middleware/trace_summary.go @@ -0,0 +1,109 @@ +// SPDX-License-Identifier: MIT + +package middleware + +import ( + "math" + "slices" + "time" +) + +// TraceSummary is the counted view of the trace buffer. +// +// It exists so a caller that wants "how many, how many failed, how slow" does +// not have to fetch every exchange and count them in the browser. The Operate +// overview needs exactly those three numbers, and the trace list is capped in +// the thousands, so shipping it across the wire to produce a single integer is +// waste that grows with the buffer. +type TraceSummary struct { + Total int `json:"total"` + Errors int `json:"errors"` + P95Millis int64 `json:"p95_ms"` + WindowHours int `json:"window_hours"` + Buckets []TraceBucket `json:"buckets"` +} + +// TraceBucket is one column of a sparkline: oldest first, so the series reads +// left to right the way a chart is drawn. +type TraceBucket struct { + Start time.Time `json:"start"` + Count int `json:"count"` + Errors int `json:"errors"` +} + +// GetTracesSummary counts the buffered exchanges over the given window. +func GetTracesSummary(window time.Duration, buckets int) TraceSummary { + return summarize(GetTraces(), window, buckets) +} + +func summarize(traces []APIExchange, window time.Duration, buckets int) TraceSummary { + if buckets < 1 { + buckets = 1 + } + now := time.Now() + cutoff := now.Add(-window) + + summary := TraceSummary{ + WindowHours: int(window.Hours()), + // Never nil: a nil slice serialises as null and breaks .map() on the + // other side, which is a silent runtime error rather than an empty chart. + Buckets: make([]TraceBucket, buckets), + } + + bucketWidth := window / time.Duration(buckets) + for i := range summary.Buckets { + summary.Buckets[i].Start = cutoff.Add(time.Duration(i) * bucketWidth) + } + + durations := make([]time.Duration, 0, len(traces)) + for _, t := range traces { + if t.Timestamp.Before(cutoff) { + continue + } + summary.Total++ + failed := isFailure(t) + if failed { + summary.Errors++ + } + durations = append(durations, t.Duration) + + // Clamp rather than skip: a request timestamped a hair in the future + // (clock skew, or arriving mid-call) still belongs in the newest column. + idx := int(t.Timestamp.Sub(cutoff) / bucketWidth) + if idx >= buckets { + idx = buckets - 1 + } + if idx < 0 { + idx = 0 + } + summary.Buckets[idx].Count++ + if failed { + summary.Buckets[idx].Errors++ + } + } + + summary.P95Millis = percentileMillis(durations, 0.95) + return summary +} + +// A 4xx is the caller getting it wrong, which is not the installation being +// unhealthy. Only 5xx and a transport-level error count against the runtime. +func isFailure(t APIExchange) bool { + return t.Error != "" || t.Response.Status >= 500 +} + +func percentileMillis(durations []time.Duration, p float64) int64 { + if len(durations) == 0 { + return 0 + } + slices.Sort(durations) + // Nearest-rank: the smallest value at or above the pth percentile. + rank := int(math.Ceil(p*float64(len(durations)))) - 1 + if rank < 0 { + rank = 0 + } + if rank >= len(durations) { + rank = len(durations) - 1 + } + return durations[rank].Milliseconds() +} diff --git a/core/http/middleware/trace_summary_test.go b/core/http/middleware/trace_summary_test.go new file mode 100644 index 000000000..afd888c2e --- /dev/null +++ b/core/http/middleware/trace_summary_test.go @@ -0,0 +1,79 @@ +// SPDX-License-Identifier: MIT + +package middleware + +import ( + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("API trace summary", func() { + exchange := func(age time.Duration, status int, dur time.Duration) APIExchange { + return APIExchange{ + Timestamp: time.Now().Add(-age), + Duration: dur, + Response: APIExchangeResponse{Status: status}, + } + } + + It("counts only what falls inside the window", func() { + traces := []APIExchange{ + exchange(1*time.Hour, 200, 10*time.Millisecond), + exchange(2*time.Hour, 200, 10*time.Millisecond), + // Older than the window: must not be counted at all. + exchange(48*time.Hour, 500, 10*time.Millisecond), + } + s := summarize(traces, 24*time.Hour, 6) + Expect(s.Total).To(Equal(2)) + Expect(s.Errors).To(BeZero()) + }) + + It("treats 5xx and a transport error as failures, but not 4xx", func() { + traces := []APIExchange{ + exchange(time.Minute, 500, time.Millisecond), + exchange(time.Minute, 503, time.Millisecond), + // A client sending a bad request is not the server failing. + exchange(time.Minute, 404, time.Millisecond), + exchange(time.Minute, 200, time.Millisecond), + } + traces[3].Error = "connection reset" + + s := summarize(traces, 24*time.Hour, 6) + Expect(s.Total).To(Equal(4)) + Expect(s.Errors).To(Equal(3)) + }) + + It("reports p95 as a real percentile rather than the slowest request", func() { + traces := make([]APIExchange, 0, 100) + for i := 1; i <= 100; i++ { + traces = append(traces, exchange(time.Minute, 200, time.Duration(i)*time.Millisecond)) + } + s := summarize(traces, 24*time.Hour, 6) + // 95th of 1..100ms, not the 100ms max. + Expect(s.P95Millis).To(BeNumerically("~", 95, 1)) + }) + + It("buckets oldest-first so a sparkline reads left to right", func() { + traces := []APIExchange{ + exchange(30*time.Minute, 200, time.Millisecond), + exchange(30*time.Minute, 200, time.Millisecond), + exchange(5*time.Hour, 200, time.Millisecond), + } + s := summarize(traces, 6*time.Hour, 6) + Expect(s.Buckets).To(HaveLen(6)) + Expect(s.Buckets[0].Count).To(Equal(1), "the 5h-old request lands in the first bucket") + Expect(s.Buckets[5].Count).To(Equal(2), "the recent pair lands in the last") + }) + + It("returns an empty, non-nil summary when nothing has been traced", func() { + s := summarize(nil, 24*time.Hour, 6) + Expect(s.Total).To(BeZero()) + Expect(s.Errors).To(BeZero()) + Expect(s.P95Millis).To(BeZero()) + // A nil slice serialises as null and breaks .map() in the browser. + Expect(s.Buckets).NotTo(BeNil()) + Expect(s.Buckets).To(HaveLen(6)) + }) +}) diff --git a/core/http/react-ui/e2e/activity-page.spec.js b/core/http/react-ui/e2e/activity-page.spec.js new file mode 100644 index 000000000..09f9b46fe --- /dev/null +++ b/core/http/react-ui/e2e/activity-page.spec.js @@ -0,0 +1,451 @@ +import { test, expect } from './coverage-fixtures.js' + +const stub = (page, { operations = [], history = [] } = {}) => Promise.all([ + page.route('**/api/operations', (route) => route.fulfill({ + contentType: 'application/json', + body: JSON.stringify({ operations }), + })), + page.route('**/api/operations/history', (route) => route.fulfill({ + contentType: 'application/json', + body: JSON.stringify({ operations: history }), + })), +]) + +test('lists live operations and cancels one from a labelled button', async ({ page }) => { + await stub(page, { + operations: [{ + id: 'gemma-3-27b-it', + name: 'gemma-3-27b-it', + jobID: 'job-gemma', + progress: 22, + taskType: 'installation', + isBackend: false, + isQueued: false, + isDeletion: false, + cancellable: true, + phase: 'downloading', + }], + }) + + let cancelledPath = '' + await page.route('**/api/operations/job-gemma/cancel', (route) => { + cancelledPath = new URL(route.request().url()).pathname + return route.fulfill({ contentType: 'application/json', body: '{}' }) + }) + + await page.goto('/app/activity') + + const card = page.locator('.operation-card').filter({ hasText: 'gemma-3-27b-it' }) + await expect(card).toBeVisible() + await expect(card).toContainText('22%') + + await card.locator('.operation-card__cancel').click() + expect(cancelledPath).toBe('/api/operations/job-gemma/cancel') +}) + +test('pauses a model download without invoking destructive cancel', async ({ page }) => { + await stub(page, { + operations: [{ + id: 'gemma-3-27b-it', + name: 'gemma-3-27b-it', + jobID: 'job-gemma', + progress: 22, + taskType: 'installation', + isBackend: false, + isQueued: false, + isDeletion: false, + cancellable: true, + phase: 'downloading', + }], + }) + + const requests = [] + await page.route('**/api/operations/job-gemma/pause', (route) => { + requests.push(new URL(route.request().url()).pathname) + return route.fulfill({ contentType: 'application/json', body: '{}' }) + }) + await page.route('**/api/operations/job-gemma/cancel', (route) => { + requests.push(new URL(route.request().url()).pathname) + return route.fulfill({ contentType: 'application/json', body: '{}' }) + }) + + await page.goto('/app/activity') + + const card = page.locator('.operation-card').filter({ hasText: 'gemma-3-27b-it' }) + await card.locator('.operation-card__pause').click() + + await expect.poll(() => requests).toEqual(['/api/operations/job-gemma/pause']) +}) + +test('separates an unacknowledged failure from the record', async ({ page }) => { + await stub(page, { + operations: [{ + id: 'sherpa-onnx', + name: 'sherpa-onnx', + jobID: 'job-sherpa', + progress: 0, + taskType: 'installation', + isBackend: true, + isQueued: false, + isDeletion: false, + cancellable: false, + error: 'no space left on device', + }], + history: [{ + id: 'bark-cpp', + name: 'bark-cpp', + jobID: 'job-bark', + isBackend: true, + taskType: 'installation', + outcome: 'failed', + error: 'checksum mismatch', + startedAt: '2026-07-28T13:40:00Z', + finishedAt: '2026-07-28T13:41:00Z', + }], + }) + + await page.goto('/app/activity') + + // Live and unacknowledged: a card that needs a decision. + await expect(page.locator('.operation-card--error')).toContainText('sherpa-onnx') + // Dismissed earlier: a row in the record. + await expect(page.locator('.activity-row')).toContainText('bark-cpp') +}) + +test('a failure never appears in both In progress and Needs attention', async ({ page }) => { + // Section membership has to be unambiguous: the same job showing twice makes + // the two failure paths (retry / dismiss) impossible to reason about. + await stub(page, { + operations: [ + { + id: 'model-a', + name: 'model-a', + jobID: 'job-a', + progress: 40, + taskType: 'installation', + isBackend: false, + isQueued: false, + isDeletion: false, + cancellable: true, + }, + { + id: 'sherpa-onnx', + name: 'sherpa-onnx', + jobID: 'job-sherpa', + progress: 0, + taskType: 'installation', + isBackend: true, + isQueued: false, + isDeletion: false, + cancellable: false, + error: 'no space left on device', + }, + ], + }) + + await page.goto('/app/activity') + + await expect(page.locator('.operation-card')).toHaveCount(2) + await expect(page.locator('.operation-card').filter({ hasText: 'sherpa-onnx' })).toHaveCount(1) +}) + +test('retrying a failed backend install dismisses it before reinstalling', async ({ page }) => { + // Order is load-bearing: a bare reinstall overwrites the opcache entry + // without going through recordTerminal, so the failure would never reach the + // record. + await stub(page, { + operations: [{ + id: 'sherpa-onnx', + name: 'sherpa-onnx', + fullName: 'localai@sherpa-onnx', + jobID: 'job-sherpa', + progress: 0, + taskType: 'installation', + isBackend: true, + isQueued: false, + isDeletion: false, + cancellable: false, + error: 'no space left on device', + }], + }) + + const calls = [] + await page.route('**/api/operations/job-sherpa/dismiss', (route) => { + calls.push('dismiss') + return route.fulfill({ contentType: 'application/json', body: '{}' }) + }) + await page.route('**/api/backends/install/**', (route) => { + calls.push(new URL(route.request().url()).pathname) + return route.fulfill({ contentType: 'application/json', body: '{}' }) + }) + + await page.goto('/app/activity') + + await page.locator('.operation-card__retry').click() + + await expect.poll(() => calls).toEqual(['dismiss', '/api/backends/install/localai@sherpa-onnx']) +}) + +test('retry dismisses the job it was pressed on, not another sharing its id', async ({ page }) => { + // /api/operations strips the "node::" prefix, so a local install and a + // node-scoped install of one backend arrive with the same id and different + // jobIDs. Dismissing by id retired whichever came first, which both left the + // acted-on failure live and silently retired an unrelated one. + const failed = (over) => ({ + id: 'sherpa-onnx', + name: 'sherpa-onnx', + fullName: 'sherpa-onnx', + progress: 0, + taskType: 'installation', + isBackend: true, + isQueued: false, + isDeletion: false, + cancellable: false, + error: 'no space left on device', + ...over, + }) + await stub(page, { + operations: [ + failed({ jobID: 'job-local' }), + failed({ jobID: 'job-node', nodeID: 'node-1' }), + ], + }) + + const calls = [] + await page.route('**/api/operations/*/dismiss', (route) => { + calls.push(new URL(route.request().url()).pathname) + return route.fulfill({ contentType: 'application/json', body: '{}' }) + }) + await page.route('**/api/nodes/*/backends/install', (route) => { + calls.push(new URL(route.request().url()).pathname) + return route.fulfill({ contentType: 'application/json', body: '{}' }) + }) + + await page.goto('/app/activity') + + // Nothing on screen tells the two cards apart, which is the point: they + // share a name and an id, and only the jobID behind each one differs. The + // node-scoped job is second in the payload, so it is the second card. + await expect(page.locator('.operation-card')).toHaveCount(2) + await page.locator('.operation-card').nth(1).locator('.operation-card__retry').click() + + await expect.poll(() => calls).toEqual([ + '/api/operations/job-node/dismiss', + '/api/nodes/node-1/backends/install', + ]) +}) + +test('the dismiss control also acts on the job it belongs to', async ({ page }) => { + // Same hazard as retry: the card's X passed the display id too. + const failed = (over) => ({ + id: 'sherpa-onnx', + name: 'sherpa-onnx', + fullName: 'sherpa-onnx', + progress: 0, + taskType: 'installation', + isBackend: true, + isQueued: false, + isDeletion: false, + cancellable: false, + error: 'no space left on device', + ...over, + }) + await stub(page, { + operations: [ + failed({ jobID: 'job-local' }), + failed({ jobID: 'job-node', nodeID: 'node-1' }), + ], + }) + + const dismissed = [] + await page.route('**/api/operations/*/dismiss', (route) => { + dismissed.push(new URL(route.request().url()).pathname) + return route.fulfill({ contentType: 'application/json', body: '{}' }) + }) + + await page.goto('/app/activity') + + await expect(page.locator('.operation-card')).toHaveCount(2) + await page.locator('.operation-card').nth(1).locator('.operation-card__hide').click() + + await expect.poll(() => dismissed).toEqual(['/api/operations/job-node/dismiss']) +}) + +test('a filter matching nothing does not claim the instance is empty', async ({ page }) => { + // Three model records on file: telling the user nothing has ever run, while + // the header counts those same three, is simply false. + await stub(page, { + history: [1, 2, 3].map((n) => ({ + id: `model-${n}`, + name: `model-${n}`, + jobID: `job-${n}`, + isBackend: false, + taskType: 'installation', + outcome: 'completed', + startedAt: '2026-07-28T13:40:00Z', + finishedAt: '2026-07-28T13:40:20Z', + })), + }) + + await page.goto('/app/activity') + await expect(page.locator('.activity-row')).toHaveCount(3) + + await page.locator('.activity-chip', { hasText: 'Backends' }).click() + + await expect(page.locator('.activity-empty--filtered')).toBeVisible() + await expect(page.locator('.activity-empty')).not.toContainText('No operations since startup') + // And the way back out is on screen. + await page.locator('.activity-empty--filtered button').click() + await expect(page.locator('.activity-row')).toHaveCount(3) +}) + +test('the summary drops a zero count instead of reporting it', async ({ page }) => { + await stub(page, { + operations: [{ + id: 'model-a', + name: 'model-a', + jobID: 'job-a', + progress: 40, + taskType: 'installation', + isBackend: false, + isQueued: false, + isDeletion: false, + cancellable: true, + }], + }) + + await page.goto('/app/activity') + + const supporting = page.locator('.page-header__supporting') + await expect(supporting).toHaveText('1 operation running.') + await expect(supporting).not.toContainText('0') +}) + +test('a cancelled deletion reports the cancellation, not a removal', async ({ page }) => { + await stub(page, { + history: [{ + id: 'model-a', + name: 'model-a', + jobID: 'job-a', + isBackend: false, + taskType: 'deletion', + outcome: 'cancelled', + startedAt: '2026-07-28T13:40:00Z', + finishedAt: '2026-07-28T13:40:02Z', + }], + }) + + await page.goto('/app/activity') + + await expect(page.locator('.activity-row')).toContainText('cancelled') + await expect(page.locator('.activity-row')).not.toContainText('removed') +}) + +test('an implausible or zero duration never reaches the row', async ({ page }) => { + await stub(page, { + history: [ + { + id: 'zero-span', + name: 'zero-span', + jobID: 'job-zero', + isBackend: false, + taskType: 'installation', + outcome: 'completed', + // recordTerminal seeds StartedAt = FinishedAt and only overwrites it + // with a real stamp, so this is an ordinary arrival. + startedAt: '2026-07-28T13:40:00Z', + finishedAt: '2026-07-28T13:40:00Z', + }, + { + id: 'zero-stamp', + name: 'zero-stamp', + jobID: 'job-stamp', + isBackend: false, + taskType: 'installation', + outcome: 'completed', + startedAt: '0001-01-01T00:00:00Z', + finishedAt: '2026-07-28T13:41:00Z', + }, + ], + }) + + await page.goto('/app/activity') + + const zeroSpan = page.locator('.activity-row').filter({ hasText: 'zero-span' }) + await expect(zeroSpan).toContainText('installed in < 1s') + + // A zero-value Go stamp is not a duration. The row says what happened and + // stops, rather than stating a span of millennia as fact. + const zeroStamp = page.locator('.activity-row').filter({ hasText: 'zero-stamp' }) + await expect(zeroStamp).toContainText('installed') + await expect(zeroStamp).not.toContainText('installed in') +}) + +test('a failed removal offers no retry, because retry only means install', async ({ page }) => { + await stub(page, { + operations: [{ + id: 'model-a', + name: 'model-a', + fullName: 'model-a', + jobID: 'job-a', + progress: 0, + taskType: 'deletion', + isBackend: false, + isQueued: false, + isDeletion: true, + cancellable: false, + error: 'file is busy', + }], + }) + + await page.goto('/app/activity') + + await expect(page.locator('.operation-card--error')).toBeVisible() + await expect(page.locator('.operation-card__retry')).toHaveCount(0) + // And it must not claim an install was attempted. + await expect(page.locator('.operation-card--error')).not.toContainText('install') +}) + +test('filters the record down to backends', async ({ page }) => { + await stub(page, { + history: [ + { + id: 'gemma-3-27b-it', + name: 'gemma-3-27b-it', + jobID: 'job-gemma', + isBackend: false, + taskType: 'installation', + outcome: 'completed', + startedAt: '2026-07-28T13:40:00Z', + finishedAt: '2026-07-28T13:41:30Z', + }, + { + id: 'bark-cpp', + name: 'bark-cpp', + jobID: 'job-bark', + isBackend: true, + taskType: 'installation', + outcome: 'completed', + startedAt: '2026-07-28T13:40:00Z', + finishedAt: '2026-07-28T13:40:20Z', + }, + ], + }) + + await page.goto('/app/activity') + await expect(page.locator('.activity-row')).toHaveCount(2) + + await page.locator('.activity-chip', { hasText: 'Backends' }).click() + + await expect(page.locator('.activity-row')).toHaveCount(1) + await expect(page.locator('.activity-row')).toContainText('bark-cpp') +}) + +test('shows the empty state when nothing has run', async ({ page }) => { + await stub(page) + + await page.goto('/app/activity') + + await expect(page.locator('.page-title')).toBeVisible() + await expect(page.locator('.activity-empty')).toBeVisible() +}) diff --git a/core/http/react-ui/e2e/admin-console.spec.js b/core/http/react-ui/e2e/admin-console.spec.js index 1a039eba3..6bd459feb 100644 --- a/core/http/react-ui/e2e/admin-console.spec.js +++ b/core/http/react-ui/e2e/admin-console.spec.js @@ -5,7 +5,9 @@ test.describe('Admin console', () => { await page.goto('/app/backends') const rail = page.locator('.console-rail') await expect(rail).toBeVisible() - for (const group of ['Inference', 'Cluster', 'Observability', 'Access', 'System']) { + // Four groups since the overview landed: Inference folded into Runtime + // (both are "the runtime right now"), Access and System into Administration. + for (const group of ['Runtime', 'Cluster', 'Observability', 'Administration']) { await expect(rail.locator('.console-group-title', { hasText: group })).toBeVisible() } }) diff --git a/core/http/react-ui/e2e/alias-template.spec.js b/core/http/react-ui/e2e/alias-template.spec.js index f3b1a0ca0..e9b13ba61 100644 --- a/core/http/react-ui/e2e/alias-template.spec.js +++ b/core/http/react-ui/e2e/alias-template.spec.js @@ -69,9 +69,9 @@ test.describe('Manage - alias badge', () => { test('renders a read-only alias -> target badge on aliased rows', async ({ page }) => { await page.goto('/app/manage') - await expect(page.locator('.table')).toBeVisible({ timeout: 10_000 }) - - // The aliased row shows the target; the plain model row does not. + // The badge moved off the row and into the pane: it is a fact about the + // model, and the rail line is spent on state. + await page.locator('[data-entity="gpt-4"]').click() await expect(page.getByText('alias -> fast-llm')).toBeVisible({ timeout: 10_000 }) }) }) diff --git a/core/http/react-ui/e2e/backends-management.spec.js b/core/http/react-ui/e2e/backends-management.spec.js index 11b179189..34408c92b 100644 --- a/core/http/react-ui/e2e/backends-management.spec.js +++ b/core/http/react-ui/e2e/backends-management.spec.js @@ -1,6 +1,9 @@ import { test, expect } from './coverage-fixtures.js' // Backends admin page (src/pages/Backends.jsx). +const PANE = '[data-testid="backends-pane"]' +const railItem = (page, name) => page.locator(`[data-entity="${name}"]`) + test.describe('Backends management page', () => { test.beforeEach(async ({ page }) => { await page.goto('/app/backends') @@ -49,11 +52,14 @@ test.describe('Backends management page - Markdown descriptions', () => { }) }) await page.goto('/app/backends') - await expect(page.locator('th', { hasText: 'Description' })).toBeVisible({ timeout: 10_000 }) + // Rendered means the rail has entries. The old gate waited on a column + // header, and there are no columns now. + await expect(railItem(page, 'markdown-backend')).toBeVisible({ timeout: 10_000 }) }) - test('table cell shows the description as clean text, not raw Markdown', async ({ page }) => { - const cell = page.locator('tr', { hasText: 'markdown-backend' }).locator('span[title]', { hasText: 'InsightFace' }) + test('the pane lede shows the description as clean text, not raw Markdown', async ({ page }) => { + await railItem(page, 'markdown-backend').click() + const cell = page.locator('.detail-pane__lede') await expect(cell).toHaveText(STRIPPED_DESCRIPTION) // The syntax itself must be gone, not merely rendered somewhere. @@ -65,15 +71,77 @@ test.describe('Backends management page - Markdown descriptions', () => { await expect(cell.locator('h1')).toHaveCount(0) }) - test('title tooltip carries the stripped text, not raw Markdown', async ({ page }) => { - const cell = page.locator('tr', { hasText: 'markdown-backend' }).locator('span[title]', { hasText: 'InsightFace' }) - - await expect(cell).toHaveAttribute('title', STRIPPED_DESCRIPTION) + test("the lede's tooltip carries the stripped text, not raw Markdown", async ({ page }) => { + await railItem(page, 'markdown-backend').click() + await expect(page.locator('.detail-pane__lede')).toHaveAttribute('title', STRIPPED_DESCRIPTION) }) - test('a backend with no description still shows the placeholder', async ({ page }) => { - const row = page.locator('tr', { hasText: 'plain-backend' }) - - await expect(row.locator('span[title=""]')).toHaveText('-') + test('a backend with no description renders no lede rather than a blank one', async ({ page }) => { + // The table needed a placeholder because an empty cell in a grid of full + // ones reads as a fault. The pane has no grid to keep aligned, so it omits + // the line - but must never print "undefined". + await railItem(page, 'plain-backend').click() + await expect(page.locator(PANE)).toContainText('plain-backend') + await expect(page.locator('.detail-pane__lede')).toHaveCount(0) + await expect(page.locator(PANE)).not.toContainText('undefined') + }) +}) + +test.describe('Backends gallery - split view', () => { + test.beforeEach(async ({ page }) => { + await page.route('**/api/backends*', (route) => { + route.fulfill({ + contentType: 'application/json', + body: JSON.stringify({ + backends: [ + { name: 'llama-cpp', description: 'GGUF inference', installed: true, version: '1.52.0', license: 'MIT', tags: ['chat'] }, + { name: 'whisper', description: 'Speech to text', installed: true, version: '1.8.2', license: 'MIT', tags: ['transcript'] }, + { name: 'diffusers', description: 'Image generation', installed: false, license: 'Apache-2.0', tags: ['image'] }, + ], + }), + }) + }) + await page.goto('/app/backends') + await expect(railItem(page, 'llama-cpp')).toBeVisible({ timeout: 10_000 }) + }) + + test('the gallery renders no table', async ({ page }) => { + await expect(page.locator('[data-testid="backends"]')).toBeVisible() + await expect(page.locator('table thead th')).toHaveCount(0) + }) + + test('with nothing selected the pane describes the host', async ({ page }) => { + await expect(page.locator(PANE)).toContainText('This host') + await expect(page.locator('[data-testid="backends-back"]')).toHaveCount(0) + }) + + test('choosing a backend turns the pane into its detail, and back returns', async ({ page }) => { + await railItem(page, 'llama-cpp').click() + await expect(page.locator(PANE)).toContainText('llama-cpp') + await expect(page.locator(PANE)).toContainText('MIT') + await expect(page.locator(PANE)).not.toContainText('This host') + + await page.locator('[data-testid="backends-back"]').click() + await expect(page.locator(PANE)).toContainText('This host') + }) + + test('the selection lives in the URL and survives a reload', async ({ page }) => { + await railItem(page, 'whisper').click() + await expect(page).toHaveURL(/[?&]backend=whisper/) + await page.reload() + await expect(railItem(page, 'whisper')).toBeVisible({ timeout: 10_000 }) + await expect(page.locator('[data-testid="backends-back"]')).toBeVisible() + }) + + + test('the rail groups while browsing and flattens on a query', async ({ page }) => { + await expect(page.locator('[data-testid^="backends-rail-group-"]').first()).toBeVisible() + await page.locator('input[placeholder*="Search backends"]').fill('llama') + await expect(page.locator('[data-testid^="backends-rail-group-"]')).toHaveCount(0) + }) + + test('an installed backend states its version, an absent one says so', async ({ page }) => { + await expect(railItem(page, 'llama-cpp')).toContainText('v1.52.0') + await expect(railItem(page, 'diffusers')).toContainText('not installed') }) }) diff --git a/core/http/react-ui/e2e/backends-notice.spec.js b/core/http/react-ui/e2e/backends-notice.spec.js new file mode 100644 index 000000000..cf7b5cbeb --- /dev/null +++ b/core/http/react-ui/e2e/backends-notice.spec.js @@ -0,0 +1,23 @@ +import { test, expect } from './coverage-fixtures.js' + +// A notice is a hairline with a coloured left edge, not a filled panel. A tint +// makes every notice shout at the weight of an error, which is how notices stop +// being read — and it is the same treatment the Operate overview uses for the +// rows that want a decision. + +test('the backends notice is an edge, not a filled card', async ({ page }) => { + // The upgrade banner is the notice worth pinning, so make one exist. + await page.route('**/api/backends/upgrades', route => route.fulfill({ + json: { 'llama-cpp': { backend_name: 'llama-cpp', installed_version: '0.9.4', available_version: '0.9.7' } }, + })) + await page.goto('/app/backends') + const notice = page.locator('.bk-notice', { hasText: /update/i }).first() + await expect(notice).toBeVisible() + const s = await notice.evaluate(el => { + const cs = getComputedStyle(el) + return { bg: cs.backgroundColor, left: parseFloat(cs.borderLeftWidth), top: parseFloat(cs.borderTopWidth) } + }) + expect(s.bg).toMatch(/rgba\(0, 0, 0, 0\)|transparent/) + expect(s.left).toBeGreaterThanOrEqual(3) + expect(s.top).toBeLessThanOrEqual(1) +}) diff --git a/core/http/react-ui/e2e/chat-transcript.spec.js b/core/http/react-ui/e2e/chat-transcript.spec.js new file mode 100644 index 000000000..8164e8652 --- /dev/null +++ b/core/http/react-ui/e2e/chat-transcript.spec.js @@ -0,0 +1,55 @@ +import { test, expect } from './coverage-fixtures.js' + +// Chat reads as a transcript rather than a bubble thread (mock 04). + +const CHAT = { + chats: [{ + id: 'c1', name: 'Transcript', model: 'mock-model', + history: [ + { role: 'user', content: 'Which backends do I have?' }, + { role: 'assistant', content: 'Seven are installed.' }, + ], + }], + activeChatId: 'c1', +} + +test.describe('Chat transcript', () => { + test.beforeEach(async ({ page }) => { + await page.addInitScript(chat => { + localStorage.setItem('localai_chats_data', JSON.stringify(chat)) + }, CHAT) + await page.goto('/app/chat') + }) + + test('neither role is a filled, rounded bubble', async ({ page }) => { + const user = page.locator('.chat-message-user .chat-message-content').first() + await expect(user).toBeVisible() + const cs = await user.evaluate(el => { + const s = getComputedStyle(el) + return { radius: s.borderTopLeftRadius, shadow: s.boxShadow } + }) + // A rounded filled bubble carries the speaker in shape and side; a + // transcript carries it in words, which survives being read aloud. + expect(cs.radius).toBe('0px') + expect(cs.shadow).toBe('none') + }) + + test('both turns run full width in one column, not left and right', async ({ page }) => { + const user = page.locator('.chat-message-user').first() + const assistant = page.locator('.chat-message-assistant').first() + const [u, a] = [await user.boundingBox(), await assistant.boundingBox()] + expect(Math.abs(u.x - a.x)).toBeLessThan(2) + }) + + test('every turn says who is speaking', async ({ page }) => { + await expect(page.locator('.chat-message-user .chat-message-model')).toHaveText('You') + await expect(page.locator('.chat-message-assistant .chat-message-model').first()) + .toHaveText('mock-model') + }) + + test('turns are separated by a rule', async ({ page }) => { + const border = await page.locator('.chat-message').first() + .evaluate(el => getComputedStyle(el).borderBottomStyle) + expect(border).toBe('solid') + }) +}) diff --git a/core/http/react-ui/e2e/chrome-audit.spec.js b/core/http/react-ui/e2e/chrome-audit.spec.js new file mode 100644 index 000000000..728701003 --- /dev/null +++ b/core/http/react-ui/e2e/chrome-audit.spec.js @@ -0,0 +1,50 @@ +import { test, expect } from './coverage-fixtures.js' + +// A standing guard against the two defects an earlier automated edit left +// scattered through the pages: icons stripped of their fa-* class (which render +// nothing at all), and controls left with the user agent's own chrome, which is +// a pale grey button on a dark ground. +const ROUTES = [ + '/app', '/app/chat', '/app/models', '/app/studio', '/app/talk', + '/app/agents', '/app/skills', '/app/collections', '/app/agent-jobs', + '/app/fine-tune', '/app/quantize', '/app/face', '/app/voice', + '/app/manage', '/app/backends', '/app/activity', '/app/operate', + '/app/settings', '/app/traces', '/app/usage', '/app/nodes', '/app/p2p', + '/app/voice-library', '/app/voice-library/new', '/app/account', +] + +test('no page renders a dead icon or a default-chrome control', async ({ page }) => { + // One test walks every route, so its budget has to scale with the list rather + // than sit on Playwright's per-test default of 30s. At 25 routes that default + // allows ~1.2s per navigation, which holds on a developer machine and does + // not on a loaded CI runner: the suite went red on the commit that added this + // spec, timing out mid-loop at waitForTimeout rather than at any single goto, + // which is what cumulative slowness looks like as opposed to one hung route. + // Six seconds a route absorbs a slow runner and still fails promptly if a + // route really does hang. + test.setTimeout(ROUTES.length * 6_000) + + const findings = [] + for (const route of ROUTES) { + await page.goto(route) + await page.waitForTimeout(400) + const found = await page.evaluate(() => { + const out = [] + for (const el of document.querySelectorAll('button, a')) { + if (el.getBoundingClientRect().width === 0) continue + const cs = getComputedStyle(el) + if (cs.borderTopStyle === 'outset' || cs.backgroundColor === 'rgb(239, 239, 239)') { + out.push(`default-chrome: "${(el.textContent || '').trim().slice(0, 24)}" [${el.className}]`) + } + } + for (const i of document.querySelectorAll('i')) { + if (!/\bfa-/.test((i.className || '').toString())) { + out.push(`dead-icon: [${i.className}]`) + } + } + return [...new Set(out)] + }) + for (const f of found) findings.push(`${route} — ${f}`) + } + expect(findings).toEqual([]) +}) diff --git a/core/http/react-ui/e2e/console-narrow.spec.js b/core/http/react-ui/e2e/console-narrow.spec.js new file mode 100644 index 000000000..2497c657b --- /dev/null +++ b/core/http/react-ui/e2e/console-narrow.spec.js @@ -0,0 +1,101 @@ +import { test, expect } from './coverage-fixtures.js' + +// Small-screen behaviour of the Operate console and the dashboard stat cards. +// +// Both defects here are about a narrow viewport but neither is only a narrow +// viewport problem: the stat cards were being laid out by the wrong rule at +// every width, and the rail's height was never bounded. + +test.describe('Operate console on a narrow screen', () => { + test('expanding the rail leaves the page still on screen', async ({ page }) => { + await page.setViewportSize({ width: 390, height: 800 }) + await page.goto('/app/manage') + + const toggle = page.locator('.console-rail-toggle') + await expect(toggle).toBeVisible() + await toggle.click() + await expect(page.locator('.console-rail-groups')).toBeVisible() + + // Thirteen destinations in one column is taller than a phone. If opening + // the menu pushes the page's own heading past the fold, the menu has + // replaced the page instead of annotating it. + // Manage titles itself with .view-bar__title rather than .page-title. + const heading = page.locator('.page-title, .view-bar__title').first() + const box = await heading.boundingBox() + expect(box).not.toBeNull() + expect(box.y).toBeLessThan(800) + }) + + test('the rail scrolls internally rather than growing without bound', async ({ page }) => { + await page.setViewportSize({ width: 390, height: 800 }) + await page.goto('/app/manage') + await page.locator('.console-rail-toggle').click() + + const groups = page.locator('.console-rail-groups') + await expect(groups).toBeVisible() + const height = await groups.evaluate(el => el.getBoundingClientRect().height) + expect(height).toBeLessThan(800) + }) +}) + +test.describe('Headline figures', () => { + // Host used shadowed StatCards; it now shares the Operate overview's hairline + // figure strip, so the guard is that its labels stay legible, not that it + // keeps a card gap. + for (const width of [768, 1024]) { + test(`Host figure labels are not clipped at ${width}px`, async ({ page }) => { + await page.setViewportSize({ width, height: 1000 }) + await page.goto('/app/manage') + const labels = page.locator('.stat-strip__label') + await expect(labels.first()).toBeVisible() + const clipped = await labels.evaluateAll(els => + els.filter(el => el.scrollWidth > el.clientWidth + 1).map(el => el.textContent)) + expect(clipped).toEqual([]) + }) + } + + test('a Host figure routes into the thing it counts', async ({ page }) => { + await page.setViewportSize({ width: 1280, height: 1000 }) + await page.goto('/app/manage') + const cell = page.locator('.stat-strip__cell').first() + await expect(cell).toBeVisible() + // A count is worth more when it is also the way to what it counted. + await expect(cell).toHaveJSProperty('tagName', 'BUTTON') + }) + + test('the figure strip keeps its height inside the flex column', async ({ page }) => { + // .page--app is a flex column whose split view takes flex:1, so a child + // with no intrinsic minimum gets shrunk to nothing. This strip did exactly + // that and rendered 2px tall with four invisible cells. + await page.setViewportSize({ width: 1440, height: 900 }) + await page.goto('/app/manage') + const strip = page.locator('.manage-summary') + await expect(strip).toBeVisible() + const h = await strip.evaluate(el => el.getBoundingClientRect().height) + expect(h).toBeGreaterThan(40) + }) +}) + +test.describe('Headline figure contrast', () => { + test('every figure is legible against the cell it sits on', async ({ page }) => { + // A