diff --git a/.agents/building-and-testing.md b/.agents/building-and-testing.md index 3cf85c0dc283..a30aa4754a07 100644 --- a/.agents/building-and-testing.md +++ b/.agents/building-and-testing.md @@ -21,6 +21,9 @@ Let's say the user wants to build a particular backend for a given platform. For The core Go suites (`./pkg`, `./core`, plus the in-process integration suite `./tests/e2e`) are covered by a **strict, monotonic coverage ratchet**: - `make test-coverage` — runs the suites with `covermode=atomic` instrumentation and writes a merged profile to `coverage/coverage.out`. Uses the same prerequisites as `make test`. + - Prints per-root wall time and the slowest specs/hooks exceeding `COVERAGE_SLOW_SPEC_THRESHOLD` (default 3 seconds, capped by `COVERAGE_SLOW_SPEC_LIMIT`, default 25 per root); machine-readable root timings are written to `coverage/timings.tsv`. + - Verbose Ginkgo output is written to `coverage/logs/.log`, with the prior run retained as `.log.previous`. The terminal prints one status line per root and a short failure extract. If any suite fails, no merged profile is produced and the percentage ratchet is explicitly not run. A lock under `coverage/` rejects concurrent runs, which would otherwise corrupt their shared profiles and logs. + - Suites run in parallel by default and each recursive root invocation has a five-minute budget. Override auto-detected parallelism with `COVERAGE_PROCS`; tune diagnostics with `COVERAGE_SUITE_TIMEOUT` and `COVERAGE_PROGRESS_AFTER`. A timeout is a performance failure to investigate, not a reason to raise the committed default. - **`--coverpkg` (`COVERAGE_COVERPKG = core/...,pkg/...`):** coverage is attributed to the core+pkg packages, not just the package under test. This is what lets the in-process `tests/e2e` suite (which drives the real HTTP server over loopback via `application.New`) credit the `core/http/endpoints/...` handlers it exercises — folding it in roughly doubled endpoint coverage (e.g. `endpoints/openai` 13.6% → 52%). The denominator is therefore *all* of `core`+`pkg` (minus generated proto, dropped via `COVERAGE_EXCLUDE_RE`), so the number isn't comparable to a plain per-package figure. - **Integration suites (`COVERAGE_E2E_ROOTS = ./tests/e2e`)** run non-recursively (excludes `tests/e2e/distributed`, which needs containers) with `--label-filter=!real-models` (those need a downloaded model) against the mock backend built by `prepare-test`. `tests/integration` is deliberately excluded — it needs `make backends/local-store`, which the coverage CI job doesn't build. - **Flake note:** folding integration tests into a *strict* gate means a hard e2e failure (or a spec that silently stops running) can fail the coverage gate, not just the test. `--flake-attempts` absorbs transient retryable failures; covermode=atomic keeps line coverage deterministic otherwise. diff --git a/.docker/llama-cpp-compile.sh b/.docker/llama-cpp-compile.sh index 647a1c44828e..112d43c160ed 100755 --- a/.docker/llama-cpp-compile.sh +++ b/.docker/llama-cpp-compile.sh @@ -28,6 +28,10 @@ if [ -z "${BUILD_TYPE:-}" ]; then # variants with it (the host never *selects* SME unless it has it, but every variant must # still compile). if [ "${TARGETARCH}" = "arm64" ]; then + # The prebuilt base inherits default ports.ubuntu.com sources; honor the + # APT_*_MIRROR build args here like the from-source path does, so this + # apt step survives a mirror outage. + sh /LocalAI/.docker/apt-mirror.sh || true apt-get update -qq && apt-get install -y -qq gcc-14 g++-14 export CC=gcc-14 CXX=g++-14 fi diff --git a/.githooks/pre-commit b/.githooks/pre-commit index 025ecfafd078..6cb614d633ac 100755 --- a/.githooks/pre-commit +++ b/.githooks/pre-commit @@ -34,6 +34,9 @@ if [ "$go_changed" -eq 0 ] && [ "$ui_changed" -eq 0 ] && [ "$rt_changed" -eq 0 ] fi if [ "$go_changed" -eq 1 ]; then + echo "pre-commit ▶ gRPC generation (make protogen-go)" + make protogen-go + # Resolve the ref golangci-lint's new-from-merge-base should compare # against. .golangci.yml pins origin/master, which is correct in CI # (origin == the canonical repo) but wrong from a fork clone, where diff --git a/.github/workflows/external-probes.yml b/.github/workflows/external-probes.yml new file mode 100644 index 000000000000..681f417e4ae5 --- /dev/null +++ b/.github/workflows/external-probes.yml @@ -0,0 +1,38 @@ +--- +name: external compatibility probes + +on: + workflow_dispatch: + schedule: + - cron: '23 4 * * 1' + +permissions: + contents: read + +jobs: + external-probe-huggingface-xet: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + - uses: actions/setup-go@v5 + with: + go-version: '1.26.x' + cache: false + - name: Probe Hugging Face Xet compatibility + run: LOCALAI_HF_XET_SMOKE=1 go test ./pkg/huggingface-api -ginkgo.focus='pinned public Xet fixture' -count=1 + + external-probe-sigstore: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + - uses: actions/setup-go@v5 + with: + go-version: '1.26.x' + cache: false + - name: Probe public Sigstore compatibility + env: + LOCALAI_COSIGN_LIVE: '1' + LOCALAI_COSIGN_LIVE_IMAGE: ${{ vars.LOCALAI_COSIGN_LIVE_IMAGE }} + LOCALAI_COSIGN_LIVE_ISSUER: ${{ vars.LOCALAI_COSIGN_LIVE_ISSUER }} + LOCALAI_COSIGN_LIVE_IDENTITY_REGEX: ${{ vars.LOCALAI_COSIGN_LIVE_IDENTITY_REGEX }} + run: go test ./pkg/oci/cosignverify -ginkgo.focus='VerifyImage' -count=1 diff --git a/.github/workflows/test-resource-refresh.yml b/.github/workflows/test-resource-refresh.yml new file mode 100644 index 000000000000..24e19f341326 --- /dev/null +++ b/.github/workflows/test-resource-refresh.yml @@ -0,0 +1,76 @@ +--- +name: refresh offline test resources + +on: + workflow_dispatch: + schedule: + - cron: '17 3 * * 1' + +permissions: + contents: read + issues: write + packages: write + +jobs: + refresh: + strategy: + fail-fast: false + matrix: + resource-set: [default, distributed-e2e, aio] + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + - uses: actions/setup-go@v5 + with: + go-version: '1.26.x' + cache: true + - uses: oras-project/setup-oras@v1 + - name: Verify upstream resources and build compressed cache + id: refresh + continue-on-error: true + env: + LOCALAI_TEST_RESOURCES_ONLINE: '1' + run: | + set -o pipefail + make update-offline-test-cache TEST_RESOURCE_SET=${{ matrix.resource-set }} 2>&1 | tee resource-refresh.log + - name: Upload investigation evidence + if: steps.refresh.outcome == 'failure' + uses: actions/upload-artifact@v4 + with: + name: test-resource-investigation-${{ matrix.resource-set }}-${{ github.run_id }} + path: resource-refresh.log + - name: Open or update investigation issue + if: steps.refresh.outcome == 'failure' + env: + GH_TOKEN: ${{ secrets.LOCALAI_BOT_TOKEN || github.token }} + RESOURCE_SET: ${{ matrix.resource-set }} + RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }} + run: | + title="test resource integrity investigation: ${RESOURCE_SET}" + body=$(printf '%s\n\n%s\n' \ + "The scheduled offline-resource refresh failed for \`${RESOURCE_SET}\`." \ + "Do not update the manifest digest blindly. Download the evidence artifact from ${RUN_URL}, compare upstream checksums/signatures and release notes, inspect redirects, and search the [GitHub Advisory Database](https://github.com/advisories) and [OSV](https://osv.dev). Retry from a declared mirror to distinguish source drift from corruption.") + existing=$(gh issue list --state open --search "${title} in:title" --json number --jq '.[0].number // empty') + if [ -n "$existing" ]; then + gh issue comment "$existing" --body "$body" + else + gh issue create --title "$title" --body "$body" + fi + - name: Log in to GHCR + if: steps.refresh.outcome == 'success' + run: echo "${{ github.token }}" | oras login ghcr.io -u "${{ github.actor }}" --password-stdin + - name: Publish compressed cache as an OCI artifact + if: steps.refresh.outcome == 'success' + env: + RESOURCE_SET: ${{ matrix.resource-set }} + run: | + repository=$(printf '%s' "ghcr.io/${GITHUB_REPOSITORY}/localai-test-resources" | tr '[:upper:]' '[:lower:]') + digest=$(jq -r --arg set "$RESOURCE_SET" '.bundles[$set] | sub("sha256:"; "sha256-")' test-resources/manifests/lock.json) + oras push \ + --artifact-type application/vnd.localai.test-resources.v1 \ + "${repository}:${RESOURCE_SET},${RESOURCE_SET}-${digest}" \ + ".cache/test-resources/bundles/${RESOURCE_SET}.tar.zst:application/vnd.localai.test-resources.bundle.v1+zstd" \ + "test-resources/manifests/${RESOURCE_SET}.json:application/vnd.localai.test-resources.manifest.v1+json" + - name: Fail after preserving evidence + if: steps.refresh.outcome == 'failure' + run: exit 1 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 7e702a7dd2a3..191d64dcd88c 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -53,13 +53,29 @@ jobs: node-version: '22' - name: Build React UI run: make react-ui + - name: Record and pack declared test resources + run: LOCALAI_TEST_RESOURCES_ONLINE=1 make update-offline-test-cache TEST_RESOURCE_SET=default + - name: Transfer local test-resource bundle + uses: actions/upload-artifact@v4 + with: + name: test-resources-default-${{ github.run_id }} + path: | + .cache/test-resources/bundles/default.tar.zst + test-resources/manifests/lock.json + - name: Clear recorded resource cache + run: rm -rf .cache/test-resources + - name: Restore local test-resource bundle + uses: actions/download-artifact@v4 + with: + name: test-resources-default-${{ github.run_id }} + path: . # Runs the core suite with coverage and fails if total coverage dropped # below the committed baseline (coverage-baseline.txt). The gate is # strict — any decrease fails. Raise the baseline with # `make test-coverage-baseline` and commit it when coverage rises. - name: Test (with coverage gate) run: | - PATH="$PATH:/root/go/bin" make --jobs 5 --output-sync=target test-coverage-check + LOCALAI_TEST_KERNEL_ENFORCE=1 PATH="$PATH:/root/go/bin" make --jobs 5 --output-sync=target test-coverage-check - name: Upload coverage report if: ${{ always() }} uses: actions/upload-artifact@v4 @@ -113,7 +129,7 @@ jobs: # Used to run the newer GNUMake version from brew that supports --output-sync export PATH="/opt/homebrew/opt/make/libexec/gnubin:$PATH" PATH="$PATH:$HOME/go/bin" make protogen-go - PATH="$PATH:$HOME/go/bin" BUILD_TYPE="GITHUB_CI_HAS_BROKEN_METAL" CMAKE_ARGS="-DGGML_F16C=OFF -DGGML_AVX512=OFF -DGGML_AVX2=OFF -DGGML_FMA=OFF" make --jobs 4 --output-sync=target test + PATH="$PATH:$HOME/go/bin" BUILD_TYPE="GITHUB_CI_HAS_BROKEN_METAL" CMAKE_ARGS="-DGGML_F16C=OFF -DGGML_AVX512=OFF -DGGML_AVX2=OFF -DGGML_FMA=OFF" make --jobs 4 --output-sync=target TEST_RESOURCE_SET=default-darwin test - name: Setup tmate session if tests fail if: ${{ failure() }} uses: mxschmitt/action-tmate@v3.23 diff --git a/.github/workflows/tests-aio.yml b/.github/workflows/tests-aio.yml index f8d3d34f077c..4ca29730cbdc 100644 --- a/.github/workflows/tests-aio.yml +++ b/.github/workflows/tests-aio.yml @@ -76,7 +76,9 @@ jobs: PATH="$PATH:$HOME/go/bin" make protogen-go - name: Test run: | - PATH="$PATH:$HOME/go/bin" make backends/local-store backends/silero-vad backends/llama-cpp backends/whisper backends/piper backends/stablediffusion-ggml docker-build-e2e e2e-aio + PATH="$PATH:$HOME/go/bin" make backends/local-store backends/silero-vad backends/llama-cpp backends/whisper backends/piper backends/stablediffusion-ggml docker-build-e2e + LOCALAI_TEST_RESOURCES_ONLINE=1 PATH="$PATH:$HOME/go/bin" make update-offline-test-cache TEST_RESOURCE_SET=aio + LOCALAI_BACKEND_DIR="$GITHUB_WORKSPACE/backends" LOCALAI_MODELS_DIR="$GITHUB_WORKSPACE/tests/e2e-aio/models" LOCALAI_IMAGE_TAG=tests LOCALAI_IMAGE=local-ai PATH="$PATH:$HOME/go/bin" make run-e2e-aio - name: Setup tmate session if tests fail if: ${{ failure() }} uses: mxschmitt/action-tmate@v3.23 diff --git a/Makefile b/Makefile index 16881ec22cae..53b311ad5d12 100644 --- a/Makefile +++ b/Makefile @@ -99,11 +99,20 @@ COVERAGE_COVERPKG?=github.com/mudler/LocalAI/core/...,github.com/mudler/LocalAI/ ## the coverage CI job doesn't do. COVERAGE_E2E_ROOTS?=./tests/e2e COVERAGE_E2E_LABELS?=!real-models +COVERAGE_PROCS?=0 +COVERAGE_SUITE_TIMEOUT?=5m +COVERAGE_PROGRESS_AFTER?=30s +COVERAGE_SLOW_SPEC_THRESHOLD?=3 +COVERAGE_SLOW_SPEC_LIMIT?=25 ## Drop generated protobuf from the denominator (it has no tests by design). COVERAGE_EXCLUDE_RE?=grpc/proto/.*[.]pb[.]go +TEST_RESOURCE_SET?=default +TEST_RESOURCE_CACHE?=$(abspath ./.cache/test-resources) +TEST_RESOURCE_MANIFESTS?=$(abspath ./test-resources/manifests) +OFFLINE_RUN=$(abspath ./scripts/run-test-offline.sh) -.PHONY: all test test-coverage test-coverage-baseline test-coverage-check test-backend-cpp test-build-scripts test-ui test-ui-coverage-baseline test-ui-coverage-check install-hooks build vendor lint lint-all +.PHONY: all test prepare-offline-test-cache update-offline-test-cache test-network-lint test-coverage test-coverage-baseline test-coverage-check test-backend-cpp test-build-scripts test-ui test-ui-coverage-baseline test-ui-coverage-check install-hooks build vendor lint lint-all all: help @@ -195,11 +204,24 @@ prepare-test: protogen-go build-mock-backend ## now drives the mock-backend binary built by build-mock-backend; real-backend ## inference moved into tests/e2e-backends/ (per-backend, path-filtered) and ## tests/e2e-aio/ (nightly). -test: prepare-test +prepare-offline-test-cache: + @test -n "$(TEST_RESOURCE_SET)" || { echo 'TEST_RESOURCE_SET is required, for example: make prepare-offline-test-cache TEST_RESOURCE_SET=default'; exit 2; } + $(GOCMD) run ./cmd/test-resources prepare "$(TEST_RESOURCE_SET)" "$(TEST_RESOURCE_MANIFESTS)" "$(TEST_RESOURCE_CACHE)" + +test-network-lint: + scripts/test-network-lint.sh + +update-offline-test-cache: + @test -n "$(TEST_RESOURCE_SET)" || { echo 'TEST_RESOURCE_SET is required, for example: make update-offline-test-cache TEST_RESOURCE_SET=default'; exit 2; } + @test "$$LOCALAI_TEST_RESOURCES_ONLINE" = 1 || { echo 'Set LOCALAI_TEST_RESOURCES_ONLINE=1 to enter explicit online record mode'; exit 2; } + $(GOCMD) run ./cmd/test-resources update "$(TEST_RESOURCE_SET)" "$(TEST_RESOURCE_MANIFESTS)" "$(TEST_RESOURCE_CACHE)" + +test: TEST_RESOURCE_SET=default +test: test-network-lint prepare-test @echo 'Running tests' export GO_TAGS="debug" OPUS_SHIM_LIBRARY=$(abspath ./pkg/opus/shim/libopusshim.so) \ - $(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --flake-attempts $(TEST_FLAKES) --fail-fast -v -r $(TEST_PATHS) + $(OFFLINE_RUN) $(TEST_RESOURCE_SET) $(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --flake-attempts $(TEST_FLAKES) --fail-fast -v -r $(TEST_PATHS) ## Compiles and runs the standalone C++ unit tests for the backends (pure ## helpers that depend only on the stdlib + nlohmann/json, no full backend @@ -226,15 +248,21 @@ test-ci-scripts: ## and writes a merged profile to $(COVERAGE_PROFILE). Deliberately omits ## --fail-fast so a single failure doesn't truncate the coverage number, and ## uses covermode=atomic so the result is deterministic. Prints the total. -test-coverage: prepare-test - @echo 'Running tests with coverage' +test-coverage: TEST_RESOURCE_SET=default +test-coverage: test-network-lint prepare-test + @echo 'Running tests with coverage (test failures stop before the percentage ratchet)' GINKGO_TAGS="$(COVERAGE_TAGS)" \ COVERAGE_COVERPKG="$(COVERAGE_COVERPKG)" \ COVERAGE_E2E_ROOTS="$(COVERAGE_E2E_ROOTS)" \ COVERAGE_E2E_LABELS="$(COVERAGE_E2E_LABELS)" \ + COVERAGE_PROCS="$(COVERAGE_PROCS)" \ + COVERAGE_SUITE_TIMEOUT="$(COVERAGE_SUITE_TIMEOUT)" \ + COVERAGE_PROGRESS_AFTER="$(COVERAGE_PROGRESS_AFTER)" \ + COVERAGE_SLOW_SPEC_THRESHOLD="$(COVERAGE_SLOW_SPEC_THRESHOLD)" \ + COVERAGE_SLOW_SPEC_LIMIT="$(COVERAGE_SLOW_SPEC_LIMIT)" \ COVERAGE_EXCLUDE_RE='$(COVERAGE_EXCLUDE_RE)' \ OPUS_SHIM_LIBRARY=$(abspath ./pkg/opus/shim/libopusshim.so) \ - scripts/run-coverage.sh $(COVERAGE_DIR) $(COVERAGE_PROFILE) $(TEST_FLAKES) $(COVERAGE_ROOTS) + $(OFFLINE_RUN) $(TEST_RESOURCE_SET) scripts/run-coverage.sh $(COVERAGE_DIR) $(COVERAGE_PROFILE) $(TEST_FLAKES) $(COVERAGE_ROOTS) @$(GOCMD) tool cover -html=$(COVERAGE_PROFILE) -o $(COVERAGE_DIR)/coverage.html @$(GOCMD) tool cover -func=$(COVERAGE_PROFILE) | tail -n1 @@ -250,6 +278,7 @@ test-coverage-baseline: test-coverage ## run-to-run jitter from the in-process tests/e2e suite folded in via ## --coverpkg (timing-dependent which handler lines execute). test-coverage-check: test-coverage + @echo 'Running coverage percentage ratchet' @scripts/coverage-check.sh $(COVERAGE_PROFILE) $(COVERAGE_BASELINE) ######################################################## @@ -326,16 +355,18 @@ e2e-aio: LOCALAI_IMAGE=local-ai \ $(MAKE) run-e2e-aio +run-e2e-aio: TEST_RESOURCE_SET=aio run-e2e-aio: protogen-go @echo 'Running e2e AIO tests' - $(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --flake-attempts $(TEST_FLAKES) -v -r ./tests/e2e-aio + $(OFFLINE_RUN) $(TEST_RESOURCE_SET) $(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --flake-attempts $(TEST_FLAKES) -v -r ./tests/e2e-aio # Distributed architecture e2e (PostgreSQL + NATS via testcontainers). # Includes NatsJWT specs (JWT-enabled NATS). Requires Docker. # VLLMMultinode is excluded here; use test-e2e-vllm-multinode for that. +test-e2e-distributed: TEST_RESOURCE_SET=distributed-e2e test-e2e-distributed: protogen-go @echo 'Running distributed e2e tests (label Distributed, incl. NatsJWT)' - $(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --label-filter='Distributed && !VLLMMultinode' --flake-attempts $(TEST_FLAKES) -v -r ./tests/e2e/distributed + $(OFFLINE_RUN) $(TEST_RESOURCE_SET) $(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --label-filter='Distributed && !VLLMMultinode' --flake-attempts $(TEST_FLAKES) -v -r ./tests/e2e/distributed # vLLM multi-node DP smoke (CPU). Builds local-ai:tests and the # cpu-vllm backend from the current working tree, then drives a diff --git a/backend/Dockerfile.llama-cpp b/backend/Dockerfile.llama-cpp index 8e725ef623a1..2f21aaa8ed72 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/backend.proto b/backend/backend.proto index 25f45eac89cf..821838cea626 100644 --- a/backend/backend.proto +++ b/backend/backend.proto @@ -181,6 +181,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 +500,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 diff --git a/backend/cpp/bonsai/Makefile b/backend/cpp/bonsai/Makefile index 82ba9660f242..9707372a9a11 100644 --- a/backend/cpp/bonsai/Makefile +++ b/backend/cpp/bonsai/Makefile @@ -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/llama-cpp/disable-score-task.sh b/backend/cpp/llama-cpp/disable-score-task.sh new file mode 100644 index 000000000000..164d4a57f7cf --- /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 88b47a654cbf..44e3b71da7ad 100644 --- a/backend/cpp/llama-cpp/grpc-server.cpp +++ b/backend/cpp/llama-cpp/grpc-server.cpp @@ -151,40 +151,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; @@ -724,6 +690,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 { @@ -1364,6 +1346,17 @@ static void params_parse(server_context& /*ctx_server*/, const backend::ModelOpt } } +#ifndef LOCALAI_LLAMA_CPP_NO_SCORE_TASK + // Score-task suffix forking: reserve seq ids (and recurrent-state cells) + // beyond the slots so one scoring call decodes all candidate tails in a + // single batch (SERVER_TASK_TYPE_SCORE, patches/). Requires the unified + // KV cache — with per-sequence streams the extra ids would shrink every + // sequence's context to n_ctx / n_seq_max. Decided after both option + // passes so an explicit kv_unified:false wins and disables forking. + params.score_enabled = request->enablescore(); + params.n_seq_score_forks = params.score_enabled && params.kv_unified ? SERVER_SCORE_FORK_SEQS : 0; +#endif + // Terminate/pad the override vectors only after BOTH the named-option loop // and the generic passthrough (common_params_parse above) have pushed their // real entries, so back() is the null sentinel the model loader asserts on. @@ -1450,6 +1443,16 @@ class BackendServiceImpl final : public backend::Backend::Service { 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); @@ -1652,7 +1655,6 @@ class BackendServiceImpl final : public backend::Backend::Service { 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()); @@ -2221,7 +2223,6 @@ class BackendServiceImpl final : public backend::Backend::Service { 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; @@ -2755,7 +2756,6 @@ class BackendServiceImpl final : public backend::Backend::Service { 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; @@ -2865,7 +2865,6 @@ class BackendServiceImpl final : public backend::Backend::Service { 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(); @@ -2942,37 +2941,16 @@ class BackendServiceImpl final : public backend::Backend::Service { // 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; @@ -2981,40 +2959,21 @@ class BackendServiceImpl final : public backend::Backend::Service { 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 @@ -3023,6 +2982,15 @@ class BackendServiceImpl final : public backend::Backend::Service { 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); @@ -3039,9 +3007,135 @@ class BackendServiceImpl final : public backend::Backend::Service { 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()) { @@ -3049,101 +3143,57 @@ class BackendServiceImpl final : public backend::Backend::Service { } 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) + ")"); - } - - // 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); + const auto & sfx_lp = cand_logprobs[inc++]; + for (int32_t j = std::max(0, divergence - n_shared); j < (int32_t) sfx_lp.size(); j++) { + lp.push_back(sfx_lp[j]); + } + + if ((int32_t) 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 { @@ -3154,7 +3204,6 @@ class BackendServiceImpl final : public backend::Backend::Service { 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; @@ -3176,7 +3225,6 @@ class BackendServiceImpl final : public backend::Backend::Service { 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/patches/0001-add-server-task-type-score.patch b/backend/cpp/llama-cpp/patches/0001-add-server-task-type-score.patch new file mode 100644 index 000000000000..b893888235f2 --- /dev/null +++ b/backend/cpp/llama-cpp/patches/0001-add-server-task-type-score.patch @@ -0,0 +1,599 @@ +diff --git a/common/common.cpp b/common/common.cpp +index 8f13217..fc584e1 100644 +--- a/common/common.cpp ++++ b/common/common.cpp +@@ -1591,8 +1591,10 @@ struct llama_context_params common_context_params_to_llama(const common_params & + auto cparams = llama_context_default_params(); + + cparams.n_ctx = params.n_ctx; +- cparams.n_seq_max = params.n_parallel; +- cparams.n_rs_seq = params.speculative.need_n_rs_seq(); ++ // score-task forks need seq ids (and recurrent-state cells) of their ++ // own beyond the parallel slots ++ cparams.n_seq_max = params.n_parallel + params.n_seq_score_forks; ++ cparams.n_rs_seq = std::max(params.speculative.need_n_rs_seq(), (uint32_t) std::max(0, params.n_rs_seq)); + cparams.n_outputs_max = std::max(params.n_outputs_max, 0); + cparams.n_batch = params.n_batch; + cparams.n_ubatch = params.n_ubatch; +diff --git a/common/common.h b/common/common.h +index bffc176..e313bd6 100644 +--- a/common/common.h ++++ b/common/common.h +@@ -455,6 +455,9 @@ struct common_params { + int32_t n_keep = 0; // number of tokens to keep from initial prompt + int32_t n_chunks = -1; // max number of chunks to process (-1 = unlimited) + int32_t n_parallel = 1; // number of parallel sequences to decode ++ int32_t n_seq_score_forks = 0; // extra seq ids beyond n_parallel, reserved for server score-task forks ++ int32_t n_rs_seq = 0; // recurrent-state rollback snapshots per seq (hybrid models cannot rewind without them; lets score tasks reuse a cached prompt across probe changes) ++ bool score_enabled = false; // reserve server resources for the Score task type + int32_t n_sequences = 1; // number of sequences to decode + int32_t n_outputs_max = 0; // max outputs in a batch (0 = n_batch) + int32_t grp_attn_n = 1; // group-attention factor +diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt +index 780df32..1d2fe8f 100644 +--- a/tools/CMakeLists.txt ++++ b/tools/CMakeLists.txt +@@ -41,3 +41,4 @@ else() + add_subdirectory(fit-params) + add_subdirectory(results) + endif() ++add_subdirectory(grpc-server) +diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp +index 715477e..de5bed8 100644 +--- a/tools/server/server-context.cpp ++++ b/tools/server/server-context.cpp +@@ -49,7 +49,16 @@ static uint32_t server_n_outputs_max(const common_params & params) { + + const uint32_t n_outputs_per_seq = 1 + common_speculative_n_max(¶ms.speculative); + +- const uint64_t n_outputs = (uint64_t) params.n_parallel * n_outputs_per_seq; ++ // score tasks (SERVER_TASK_TYPE_SCORE) output logits for every candidate ++ // token, so reserve room for a bounded candidate tail per parallel slot ++ if (!params.score_enabled) { ++ return std::max(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/turboquant/Makefile b/backend/cpp/turboquant/Makefile index cfc2593c0670..f376e6e303bf 100644 --- a/backend/cpp/turboquant/Makefile +++ b/backend/cpp/turboquant/Makefile @@ -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/cmd/test-resources/main.go b/cmd/test-resources/main.go new file mode 100644 index 000000000000..133d3cc27189 --- /dev/null +++ b/cmd/test-resources/main.go @@ -0,0 +1,393 @@ +// SPDX-License-Identifier: MIT + +package main + +import ( + "crypto/sha256" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "time" + + "github.com/mudler/LocalAI/core/services/cloudproxy/mitm" + "github.com/mudler/LocalAI/internal/testresources" + "github.com/mudler/LocalAI/pkg/httpclient" +) + +func main() { + if err := run(os.Args[1:]); err != nil { + fmt.Fprintln(os.Stderr, "test-resources:", err) + os.Exit(1) + } +} + +func run(args []string) error { + if len(args) >= 6 && args[0] == "run" && args[4] == "--" { + return runOffline(args[1], args[2], args[3], args[5:]) + } + if len(args) != 4 { + return errors.New("usage: test-resources RESOURCE_SET MANIFEST_DIR CACHE_DIR | test-resources run RESOURCE_SET MANIFEST_DIR CACHE_DIR -- COMMAND") + } + target, manifestDir, cacheDir := args[1], args[2], args[3] + if args[0] == "update" { + return update(target, manifestDir, cacheDir) + } + if args[0] != "prepare" { + return errors.New("usage: test-resources RESOURCE_SET MANIFEST_DIR CACHE_DIR") + } + return prepare(target, manifestDir, cacheDir) +} + +func prepare(target, manifestDir, cacheDir string) error { + manifest, err := testresources.LoadManifest(filepath.Join(manifestDir, target+".json")) + if err != nil { + return fmt.Errorf("%w; run `make update-offline-test-cache TEST_RESOURCE_SET=%s`", err, target) + } + if manifest.Target != target { + return fmt.Errorf("manifest target %q does not match %q", manifest.Target, target) + } + lock, err := testresources.LoadLock(filepath.Join(manifestDir, "lock.json")) + if err != nil { + return err + } + locked, ok := lock.Bundles[target] + if !ok { + return fmt.Errorf("cache bundle is not locked for resource set %q; run `make update-offline-test-cache TEST_RESOURCE_SET=%s`", target, target) + } + if digest, ok := strings.CutPrefix(locked, "sha256:"); ok { + bundlePath := filepath.Join(cacheDir, "bundles", target+".tar.zst") + if _, err := os.Stat(bundlePath); errors.Is(err, os.ErrNotExist) { + bundlePath = filepath.Join(cacheDir, "bundles", target+".tar") + } + if err := testresources.RestoreBundle(cacheDir, bundlePath, digest); err != nil { + return preparationError(target, err) + } + } + materialized := filepath.Join(cacheDir, "materialized", target) + if err := os.MkdirAll(materialized, 0o755); err != nil { + return err + } + index, err := testresources.LoadHTTPIndex(cacheDir) + if err != nil { + return preparationError(target, err) + } + for _, resource := range manifest.HTTP { + _, err := testresources.VerifyBlob(cacheDir, resource.SHA256) + if err != nil { + return preparationError(target, err) + } + entry, ok := index[testresources.RequestKey(resource.Method, resource.URL, resource.Headers())] + if !ok || entry.Digest != resource.SHA256 { + return preparationError(target, fmt.Errorf("HTTP cache entry missing or mismatched: %s %s", resource.Method, resource.URL)) + } + } + for _, resource := range manifest.Files { + path, err := testresources.VerifyBlob(cacheDir, resource.SHA256) + if err != nil { + return preparationError(target, err) + } + environmentPath := path + if resource.Destination != "" { + destination := filepath.Join(materialized, resource.Destination) + if err := copyFile(path, destination); err != nil { + return err + } + environmentPath = destination + } + if resource.Environment != "" { + if err := os.Setenv(resource.Environment, environmentPath); err != nil { + return err + } + } + } + for _, resource := range manifest.Images { + path, err := testresources.VerifyBlob(cacheDir, resource.SHA256) + if err != nil { + return preparationError(target, err) + } + cmd := exec.Command("docker", "load", "--input", path) + cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr + if err := cmd.Run(); err != nil { + return fmt.Errorf("load declared image %s: %w (verify Docker is running and this user can access its socket)", resource.Reference, err) + } + } + return nil +} + +func update(target, manifestDir, cacheDir string) error { + if os.Getenv("LOCALAI_TEST_RESOURCES_ONLINE") != "1" { + return errors.New("update requires explicit online record mode: LOCALAI_TEST_RESOURCES_ONLINE=1") + } + manifest, err := testresources.LoadManifest(filepath.Join(manifestDir, target+".json")) + if err != nil { + return err + } + client := httpclient.New() + client.CheckRedirect = func(_ *http.Request, _ []*http.Request) error { return http.ErrUseLastResponse } + index, err := testresources.LoadHTTPIndex(cacheDir) + if err != nil { + return err + } + for _, resource := range manifest.HTTP { + entry, err := fetchHTTPWithMirrors(client, resource, cacheDir) + if err != nil { + return err + } + index[testresources.RequestKey(resource.Method, resource.URL, resource.Headers())] = entry + } + if err := testresources.WriteHTTPIndex(cacheDir, index); err != nil { + return err + } + for _, resource := range manifest.Files { + if err := fetchWithMirrors(client, resource.URL, resource.Mirrors, resource.SHA256, cacheDir); err != nil { + return err + } + } + for _, resource := range manifest.Images { + if err := pullAndPack(resource.Reference, resource.SHA256, cacheDir); err != nil { + return err + } + } + bundlePath := filepath.Join(cacheDir, "bundles", target+".tar.zst") + digest, err := testresources.PackBundle(cacheDir, bundlePath, manifest) + if err != nil { + return err + } + lockPath := filepath.Join(manifestDir, "lock.json") + lock, err := testresources.LoadLock(lockPath) + if err != nil { + return err + } + lock.Bundles[target] = "sha256:" + digest + return testresources.WriteLock(lockPath, lock) +} + +func fetchHTTPWithMirrors(client *http.Client, resource testresources.HTTP, cacheDir string) (testresources.HTTPEntry, error) { + urls := append([]string{resource.URL}, resource.Mirrors...) + var failures []error + for _, candidate := range urls { + for attempt := 1; attempt <= 2; attempt++ { + started := time.Now() + entry, err := fetchHTTP(client, resource, candidate, cacheDir) + fmt.Fprintf(os.Stderr, "test-resources: download %s attempt %d took %s\n", candidate, attempt, time.Since(started).Round(time.Millisecond)) + if err == nil { + return entry, nil + } + failures = append(failures, fmt.Errorf("%s attempt %d: %w", candidate, attempt, err)) + } + } + return testresources.HTTPEntry{}, resourceChangeError(resource.URL, resource.SHA256, failures) +} + +func fetchHTTP(client *http.Client, resource testresources.HTTP, sourceURL, cacheDir string) (testresources.HTTPEntry, error) { + request, err := http.NewRequest(resource.Method, sourceURL, nil) + if err != nil { + return testresources.HTTPEntry{}, err + } + request.Header = resource.Headers() + response, err := client.Do(request) + if err != nil { + return testresources.HTTPEntry{}, fmt.Errorf("fetch %s: %w", resource.URL, err) + } + defer func() { _ = response.Body.Close() }() + size, err := storeVerified(response.Body, resource.SHA256, cacheDir) + if err != nil { + return testresources.HTTPEntry{}, err + } + return testresources.HTTPEntry{Digest: resource.SHA256, Size: size, Status: response.StatusCode, Header: testresources.SanitizeHeaders(response.Header)}, nil +} + +func fetchWithMirrors(client *http.Client, primary string, mirrors []string, expected, cacheDir string) error { + urls := append([]string{primary}, mirrors...) + var failures []error + for _, candidate := range urls { + for attempt := 1; attempt <= 2; attempt++ { + started := time.Now() + err := fetch(client, candidate, expected, cacheDir) + fmt.Fprintf(os.Stderr, "test-resources: download %s attempt %d took %s\n", candidate, attempt, time.Since(started).Round(time.Millisecond)) + if err == nil { + return nil + } + failures = append(failures, fmt.Errorf("%s attempt %d: %w", candidate, attempt, err)) + } + } + return resourceChangeError(primary, expected, failures) +} + +func resourceChangeError(resourceURL, expected string, failures []error) error { + return fmt.Errorf("resource verification failed for %s (expected sha256:%s) after retrying every declared mirror: %w\nsecurity review required before changing the manifest: compare the upstream release checksum/signature and changelog, inspect redirects, and search https://github.com/advisories and https://osv.dev; a mismatch may be an upstream release, mirror corruption, or a supply-chain incident", resourceURL, expected, errors.Join(failures...)) +} + +func fetch(client *http.Client, rawURL, expected, cacheDir string) error { + request, err := http.NewRequest(http.MethodGet, rawURL, nil) + if err != nil { + return err + } + response, err := client.Do(request) + if err != nil { + return fmt.Errorf("fetch %s: %w", rawURL, err) + } + if response.StatusCode != http.StatusOK { + closeErr := response.Body.Close() + if closeErr != nil { + return errors.Join(fmt.Errorf("fetch %s: status %s", rawURL, response.Status), closeErr) + } + return fmt.Errorf("fetch %s: status %s", rawURL, response.Status) + } + _, storeErr := storeVerified(response.Body, expected, cacheDir) + return errors.Join(storeErr, response.Body.Close()) +} + +func storeVerified(reader io.Reader, expected, cacheDir string) (int64, error) { + directory := filepath.Join(cacheDir, "blobs", "sha256") + if err := os.MkdirAll(directory, 0o755); err != nil { + return 0, err + } + temporary, err := os.CreateTemp(directory, ".record-*") + if err != nil { + return 0, err + } + temporaryName := temporary.Name() + defer func() { _ = os.Remove(temporaryName) }() + hash := sha256.New() + size, copyErr := io.Copy(io.MultiWriter(temporary, hash), reader) + closeErr := temporary.Close() + if err := errors.Join(copyErr, closeErr); err != nil { + return 0, err + } + actual := fmt.Sprintf("%x", hash.Sum(nil)) + if actual != expected { + return 0, fmt.Errorf("resource digest mismatch: expected sha256:%s, got sha256:%s", expected, actual) + } + if err := os.Rename(temporaryName, testresources.BlobPath(cacheDir, expected)); err != nil { + return 0, err + } + return size, nil +} + +func pullAndPack(reference, expected, cacheDir string) error { + if !strings.Contains(reference, "@sha256:") { + return fmt.Errorf("refusing mutable image reference %s", reference) + } + if err := exec.Command("docker", "pull", reference).Run(); err != nil { + return fmt.Errorf("pull image %s: %w", reference, err) + } + cmd := exec.Command("docker", "save", reference) + stdout, err := cmd.StdoutPipe() + if err != nil { + return err + } + if err := cmd.Start(); err != nil { + return err + } + _, storeErr := storeVerified(stdout, expected, cacheDir) + waitErr := cmd.Wait() + return errors.Join(storeErr, waitErr) +} + +func preparationError(target string, err error) error { + return fmt.Errorf("%w; run `make prepare-offline-test-cache TEST_RESOURCE_SET=%s` during the network-enabled preparation phase", err, target) +} + +func copyFile(source, destination string) error { + if err := os.MkdirAll(filepath.Dir(destination), 0o755); err != nil { + return err + } + in, err := os.Open(source) + if err != nil { + return err + } + out, err := os.OpenFile(destination, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o644) + if err != nil { + _ = in.Close() + return err + } + _, copyErr := io.Copy(out, in) + return errors.Join(copyErr, in.Close(), out.Close()) +} + +func runOffline(target, manifestDir, cacheDir string, command []string) error { + manifest, err := testresources.LoadManifest(filepath.Join(manifestDir, target+".json")) + if err != nil { + return err + } + if err := prepare(target, manifestDir, cacheDir); err != nil { + return err + } + dockerNetwork := "" + if runtime.GOOS == "linux" && (len(manifest.Images) > 0 || target == "aio") { + dockerNetwork = fmt.Sprintf("localai-test-%d", os.Getpid()) + create := exec.Command("docker", "network", "create", "--internal", dockerNetwork) + create.Stdout, create.Stderr = io.Discard, os.Stderr + if err := create.Run(); err != nil { + return fmt.Errorf("create internal test Docker network: %w", err) + } + defer func() { _ = exec.Command("docker", "network", "rm", dockerNetwork).Run() }() + } + index, err := testresources.LoadHTTPIndex(cacheDir) + if err != nil { + return err + } + hosts := make([]string, 0, len(manifest.HTTP)) + seen := map[string]bool{} + for _, resource := range manifest.HTTP { + parsed, err := url.Parse(resource.URL) + if err != nil { + return err + } + if parsed.Hostname() != "" && !seen[parsed.Hostname()] { + hosts = append(hosts, parsed.Hostname()) + seen[parsed.Hostname()] = true + } + } + caDir := filepath.Join(cacheDir, "ca") + ca, err := mitm.LoadOrCreateCA(caDir) + if err != nil { + return err + } + server, err := mitm.NewServer(mitm.Config{ + Addr: "127.0.0.1:0", CA: ca, InterceptHosts: hosts, AllowPlainHTTP: true, InterceptAll: true, + Handler: func(w http.ResponseWriter, r *http.Request, _ string) { + key := testresources.RequestKey(r.Method, r.URL.String(), r.Header) + entry, ok := index[key] + if !ok { + http.Error(w, "undeclared test HTTP request: "+key, http.StatusGatewayTimeout) + return + } + if err := testresources.ReplayResponse(w, cacheDir, entry); err != nil { + http.Error(w, err.Error(), http.StatusBadGateway) + } + }, + }) + if err != nil { + return err + } + if err := server.Start(); err != nil { + return err + } + defer server.Stop() + proxyURL := "http://" + server.Addr() + caPath := filepath.Join(caDir, "ca.crt") + env := append(os.Environ(), + "LOCALAI_TEST_OFFLINE=1", "HTTP_PROXY="+proxyURL, "HTTPS_PROXY="+proxyURL, + "ALL_PROXY="+proxyURL, "http_proxy="+proxyURL, "https_proxy="+proxyURL, + "all_proxy="+proxyURL, "SSL_CERT_FILE="+caPath, "CURL_CA_BUNDLE="+caPath, + "REQUESTS_CA_BUNDLE="+caPath, "GIT_SSL_CAINFO="+caPath, "NODE_EXTRA_CA_CERTS="+caPath, + "NO_PROXY=localhost,127.0.0.0/8,::1,172.16.0.0/12,192.168.0.0/16", + "no_proxy=localhost,127.0.0.0/8,::1,172.16.0.0/12,192.168.0.0/16", + "TESTCONTAINERS_RYUK_DISABLED=true", + ) + if dockerNetwork != "" { + env = append(env, "LOCALAI_TEST_DOCKER_NETWORK="+dockerNetwork) + } + cmd := exec.Command(command[0], command[1:]...) + cmd.Env, cmd.Stdin, cmd.Stdout, cmd.Stderr = env, os.Stdin, os.Stdout, os.Stderr + return cmd.Run() +} diff --git a/core/application/router_factories.go b/core/application/router_factories.go index 879c43a835ee..2bee523ea3bf 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 diff --git a/core/application/router_factories_test.go b/core/application/router_factories_test.go index 5a6988a88fba..91b26d491fdb 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/backend/options.go b/core/backend/options.go index 72f49f1eda47..95c12c6ed27e 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), diff --git a/core/backend/options_internal_test.go b/core/backend/options_internal_test.go index bf4258bbfdc2..19e62504c58c 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 103d36efc577..3525da2991c8 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 f92d2b015e5f..5a0fc0f2c43f 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 7ce795f6098e..174e86500102 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/config/backend_capabilities.go b/core/config/backend_capabilities.go index 6605bcf39e81..a8ca571d29c9 100644 --- a/core/config/backend_capabilities.go +++ b/core/config/backend_capabilities.go @@ -30,6 +30,7 @@ const ( UsecaseFaceRecognition = "face_recognition" UsecaseSpeakerRecognition = "speaker_recognition" UsecaseTokenClassify = "token_classify" + UsecaseScore = "score" ) // GRPCMethod identifies a Backend service RPC from backend.proto. @@ -60,6 +61,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 @@ -192,6 +194,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. @@ -241,8 +248,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", diff --git a/core/config/meta/registry.go b/core/config/meta/registry.go index 2ad64f4d3f0f..be6b7516d4eb 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": { diff --git a/core/config/model_config.go b/core/config/model_config.go index 7ca24be64dc9..664354762c7f 100644 --- a/core/config/model_config.go +++ b/core/config/model_config.go @@ -669,6 +669,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 +692,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 +1051,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 +1510,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 +1640,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 @@ -1691,9 +1756,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 { @@ -1883,8 +1948,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 bd9e93865035..10b0517e2648 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 961693b015a6..55e431c6b943 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 1b7a10f455a9..19bbebd3997c 100644 --- a/core/config/model_config_test.go +++ b/core/config/model_config_test.go @@ -1,8 +1,6 @@ package config import ( - "io" - "net/http" "os" "path/filepath" @@ -127,21 +125,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{ @@ -271,17 +267,7 @@ parameters: Expect(valid).To(BeTrue()) Expect(err).NotTo(HaveOccurred()) - // download https://raw.githubusercontent.com/mudler/LocalAI/v2.25.0/embedded/models/hermes-2-pro-mistral.yaml - httpClient := http.Client{} - resp, err := httpClient.Get("https://raw.githubusercontent.com/mudler/LocalAI/v2.25.0/embedded/models/hermes-2-pro-mistral.yaml") - Expect(err).To(BeNil()) - defer resp.Body.Close() - tmp, err = os.CreateTemp("", "config.yaml") - Expect(err).To(BeNil()) - defer os.Remove(tmp.Name()) - _, err = io.Copy(tmp, resp.Body) - Expect(err).To(BeNil()) - configs, err = readModelConfigsFromFile(tmp.Name()) + configs, err = readModelConfigsFromFile(filepath.Join("testdata", "hermes-2-pro-mistral.yaml")) config = configs[0] Expect(err).To(BeNil()) Expect(config).ToNot(BeNil()) diff --git a/core/config/testdata/hermes-2-pro-mistral.yaml b/core/config/testdata/hermes-2-pro-mistral.yaml new file mode 100644 index 000000000000..058e14ef364d --- /dev/null +++ b/core/config/testdata/hermes-2-pro-mistral.yaml @@ -0,0 +1,7 @@ +name: hermes-2-pro-mistral +backend: llama-cpp +context_size: 4096 +parameters: + model: hermes-2-pro-mistral.Q4_K_M.gguf +template: + chat: chatml diff --git a/core/gallery/backends_test.go b/core/gallery/backends_test.go index a3339861c7fc..27dc35801c7c 100644 --- a/core/gallery/backends_test.go +++ b/core/gallery/backends_test.go @@ -1,12 +1,18 @@ package gallery import ( + "archive/tar" + "bytes" "context" "encoding/json" "os" "path/filepath" "runtime" + "github.com/google/go-containerregistry/pkg/name" + "github.com/google/go-containerregistry/pkg/v1/empty" + "github.com/google/go-containerregistry/pkg/v1/mutate" + "github.com/google/go-containerregistry/pkg/v1/tarball" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/system" @@ -15,9 +21,21 @@ import ( "gopkg.in/yaml.v3" ) -const ( - testImage = "quay.io/mudler/tests:localai-backend-test" -) +func writeBackendImageFixture(path string) { + var layerTar bytes.Buffer + w := tar.NewWriter(&layerTar) + contents := []byte("#!/bin/sh\necho test backend\n") + Expect(w.WriteHeader(&tar.Header{Name: "run.sh", Mode: 0o755, Size: int64(len(contents))})).To(Succeed()) + _, err := w.Write(contents) + Expect(err).NotTo(HaveOccurred()) + Expect(w.Close()).To(Succeed()) + + layer, err := tarball.LayerFromReader(bytes.NewReader(layerTar.Bytes())) + Expect(err).NotTo(HaveOccurred()) + image, err := mutate.AppendLayers(empty.Image, layer) + Expect(err).NotTo(HaveOccurred()) + Expect(tarball.WriteToFile(path, name.MustParseReference("localai/backend-test:fixture"), image)).To(Succeed()) +} var _ = Describe("Runtime capability-based backend selection", func() { var tempDir string @@ -99,6 +117,7 @@ var _ = Describe("Gallery Backends", func() { galleries []config.Gallery ml *model.ModelLoader systemState *system.SystemState + testImage string ) BeforeEach(func() { @@ -106,11 +125,21 @@ var _ = Describe("Gallery Backends", func() { tempDir, err = os.MkdirTemp("", "gallery-test-*") Expect(err).NotTo(HaveOccurred()) - // Setup test galleries + imagePath := filepath.Join(tempDir, "backend-image.tar") + writeBackendImageFixture(imagePath) + testImage = "ocifile://" + imagePath + + galleryPath := filepath.Join(tempDir, "backend-gallery.yaml") + galleryData, err := yaml.Marshal(GalleryBackends{ + &GalleryBackend{Metadata: Metadata{Name: "test-backend"}, URI: testImage}, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(os.WriteFile(galleryPath, galleryData, 0o644)).To(Succeed()) + galleries = []config.Gallery{ { Name: "test-gallery", - URL: "https://gist.githubusercontent.com/mudler/71d5376bc2aa168873fa519fa9f4bd56/raw/0557f9c640c159fa8e4eab29e8d98df6a3d6e80f/backend-gallery.yaml", + URL: "file://" + galleryPath, }, } systemState, err = system.GetSystemState(system.WithBackendPath(tempDir)) @@ -877,7 +906,7 @@ var _ = Describe("Gallery Backends", func() { Metadata: Metadata{ Name: "test-backend", }, - URI: "quay.io/mudler/tests:localai-backend-test", + URI: testImage, Alias: "test-alias", } @@ -907,7 +936,7 @@ var _ = Describe("Gallery Backends", func() { Metadata: Metadata{ Name: "test-backend", }, - URI: "quay.io/mudler/tests:localai-backend-test", + URI: testImage, Alias: "test-alias", } @@ -931,7 +960,7 @@ var _ = Describe("Gallery Backends", func() { Metadata: Metadata{ Name: "test-backend", }, - URI: "quay.io/mudler/tests:localai-backend-test", + URI: testImage, Alias: "test-alias", } diff --git a/core/gallery/importers/discovery_options_test.go b/core/gallery/importers/discovery_options_test.go new file mode 100644 index 000000000000..163caa3659c7 --- /dev/null +++ b/core/gallery/importers/discovery_options_test.go @@ -0,0 +1,49 @@ +// SPDX-License-Identifier: MIT + +package importers_test + +import ( + "context" + "encoding/json" + "errors" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/gallery/importers" + hfapi "github.com/mudler/LocalAI/pkg/huggingface-api" +) + +type fixtureMetadata struct { + details *hfapi.ModelDetails + err error + calls []string +} + +func (f *fixtureMetadata) GetModelDetails(repo string) (*hfapi.ModelDetails, error) { + f.calls = append(f.calls, repo) + return f.details, f.err +} + +var _ = Describe("DiscoverModelConfigWithOptions", func() { + It("uses fixture metadata without creating a live client", func() { + metadata := &fixtureMetadata{details: &hfapi.ModelDetails{ + ModelID: "fixture/whisper", + PipelineTag: "automatic-speech-recognition", + Files: []hfapi.ModelFile{{Path: "ggml-model.bin"}}, + }} + config, err := importers.DiscoverModelConfigWithOptions(context.Background(), "hf://fixture/whisper", json.RawMessage(`{}`), importers.DiscoverOptions{HuggingFace: metadata}) + Expect(err).NotTo(HaveOccurred()) + Expect(config.Name).NotTo(BeEmpty()) + Expect(metadata.calls).To(Equal([]string{"fixture/whisper"})) + }) + + It("does not invoke metadata after cancellation", func() { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + metadata := &fixtureMetadata{err: errors.New("must not be returned")} + _, err := importers.DiscoverModelConfigWithOptions(ctx, "hf://fixture/model", json.RawMessage(`{}`), importers.DiscoverOptions{HuggingFace: metadata}) + Expect(err).To(MatchError(context.Canceled)) + Expect(metadata.calls).To(BeEmpty()) + }) +}) diff --git a/core/gallery/importers/importers.go b/core/gallery/importers/importers.go index 8c156144d44c..b1b6c7f354d9 100644 --- a/core/gallery/importers/importers.go +++ b/core/gallery/importers/importers.go @@ -1,6 +1,7 @@ package importers import ( + "context" "encoding/json" "errors" "fmt" @@ -26,6 +27,8 @@ import ( // this sentinel so legacy callers keep working. var ErrAmbiguousImport = errors.New("importer: ambiguous — specify preferences.backend") +var newHuggingFaceMetadata = func() HuggingFaceMetadata { return hfapi.NewClient() } + // AmbiguousImportError is the concrete error DiscoverModelConfig returns when // it can't pick an importer automatically. It carries the importer-modality // key (e.g. "tts", "asr") and the list of candidate backend names so HTTP @@ -246,15 +249,48 @@ func hasYAMLExtension(uri string) bool { } func DiscoverModelConfig(uri string, preferences json.RawMessage) (gallery.ModelConfig, error) { + return DiscoverModelConfigWithOptions(context.Background(), uri, preferences, DiscoverOptions{}) +} + +// HuggingFaceMetadata provides only the repository metadata needed during +// importer discovery. Tests can supply fixtures without constructing a live +// Hugging Face client. +type HuggingFaceMetadata interface { + GetModelDetails(string) (*hfapi.ModelDetails, error) +} + +// DiscoverOptions contains optional dependencies for model discovery. +type DiscoverOptions struct { + HuggingFace HuggingFaceMetadata +} + +// SetHuggingFaceMetadataFactoryForTest replaces the production metadata +// client factory and returns a restore function. It must only be called by a +// serial test-suite setup before discovery begins. +func SetHuggingFaceMetadataFactoryForTest(factory func() HuggingFaceMetadata) func() { + previous := newHuggingFaceMetadata + newHuggingFaceMetadata = factory + return func() { newHuggingFaceMetadata = previous } +} + +// DiscoverModelConfigWithOptions discovers a model using explicitly supplied +// dependencies. A nil metadata client retains the production behavior. +func DiscoverModelConfigWithOptions(ctx context.Context, uri string, preferences json.RawMessage, opts DiscoverOptions) (gallery.ModelConfig, error) { var err error var modelConfig gallery.ModelConfig - hf := hfapi.NewClient() + hf := opts.HuggingFace + if hf == nil { + hf = newHuggingFaceMetadata() + } hfrepoID := strings.ReplaceAll(uri, "huggingface://", "") hfrepoID = strings.ReplaceAll(hfrepoID, "hf://", "") hfrepoID = strings.ReplaceAll(hfrepoID, "https://huggingface.co/", "") + if err := ctx.Err(); err != nil { + return gallery.ModelConfig{}, err + } hfDetails, err := hf.GetModelDetails(hfrepoID) if err != nil { // maybe not a HF repository diff --git a/core/gallery/importers/importers_suite_test.go b/core/gallery/importers/importers_suite_test.go index a65b8163ad56..3e148a98c46f 100644 --- a/core/gallery/importers/importers_suite_test.go +++ b/core/gallery/importers/importers_suite_test.go @@ -1,13 +1,72 @@ package importers_test import ( + "context" + "errors" "testing" + gguf "github.com/gpustack/gguf-parser-go" + "github.com/mudler/LocalAI/core/gallery/importers" + hfapi "github.com/mudler/LocalAI/pkg/huggingface-api" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) +type metadataFixtures map[string]*hfapi.ModelDetails + +func (f metadataFixtures) GetModelDetails(repo string) (*hfapi.ModelDetails, error) { + details, ok := f[repo] + if !ok { + return nil, errors.New("metadata fixture not declared: " + repo) + } + return details, nil +} + +func file(repo, path, sha string) hfapi.ModelFile { + return hfapi.ModelFile{Path: path, SHA256: sha, URL: "https://huggingface.co/" + repo + "/resolve/main/" + path} // test-network: fixture +} + +var fixtures = metadataFixtures{ + "mudler/vibevoice.cpp-models": { + ModelID: "mudler/vibevoice.cpp-models", Author: "mudler", + Files: []hfapi.ModelFile{ + file("mudler/vibevoice.cpp-models", "vibevoice-realtime-Q4_K_M.gguf", "01"), + file("mudler/vibevoice.cpp-models", "vibevoice-asr-Q4_K_M.gguf", "02"), + file("mudler/vibevoice.cpp-models", "tokenizer.gguf", "03"), + file("mudler/vibevoice.cpp-models", "voice-Alice.gguf", "04"), + }, + }, + "UsefulSensors/moonshine-tiny": {ModelID: "UsefulSensors/moonshine-tiny", Author: "UsefulSensors", PipelineTag: "automatic-speech-recognition", Files: []hfapi.ModelFile{file("UsefulSensors/moonshine-tiny", "model.onnx", "05")}}, + "nvidia/parakeet-tdt-0.6b-v3": {ModelID: "nvidia/parakeet-tdt-0.6b-v3", Author: "nvidia", PipelineTag: "automatic-speech-recognition", Files: []hfapi.ModelFile{file("nvidia/parakeet-tdt-0.6b-v3", "parakeet.nemo", "06")}}, + "LiquidAI/LFM2.5-Audio-1.5B": {ModelID: "LiquidAI/LFM2.5-Audio-1.5B", Author: "LiquidAI"}, + "LiquidAI/LFM2-Audio-1.5B": {ModelID: "LiquidAI/LFM2-Audio-1.5B", Author: "LiquidAI"}, + "LiquidAI/LFM2.5-Audio-1.5B-GGUF": {ModelID: "LiquidAI/LFM2.5-Audio-1.5B-GGUF", Author: "LiquidAI", Files: []hfapi.ModelFile{file("LiquidAI/LFM2.5-Audio-1.5B-GGUF", "LFM2.5-Audio-Q4_K_M.gguf", "07")}}, + "hexgrad/Kokoro-82M": {ModelID: "hexgrad/Kokoro-82M", Author: "hexgrad", PipelineTag: "text-to-speech", Files: []hfapi.ModelFile{file("hexgrad/Kokoro-82M", "kokoro-v1_0.pth", "08")}}, + "Qwen/Qwen3-ASR-1.7B": {ModelID: "Qwen/Qwen3-ASR-1.7B", Author: "Qwen", PipelineTag: "automatic-speech-recognition"}, + "HirCoir/piper-voice-es-mx-lucas-melor": {ModelID: "HirCoir/piper-voice-es-mx-lucas-melor", Author: "HirCoir", PipelineTag: "text-to-speech", Files: []hfapi.ModelFile{file("HirCoir/piper-voice-es-mx-lucas-melor", "es_MX-lucas-medium.onnx", "09"), file("HirCoir/piper-voice-es-mx-lucas-melor", "es_MX-lucas-medium.onnx.json", "10")}}, + "h94/IP-Adapter-FaceID": {ModelID: "h94/IP-Adapter-FaceID", Author: "h94", PipelineTag: "text-to-image"}, + "LocalAI-io/whisper-large-v3-it-yodas-only-ggml": {ModelID: "LocalAI-io/whisper-large-v3-it-yodas-only-ggml", Author: "LocalAI-io", PipelineTag: "automatic-speech-recognition", Files: []hfapi.ModelFile{file("LocalAI-io/whisper-large-v3-it-yodas-only-ggml", "ggml-model-q4_0.bin", "11"), file("LocalAI-io/whisper-large-v3-it-yodas-only-ggml", "ggml-model-q5_0.bin", "12"), file("LocalAI-io/whisper-large-v3-it-yodas-only-ggml", "ggml-model-q8_0.bin", "13")}}, + "Systran/faster-whisper-large-v3": {ModelID: "Systran/faster-whisper-large-v3", Author: "Systran", PipelineTag: "automatic-speech-recognition", Files: []hfapi.ModelFile{file("Systran/faster-whisper-large-v3", "model.bin", "14"), file("Systran/faster-whisper-large-v3", "config.json", "15")}}, + "nari-labs/Dia-1.6B": {ModelID: "nari-labs/Dia-1.6B", Author: "nari-labs", PipelineTag: "text-to-speech"}, + "mudler/rfdetr-cpp-nano": {ModelID: "mudler/rfdetr-cpp-nano", Author: "mudler", PipelineTag: "object-detection", Files: []hfapi.ModelFile{file("mudler/rfdetr-cpp-nano", "rfdetr-nano-Q4_K_M.gguf", "16")}}, + "Qdrant/bm25": {ModelID: "Qdrant/bm25", Author: "Qdrant", PipelineTag: "sentence-similarity"}, + "pyannote/voice-activity-detection": {ModelID: "pyannote/voice-activity-detection", Author: "pyannote", PipelineTag: "automatic-speech-recognition"}, + "mudler/LocalAI-functioncall-qwen2.5-7b-v0.5-Q4_K_M-GGUF": {ModelID: "mudler/LocalAI-functioncall-qwen2.5-7b-v0.5-Q4_K_M-GGUF", Author: "mudler", Files: []hfapi.ModelFile{file("mudler/LocalAI-functioncall-qwen2.5-7b-v0.5-Q4_K_M-GGUF", "localai-functioncall-qwen2.5-7b-v0.5-q4_k_m.gguf", "4e7b7fe1d54b881f1ef90799219dc6cc285d29db24f559c8998d1addb35713d4")}}, + "Qwen/Qwen3-VL-2B-Instruct-GGUF": {ModelID: "Qwen/Qwen3-VL-2B-Instruct-GGUF", Author: "Qwen", Files: []hfapi.ModelFile{file("Qwen/Qwen3-VL-2B-Instruct-GGUF", "Qwen3VL-2B-Instruct-Q4_K_M.gguf", "17"), file("Qwen/Qwen3-VL-2B-Instruct-GGUF", "Qwen3VL-2B-Instruct-Q8_0.gguf", "18"), file("Qwen/Qwen3-VL-2B-Instruct-GGUF", "mmproj-Qwen3VL-2B-Instruct-F16.gguf", "20"), file("Qwen/Qwen3-VL-2B-Instruct-GGUF", "mmproj-Qwen3VL-2B-Instruct-Q8_0.gguf", "19")}}, +} + func TestImporters(t *testing.T) { RegisterFailHandler(Fail) RunSpecs(t, "Importers test suite") } + +var _ = BeforeSuite(func() { + restoreMetadata := importers.SetHuggingFaceMetadataFactoryForTest(func() importers.HuggingFaceMetadata { return fixtures }) + restoreMTP := importers.SetMTPProbeForTest(func(context.Context, string) (*gguf.GGUFFile, error) { + return nil, errors.New("remote GGUF probing disabled in fixture-backed importer tests") + }) + DeferCleanup(func() { + restoreMTP() + restoreMetadata() + }) +}) diff --git a/core/gallery/importers/llama-cpp.go b/core/gallery/importers/llama-cpp.go index a1f3250f15a9..36b792ebdda4 100644 --- a/core/gallery/importers/llama-cpp.go +++ b/core/gallery/importers/llama-cpp.go @@ -19,10 +19,21 @@ import ( ) var ( - _ Importer = &LlamaCPPImporter{} - _ AdditionalBackendsProvider = &LlamaCPPImporter{} + _ Importer = &LlamaCPPImporter{} + _ AdditionalBackendsProvider = &LlamaCPPImporter{} + parseRemoteGGUF = func(ctx context.Context, url string) (*gguf.GGUFFile, error) { + return gguf.ParseGGUFFileRemote(ctx, url) + } ) +// SetMTPProbeForTest replaces the remote GGUF header reader and returns a +// restore function. It must only be called during serial suite setup. +func SetMTPProbeForTest(probe func(context.Context, string) (*gguf.GGUFFile, error)) func() { + previous := parseRemoteGGUF + parseRemoteGGUF = probe + return func() { parseRemoteGGUF = previous } +} + type LlamaCPPImporter struct{} func (i *LlamaCPPImporter) Name() string { return "llama-cpp" } @@ -390,7 +401,7 @@ func maybeApplyMTPDefaults(modelConfig *config.ModelConfig, details Details, cfg } }() - f, err := gguf.ParseGGUFFileRemote(ctx, probeURL) + f, err := parseRemoteGGUF(ctx, probeURL) 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/request_test.go b/core/gallery/request_test.go index 1167569653db..29192511e59e 100644 --- a/core/gallery/request_test.go +++ b/core/gallery/request_test.go @@ -6,6 +6,7 @@ import ( "net/http/httptest" . "github.com/mudler/LocalAI/core/gallery" + "github.com/mudler/LocalAI/pkg/downloader" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) @@ -18,9 +19,8 @@ var _ = Describe("Gallery API tests", func() { URL: "github:go-skynet/model-gallery/gpt4all-j.yaml@main", }, } - e, err := GetGalleryConfigFromURL[ModelConfig](req.URL, "") - Expect(err).ToNot(HaveOccurred()) - Expect(e.Name).To(Equal("gpt4all-j")) + resolved := downloader.URI(req.URL).ResolveURL() + Expect(resolved).To(Equal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml")) // test-network: fixture }) }) diff --git a/core/http/app_test.go b/core/http/app_test.go index 36072b96e118..90c60532efb2 100644 --- a/core/http/app_test.go +++ b/core/http/app_test.go @@ -297,9 +297,7 @@ func getRequest(url string, header http.Header) (error, int, []byte) { return nil, resp.StatusCode, body } -const bertEmbeddingsURL = `https://gist.githubusercontent.com/mudler/0a080b166b87640e8644b09c2aee6e3b/raw/f0e8c26bb72edc16d9fbafbfd6638072126ff225/bert-embeddings-gallery.yaml` - -var _ = Describe("API test", func() { +var _ = Describe("API test", Serial, func() { var app *echo.Echo var client *openai.Client @@ -308,6 +306,7 @@ var _ = Describe("API test", func() { var cancel context.CancelFunc var tmpdir string var modelDir string + var bertEmbeddingsURL string // localAIApp captures the Application so AfterEach can synchronously // stop the spawned gRPC backend processes. application.New cancels // them asynchronously on context cancel, which races with test-binary @@ -332,6 +331,17 @@ var _ = Describe("API test", func() { modelDir = filepath.Join(tmpdir, "models") err = os.Mkdir(modelDir, 0750) Expect(err).ToNot(HaveOccurred()) + fixtureDir := filepath.Join(modelDir, ".fixtures") + err = os.Mkdir(fixtureDir, 0750) + Expect(err).ToNot(HaveOccurred()) + galleryFixturePath := filepath.Join(fixtureDir, "bert-embeddings-gallery.yaml") + err = os.WriteFile(galleryFixturePath, []byte("name: bert\nconfig_file: |\n name: bert\n backend: embeddings\n usage: You can test this model with curl like this\n parameters:\n model: bert\n"), 0600) + Expect(err).ToNot(HaveOccurred()) + bertEmbeddingsURL = "file://" + galleryFixturePath + // Additional files are cache inputs, not behavior under test here. Seed the + // destination so model application never reaches the public network. + err = os.WriteFile(filepath.Join(modelDir, "foo.yaml"), []byte("fixture: true\n"), 0600) + Expect(err).ToNot(HaveOccurred()) c, cancel = context.WithCancel(context.Background()) @@ -511,7 +521,7 @@ var _ = Describe("API test", func() { fmt.Println(response) resp = response return response["processed"].(bool) - }, "360s", "10s").Should(Equal(true)) + }, "30s", "50ms").Should(Equal(true)) Expect(resp["message"]).ToNot(ContainSubstring("error")) dat, err := os.ReadFile(filepath.Join(modelDir, "bert2.yaml")) @@ -556,7 +566,7 @@ var _ = Describe("API test", func() { Eventually(func() bool { response := getModelStatus("http://" + testHTTPAddr + "/models/jobs/" + uuid) return response["processed"].(bool) - }, "360s", "10s").Should(Equal(true)) + }, "30s", "50ms").Should(Equal(true)) dat, err := os.ReadFile(filepath.Join(modelDir, "bert.yaml")) Expect(err).ToNot(HaveOccurred()) @@ -580,7 +590,7 @@ var _ = Describe("API test", func() { Eventually(func() bool { response := getModelStatus("http://" + testHTTPAddr + "/models/jobs/" + uuid) return response["processed"].(bool) - }, "360s", "10s").Should(Equal(true)) + }, "30s", "50ms").Should(Equal(true)) dat, err := os.ReadFile(filepath.Join(modelDir, "bert.yaml")) Expect(err).ToNot(HaveOccurred()) @@ -632,7 +642,7 @@ parameters: response := getModelStatus("http://" + testHTTPAddr + "/models/jobs/" + uuid) resp = response return response["processed"].(bool) - }, "360s", "10s").Should(Equal(true)) + }, "360s", "50ms").Should(Equal(true)) // Check that the model was imported successfully Expect(resp["message"]).ToNot(ContainSubstring("error")) @@ -703,7 +713,7 @@ parameters: response := getModelStatus("http://" + testHTTPAddr + "/models/jobs/" + uuid) resp = response return response["processed"].(bool) - }, "360s", "10s").Should(Equal(true)) + }, "360s", "50ms").Should(Equal(true)) // Check that the model was imported successfully Expect(resp["message"]).To(ContainSubstring("error")) diff --git a/core/http/endpoints/localai/import_model_test.go b/core/http/endpoints/localai/import_model_test.go index 96c20160e7e8..2b2bbede2e64 100644 --- a/core/http/endpoints/localai/import_model_test.go +++ b/core/http/endpoints/localai/import_model_test.go @@ -10,14 +10,22 @@ import ( "github.com/labstack/echo/v4" "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/gallery/importers" . "github.com/mudler/LocalAI/core/http/endpoints/localai" "github.com/mudler/LocalAI/core/services/galleryop" + hfapi "github.com/mudler/LocalAI/pkg/huggingface-api" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/system" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) +type ambiguityMetadata struct{} + +func (ambiguityMetadata) GetModelDetails(repo string) (*hfapi.ModelDetails, error) { + return &hfapi.ModelDetails{ModelID: repo, Author: "nari-labs", PipelineTag: "text-to-speech"}, nil +} + var _ = Describe("ImportModelURIEndpoint ambiguity handling", func() { var ( @@ -26,6 +34,9 @@ var _ = Describe("ImportModelURIEndpoint ambiguity handling", func() { ) BeforeEach(func() { + restore := importers.SetHuggingFaceMetadataFactoryForTest(func() importers.HuggingFaceMetadata { return ambiguityMetadata{} }) + DeferCleanup(restore) + var err error tempDir, err = os.MkdirTemp("", "import-model-test") Expect(err).ToNot(HaveOccurred()) diff --git a/core/http/endpoints/localai/localai_suite_test.go b/core/http/endpoints/localai/localai_suite_test.go index ea415bf70008..fd64e6dd7201 100644 --- a/core/http/endpoints/localai/localai_suite_test.go +++ b/core/http/endpoints/localai/localai_suite_test.go @@ -3,6 +3,7 @@ package localai_test import ( "testing" + "github.com/mudler/LocalAI/core/services/testutil" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) @@ -11,3 +12,13 @@ func TestLocalAIEndpoints(t *testing.T) { RegisterFailHandler(Fail) RunSpecs(t, "LocalAI Endpoints test suite") } + +var _ = SynchronizedBeforeSuite(func() []byte { + return []byte(testutil.StartSharedTestDB()) +}, func(endpoint []byte) { + testutil.SetSharedTestDBEndpoint(string(endpoint)) +}) + +var _ = SynchronizedAfterSuite(func() {}, func() { + testutil.StopSharedTestDB() +}) diff --git a/core/http/endpoints/localai/router_decide_test.go b/core/http/endpoints/localai/router_decide_test.go index d63cd091ddcb..1efd11dcf870 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/ollama/helpers_internal_test.go b/core/http/endpoints/ollama/helpers_internal_test.go index 356c19a5b7ae..92469eb8f1ae 100644 --- a/core/http/endpoints/ollama/helpers_internal_test.go +++ b/core/http/endpoints/ollama/helpers_internal_test.go @@ -30,7 +30,7 @@ var _ = Describe("applyOllamaOptions num_ctx clamping (issue #11022)", func() { It("does not let an oversized num_ctx raise an existing context ceiling", func() { for _, ceiling := range []int{4096, 8192} { existing := ceiling - cfg := &config.ModelConfig{ContextSize: &existing} + cfg := &config.ModelConfig{LLMConfig: config.LLMConfig{ContextSize: &existing}} applyOllamaOptions(&schema.OllamaOptions{NumCtx: 2000000000}, cfg) Expect(cfg.ContextSize).ToNot(BeNil()) @@ -41,7 +41,7 @@ var _ = Describe("applyOllamaOptions num_ctx clamping (issue #11022)", func() { It("still honors a num_ctx smaller than the existing ceiling", func() { existing := 8192 - cfg := &config.ModelConfig{ContextSize: &existing} + cfg := &config.ModelConfig{LLMConfig: config.LLMConfig{ContextSize: &existing}} applyOllamaOptions(&schema.OllamaOptions{NumCtx: 2048}, cfg) Expect(cfg.ContextSize).ToNot(BeNil()) @@ -62,4 +62,4 @@ var _ = Describe("applyOllamaOptions num_ctx clamping (issue #11022)", func() { Expect(cfg.ContextSize).To(BeNil()) }) -}) \ No newline at end of file +}) diff --git a/core/http/endpoints/openai/realtime.go b/core/http/endpoints/openai/realtime.go index cdff7aec531a..5167af5a4e28 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" @@ -137,6 +138,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 @@ -197,14 +204,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, @@ -266,6 +274,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 @@ -540,6 +566,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 @@ -589,6 +621,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. @@ -751,7 +787,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 } @@ -777,7 +815,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 } @@ -953,6 +991,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 { @@ -1228,6 +1270,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 } @@ -1301,6 +1354,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 @@ -1362,183 +1432,290 @@ 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() + vadTick(sink, silenceThreshold) + } + } +} - // 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() - } +// 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() - session.AudioBufferLock.Lock() - allAudio := make([]byte, len(session.InputAudioBuffer)) - copy(allAudio, session.InputAudioBuffer) - session.AudioBufferLock.Unlock() + // The turn coordinator's data-heavy effects (OpenTurn/CommitTurn) + // need this tick's mode; set it before any Apply below. + sink.sv = sv - aints := sound.BytesToInt16sLE(allAudio) - if len(aints) == 0 || len(aints) < int(silenceThreshold*float64(session.InputSampleRate)) { - continue - } + // 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() + } - // Resample from InputSampleRate to 16kHz - aints = sound.ResampleInt16(aints, session.InputSampleRate, localSampleRate) + 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() - audioLength := float64(len(aints)) / localSampleRate + aints := sound.BytesToInt16sLE(allAudio) + if len(aints) == 0 || len(aints) < int(silenceThreshold*float64(session.InputSampleRate)) { + return + } - if sv != nil && lts.open() { - lts.feedNewAudio(aints) - lts.drainEvents(audioLength) - } + // Resample from InputSampleRate to 16kHz + aints = sound.ResampleInt16(aints, session.InputSampleRate, localSampleRate) - 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 - } + audioLength := float64(len(aints)) / localSampleRate - // 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() + if sv != nil && lts.open() { + lts.feedNewAudio(aints) + lts.drainEvents(audioLength) + } - // 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 - } + // 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 + } - // 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) - } + 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() - 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) + // 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 + } - // Segment still in progress when audio ended - segEndTime := segments[len(segments)-1].End - if segEndTime == 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) + } + + // 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 + } - threshold := silenceThreshold - eouPending := false - if sv != nil { - eouPending = lts.eouPending(segments) - threshold = lts.thresholdSec(eouPending, sv) - } + 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) + } - 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 - } - } + // Segment still in progress when audio ended + segEndTime := segments[len(segments)-1].End + if segEndTime == 0 { + return + } - 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() + threshold := silenceThreshold + eouPending := false + if sv != nil { + eouPending = lts.eouPending(segments) + threshold = lts.thresholdSec(eouPending, sv) + } - // 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) - } - } + 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) { @@ -1883,6 +2060,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 @@ -2203,9 +2385,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 @@ -2358,6 +2549,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{ @@ -2425,7 +2640,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). @@ -2455,12 +2670,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) @@ -2519,11 +2734,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 000000000000..c310f3f001a2 --- /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 000000000000..b8630ca8120b --- /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 fe52e1c645bb..a2c104b3c29b 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 0449daee3740..030a0c914787 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 66dfc6efe2dd..75a71ba25a85 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_turncoord.go b/core/http/endpoints/openai/realtime_turncoord.go index 30ffffc6680f..f0d599f7ea2c 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 000000000000..43f26731c6d6 --- /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 475b45e8f2e6..c9b78ae9da58 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 000000000000..43749d3c05be --- /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 000000000000..57e1ab66e1c2 --- /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 6b0a233eebbe..b847a35a75d7 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 2f75486adcc3..a18fa3645dbb 100644 --- a/core/http/endpoints/openai/types/types.go +++ b/core/http/endpoints/openai/types/types.go @@ -949,6 +949,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 { @@ -1180,6 +1185,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 000000000000..ad2a1c5ce8eb --- /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/middleware/route_model.go b/core/http/middleware/route_model.go index 470bd05f5aa2..643c6fa12d45 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 4a9be2b12fb3..0496ed9ddbf5 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/openresponses_test.go b/core/http/openresponses_test.go index f30674362534..ab5dd716a7df 100644 --- a/core/http/openresponses_test.go +++ b/core/http/openresponses_test.go @@ -28,7 +28,7 @@ import ( // the registered model is "Qwen3-VL-2B-Instruct-Q4_K_M", not the repo name. const testModel = "Qwen3-VL-2B-Instruct-Q4_K_M" -var _ = Describe("Open Responses API", func() { +var _ = Describe("Open Responses API", Serial, func() { var app *echo.Echo var localApp *application.Application var localModelDir string diff --git a/core/services/agentpool/services_suite_test.go b/core/services/agentpool/services_suite_test.go index 7a60db0d4c32..ecca27ad42d7 100644 --- a/core/services/agentpool/services_suite_test.go +++ b/core/services/agentpool/services_suite_test.go @@ -3,6 +3,7 @@ package agentpool_test import ( "testing" + "github.com/mudler/LocalAI/core/services/testutil" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) @@ -11,3 +12,13 @@ func TestServices(t *testing.T) { RegisterFailHandler(Fail) RunSpecs(t, "LocalAI services test") } + +var _ = SynchronizedBeforeSuite(func() []byte { + return []byte(testutil.StartSharedTestDB()) +}, func(endpoint []byte) { + testutil.SetSharedTestDBEndpoint(string(endpoint)) +}) + +var _ = SynchronizedAfterSuite(func() {}, func() { + testutil.StopSharedTestDB() +}) diff --git a/core/services/agents/agents_suite_test.go b/core/services/agents/agents_suite_test.go index 6cc46193b97f..29f76733cdd2 100644 --- a/core/services/agents/agents_suite_test.go +++ b/core/services/agents/agents_suite_test.go @@ -3,6 +3,7 @@ package agents import ( "testing" + "github.com/mudler/LocalAI/core/services/testutil" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) @@ -11,3 +12,13 @@ func TestAgents(t *testing.T) { RegisterFailHandler(Fail) RunSpecs(t, "Agents test suite") } + +var _ = SynchronizedBeforeSuite(func() []byte { + return []byte(testutil.StartSharedTestDB()) +}, func(endpoint []byte) { + testutil.SetSharedTestDBEndpoint(string(endpoint)) +}) + +var _ = SynchronizedAfterSuite(func() {}, func() { + testutil.StopSharedTestDB() +}) diff --git a/core/services/cloudproxy/mitm/proxy.go b/core/services/cloudproxy/mitm/proxy.go index 79f49aa648f2..4ccd744f4bae 100644 --- a/core/services/cloudproxy/mitm/proxy.go +++ b/core/services/cloudproxy/mitm/proxy.go @@ -23,15 +23,17 @@ import ( // in its intercept allowlist; non-allowlisted hosts get a plain // TCP CONNECT tunnel. type Server struct { - addr string - ca *CA - interceptHosts map[string]bool - handler InterceptHandler - connectTimeout time.Duration - dialTimeout time.Duration - upstreamTLS *tls.Config - events pii.EventStore - eventSeq atomic.Uint64 + addr string + ca *CA + interceptHosts map[string]bool + handler InterceptHandler + connectTimeout time.Duration + dialTimeout time.Duration + upstreamTLS *tls.Config + events pii.EventStore + eventSeq atomic.Uint64 + allowPlainHTTP bool + interceptAll bool listener net.Listener srv *http.Server @@ -51,6 +53,12 @@ type Config struct { CA *CA InterceptHosts []string Handler InterceptHandler + // AllowPlainHTTP is used by the deterministic test-resource proxy. + // Production listeners leave it false and continue to require CONNECT. + AllowPlainHTTP bool + // InterceptAll prevents undeclared HTTPS hosts from being tunnelled by + // strict test-resource replay. Production listeners use the host allowlist. + InterceptAll bool // EventStore optionally receives a proxy_connect event for every // CONNECT, recording the destination host and whether the proxy // intercepted or tunneled it. nil disables connect-event recording. @@ -73,6 +81,8 @@ func NewServer(cfg Config) (*Server, error) { ca: cfg.CA, interceptHosts: hosts, handler: cfg.Handler, + allowPlainHTTP: cfg.AllowPlainHTTP, + interceptAll: cfg.InterceptAll, connectTimeout: 30 * time.Second, dialTimeout: 15 * time.Second, upstreamTLS: &tls.Config{NextProtos: []string{"http/1.1"}}, @@ -126,6 +136,10 @@ func (s *Server) Stop() { func (s *Server) handle(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodConnect { + if s.allowPlainHTTP && r.URL != nil && r.URL.IsAbs() { + s.handler(w, r, r.URL.Host) + return + } http.Error(w, "this proxy only supports HTTPS via CONNECT", http.StatusMethodNotAllowed) return } @@ -168,6 +182,9 @@ func (s *Server) recordConnectEvent(host string, intercepted bool) { // shouldIntercept reports whether host is in the allowlist. An // empty allowlist tunnels everything. func (s *Server) shouldIntercept(host string) bool { + if s.interceptAll { + return true + } if len(s.interceptHosts) == 0 { return false } diff --git a/core/services/jobs/jobs_suite_test.go b/core/services/jobs/jobs_suite_test.go index 957e5ead7fb7..c183767e8d94 100644 --- a/core/services/jobs/jobs_suite_test.go +++ b/core/services/jobs/jobs_suite_test.go @@ -3,6 +3,7 @@ package jobs import ( "testing" + "github.com/mudler/LocalAI/core/services/testutil" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) @@ -11,3 +12,13 @@ func TestJobs(t *testing.T) { RegisterFailHandler(Fail) RunSpecs(t, "Jobs test suite") } + +var _ = SynchronizedBeforeSuite(func() []byte { + return []byte(testutil.StartSharedTestDB()) +}, func(endpoint []byte) { + testutil.SetSharedTestDBEndpoint(string(endpoint)) +}) + +var _ = SynchronizedAfterSuite(func() {}, func() { + testutil.StopSharedTestDB() +}) diff --git a/core/services/nodes/nodes_suite_test.go b/core/services/nodes/nodes_suite_test.go index a6a24852b00e..56cf49a68150 100644 --- a/core/services/nodes/nodes_suite_test.go +++ b/core/services/nodes/nodes_suite_test.go @@ -3,6 +3,7 @@ package nodes import ( "testing" + "github.com/mudler/LocalAI/core/services/testutil" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) @@ -11,3 +12,13 @@ func TestNodes(t *testing.T) { RegisterFailHandler(Fail) RunSpecs(t, "Nodes test suite") } + +var _ = SynchronizedBeforeSuite(func() []byte { + return []byte(testutil.StartSharedTestDB()) +}, func(endpoint []byte) { + testutil.SetSharedTestDBEndpoint(string(endpoint)) +}) + +var _ = SynchronizedAfterSuite(func() {}, func() { + testutil.StopSharedTestDB() +}) diff --git a/core/services/routing/router/score.go b/core/services/routing/router/score.go index 34beeb5c6088..657897ec77fa 100644 --- a/core/services/routing/router/score.go +++ b/core/services/routing/router/score.go @@ -6,6 +6,7 @@ import ( "fmt" "math" "strings" + "sync" "text/template" "time" @@ -98,6 +99,11 @@ type ScoreClassifierOptions struct { // sends Probe.Prompt as-is and relies on the backend's n_ctx guard. TokenCounter func(string) (int, error) MaxContextTokens int + + // CompletionReserveTokens reserves additional context beyond the longest + // scoring candidate. Classifier slot filling uses this to ensure the prompt + // scored here can be continued without overflowing the model context. + CompletionReserveTokens int } // ScoreClassifier scores every policy label as the model's actual @@ -139,6 +145,17 @@ type ScoreClassifier struct { budget *lazyBudget cache *labelSetCache + + // stablePrefix is the rendered-prompt prefix shared by every probe: + // the chat template's preamble plus the option-list system prompt, + // up to where the per-turn text begins. Computed once (the byte-wise + // common prefix of two synthetic probes) and sent with each Score + // call as a state-reuse boundary hint — on backends whose models + // cannot rewind state (hybrid/recurrent), a snapshot at this + // boundary is what keeps repeat scoring at probe-size cost instead + // of a full option-list re-prefill. + stablePrefixOnce sync.Once + stablePrefix string } // NewScoreClassifier panics on caller errors at construction (empty @@ -202,8 +219,13 @@ func NewScoreClassifier(policies []ScorePolicy, scorer backend.Scorer, opts Scor systemPrompt: systemPrompt, labelOrder: labels, candidates: candidates, - budget: &lazyBudget{tokenize: opts.TokenCounter, maxContext: opts.MaxContextTokens, extras: candidates}, - cache: newLabelSetCache(opts.CacheCap), + budget: &lazyBudget{ + tokenize: opts.TokenCounter, + maxContext: opts.MaxContextTokens, + extras: candidates, + reserve: opts.CompletionReserveTokens, + }, + cache: newLabelSetCache(opts.CacheCap), } } @@ -228,25 +250,76 @@ func renderSystemPrompt(tmpl string, policies []ScorePolicy) (string, error) { func (c *ScoreClassifier) Name() string { return ClassifierScore } -func (c *ScoreClassifier) Classify(ctx context.Context, p Probe) (Decision, error) { - start := time.Now() - +// renderProbe returns the exact prompt Classify scores for p (system +// prompt + trimmed user turns through the model's chat template), plus the +// trimmed user text used as the memo-cache key. +func (c *ScoreClassifier) renderProbe(p Probe) (prompt, userText string, err error) { // Trim oldest turns until the rendered prompt fits the classifier's // context. Cache-keyed on the trimmed text so conversations that // trim to the same tail share an entry. - userText := trimmedProbeText(p, c.budget, func(joined string) (string, error) { + userText = trimmedProbeText(p, c.budget, func(joined string) (string, error) { return c.renderer(c.systemPrompt, joined) }) + prompt, err = c.renderer(c.systemPrompt, userText) + return prompt, userText, err +} - key := cacheKey(userText) - if hit, ok := c.cache.get(key); ok { - return Decision{Labels: hit, Score: 1.0, Latency: time.Since(start)}, nil +// stablePrefixLen returns the byte length of prompt's leading run that +// is invariant across probes, by rendering two synthetic probes with no +// common text and taking their byte-wise common prefix. Clamped against +// the actual prompt so a template that (unexpectedly) varies its +// preamble degrades to a shorter hint, never a wrong one. +func (c *ScoreClassifier) stablePrefixLen(prompt string) int { + c.stablePrefixOnce.Do(func() { + a, aErr := c.renderer(c.systemPrompt, "\x02") + b, bErr := c.renderer(c.systemPrompt, "\x03") + if aErr != nil || bErr != nil { + return + } + n := 0 + for n < len(a) && n < len(b) && a[n] == b[n] { + n++ + } + c.stablePrefix = a[:n] + }) + n := 0 + limit := min(len(c.stablePrefix), len(prompt)) + for n < limit && prompt[n] == c.stablePrefix[n] { + n++ } - prompt, err := c.renderer(c.systemPrompt, userText) + return n +} + +// SlotFillPrompt returns the completion prompt for filling a chosen +// label's argument slots: the identical prompt Classify scored — so the +// backend's prompt cache is already warm with it — continued by the +// label's route JSON re-opened at its first slot field: +// +// …{"route": "up", "distance": +// +// The caller constrains the remaining tokens with a grammar and closes +// the object; keeping the field style byte-identical to the scored +// candidates keeps the continuation on-distribution. +func (c *ScoreClassifier) SlotFillPrompt(p Probe, label, firstSlot string) (string, error) { + prompt, _, err := c.renderProbe(p) + if err != nil { + return "", fmt.Errorf("score slot fill: render prompt: %w", err) + } + return prompt + `{"route": "` + escapeJSONString(label) + `", "` + escapeJSONString(firstSlot) + `": `, nil +} + +func (c *ScoreClassifier) Classify(ctx context.Context, p Probe) (Decision, error) { + start := time.Now() + + prompt, userText, err := c.renderProbe(p) if err != nil { return errDecision(start, fmt.Errorf("score classify: render prompt: %w", err)) } - results, err := c.scorer.Score(ctx, prompt, c.candidates) + key := cacheKey(userText) + if hit, ok := c.cache.get(key); ok { + return Decision{Labels: hit, Score: 1.0, Latency: time.Since(start)}, nil + } + results, err := c.scorer.Score(ctx, prompt, c.stablePrefixLen(prompt), c.candidates) if err != nil { xlog.Warn("router: score classifier failed", "error", err, "labels", c.labelOrder) return errDecision(start, fmt.Errorf("score classify: %w", err)) diff --git a/core/services/routing/router/score_test.go b/core/services/routing/router/score_test.go index 75707186efdd..d180115892c1 100644 --- a/core/services/routing/router/score_test.go +++ b/core/services/routing/router/score_test.go @@ -17,13 +17,15 @@ type stubScorer struct { results []backend.CandidateScore err error calls int - lastP string - lastC []string + lastP string + lastC []string + lastStable int } -func (s *stubScorer) Score(_ context.Context, prompt string, candidates []string) ([]backend.CandidateScore, error) { +func (s *stubScorer) Score(_ context.Context, prompt string, stablePrefixLen int, candidates []string) ([]backend.CandidateScore, error) { s.calls++ s.lastP = prompt + s.lastStable = stablePrefixLen s.lastC = append(s.lastC[:0], candidates...) if s.err != nil { return nil, s.err @@ -74,6 +76,26 @@ var _ = Describe("ScoreClassifier", func() { Expect(d.Score).To(BeNumerically(">=", 0.8), "want >= 0.8 for dominant single label") }) + It("reports the probe-invariant stable prompt prefix to the scorer", func() { + s := &stubScorer{results: []backend.CandidateScore{ + {LogProb: -0.05, NumTokens: 6}, + {LogProb: -8.0, NumTokens: 6}, + {LogProb: -10.0, NumTokens: 6}, + }} + c := NewScoreClassifier(testPolicies(), s, ScoreClassifierOptions{}) + _, err := c.Classify(context.Background(), Probe{Prompt: "first probe text"}) + Expect(err).NotTo(HaveOccurred()) + first, firstStable := s.lastP, s.lastStable + Expect(firstStable).To(BeNumerically(">", 0), "stable prefix must be non-empty") + Expect(firstStable).To(BeNumerically("<", len(first)), "stable prefix must not cover the probe") + + _, err = c.Classify(context.Background(), Probe{Prompt: "a wholly different request"}) + Expect(err).NotTo(HaveOccurred()) + Expect(s.lastStable).To(Equal(firstStable), "stable prefix is probe-invariant") + Expect(s.lastP[:firstStable]).To(Equal(first[:firstStable]), "prompts share the declared prefix") + Expect(first[firstStable:]).To(ContainSubstring("first probe text"), "the probe lives beyond the stable prefix") + }) + It("activates multiple labels", func() { // Raw mode two-way tie: code and math share ~0.49 each, chat // far behind. Both must activate so the router can pick a @@ -368,6 +390,19 @@ var _ = Describe("ScoreClassifier conversation trimming", func() { Expect(len(strings.Fields(s.lastP))).To(BeNumerically("<", 20000), "must be trimmed, not the full transcript") }) + It("reserves context for a completion after scoring", func() { + without := NewScoreClassifier(testPolicies(), &stubScorer{}, ScoreClassifierOptions{ + TokenCounter: wordCount, + MaxContextTokens: 10000, + }) + with := NewScoreClassifier(testPolicies(), &stubScorer{}, ScoreClassifierOptions{ + TokenCounter: wordCount, + MaxContextTokens: 10000, + CompletionReserveTokens: 257, + }) + Expect(with.probeTokenBudget()).To(Equal(without.probeTokenBudget() - 257)) + }) + It("keeps the newest turn whole even when it alone exceeds the budget", func() { s := &stubScorer{results: threeScores} c := NewScoreClassifier(testPolicies(), s, ScoreClassifierOptions{ diff --git a/core/services/routing/router/trim.go b/core/services/routing/router/trim.go index 50f752d9d0b7..b142373c85c0 100644 --- a/core/services/routing/router/trim.go +++ b/core/services/routing/router/trim.go @@ -123,6 +123,7 @@ type lazyBudget struct { tokenize func(string) (int, error) maxContext int extras []string + reserve int mu sync.Mutex value atomic.Int64 // 0=unset, >0=budget, -1=disabled @@ -156,7 +157,7 @@ func (l *lazyBudget) get() int { longest = n } } - b := l.maxContext - longest - tokenBudgetMargin + b := l.maxContext - longest - l.reserve - tokenBudgetMargin if b <= 0 { l.value.Store(-1) return 0 diff --git a/core/services/testutil/testdb.go b/core/services/testutil/testdb.go index 80e511201b7d..c5b7597fd23c 100644 --- a/core/services/testutil/testdb.go +++ b/core/services/testutil/testdb.go @@ -2,11 +2,16 @@ package testutil import ( "context" + "fmt" "runtime" + "sync" + "sync/atomic" "time" + "github.com/mudler/LocalAI/internal/testfixtures" "github.com/testcontainers/testcontainers-go" tcpostgres "github.com/testcontainers/testcontainers-go/modules/postgres" + tcnetwork "github.com/testcontainers/testcontainers-go/network" "github.com/testcontainers/testcontainers-go/wait" "gorm.io/driver/postgres" "gorm.io/gorm" @@ -16,24 +21,110 @@ import ( . "github.com/onsi/gomega" ) -// SetupTestDB creates a fresh PostgreSQL 16 container and returns a gorm.DB. -// The container is cleaned up via DeferCleanup when the test completes. +var ( + sharedDBMu sync.Mutex + sharedDBContainer testcontainers.Container + sharedDBEndpoint string + sharedDBSequence atomic.Uint64 +) + +// StartSharedTestDB starts one PostgreSQL container and returns its endpoint. +// Pass that endpoint to SetSharedTestDBEndpoint in every parallel test process. +func StartSharedTestDB() string { + if runtime.GOOS == "darwin" { + return "" + } + sharedDBMu.Lock() + defer sharedDBMu.Unlock() + if sharedDBContainer != nil { + return sharedDBEndpoint + } + + container, endpoint := startTestDBContainer() + sharedDBContainer = container + sharedDBEndpoint = endpoint + return endpoint +} + +// SetSharedTestDBEndpoint attaches this test process to the suite database. +func SetSharedTestDBEndpoint(endpoint string) { + sharedDBMu.Lock() + defer sharedDBMu.Unlock() + sharedDBEndpoint = endpoint +} + +// StopSharedTestDB terminates the process-scoped PostgreSQL fixture. +func StopSharedTestDB() { + sharedDBMu.Lock() + defer sharedDBMu.Unlock() + if sharedDBContainer == nil { + return + } + Expect(sharedDBContainer.Terminate(context.Background())).To(Succeed()) + sharedDBContainer = nil + sharedDBEndpoint = "" +} + +// SetupTestDB returns an isolated PostgreSQL database fixture. Suites that call +// StartSharedTestDB get a fresh schema; other suites retain a fresh container. func SetupTestDB() *gorm.DB { if runtime.GOOS == "darwin" { Skip("testcontainers requires Docker, not available on macOS CI") } + + sharedDBMu.Lock() + endpoint := sharedDBEndpoint + sharedDBMu.Unlock() + if endpoint != "" { + return setupIsolatedSchema(endpoint) + } + + pgC, endpoint := startTestDBContainer() + DeferCleanup(func() { _ = pgC.Terminate(context.Background()) }) + return openTestDB(endpoint, "") +} + +func startTestDBContainer() (testcontainers.Container, string) { ctx := context.Background() - pgC, err := tcpostgres.Run(ctx, "postgres:16", + Expect(testfixtures.RequireImage(ctx, testfixtures.Postgres16, "default")).To(Succeed()) + testNetwork, err := testfixtures.DockerNetwork() + Expect(err).NotTo(HaveOccurred()) + pgC, err := tcpostgres.Run(ctx, testfixtures.Postgres16, tcpostgres.WithDatabase("testdb"), tcpostgres.WithUsername("test"), tcpostgres.WithPassword("test"), testcontainers.WithWaitStrategyAndDeadline(60*time.Second, wait.ForLog("database system is ready to accept connections").WithOccurrence(2)), + tcnetwork.WithNetworkName([]string{"postgres"}, testNetwork), ) Expect(err).ToNot(HaveOccurred()) - DeferCleanup(func() { pgC.Terminate(context.Background()) }) - connStr, err := pgC.ConnectionString(ctx, "sslmode=disable") + endpoint, err := testfixtures.ContainerEndpoint(ctx, pgC, "5432") Expect(err).ToNot(HaveOccurred()) + return pgC, endpoint +} + +func setupIsolatedSchema(endpoint string) *gorm.DB { + schema := fmt.Sprintf("test_%d_%d", GinkgoParallelProcess(), sharedDBSequence.Add(1)) + admin := openTestDB(endpoint, "") + Expect(admin.Exec("CREATE SCHEMA " + schema).Error).ToNot(HaveOccurred()) + db := openTestDB(endpoint, schema) + DeferCleanup(func() { + if sqlDB, err := db.DB(); err == nil { + _ = sqlDB.Close() + } + Expect(admin.Exec("DROP SCHEMA " + schema + " CASCADE").Error).ToNot(HaveOccurred()) + if sqlDB, err := admin.DB(); err == nil { + _ = sqlDB.Close() + } + }) + return db +} + +func openTestDB(endpoint, schema string) *gorm.DB { + connStr := fmt.Sprintf("postgres://test:test@%s/testdb?sslmode=disable", endpoint) + if schema != "" { + connStr += "&search_path=" + schema + } db, err := gorm.Open(postgres.Open(connStr), &gorm.Config{ Logger: logger.Default.LogMode(logger.Silent), }) diff --git a/core/services/worker/free_timeout_test.go b/core/services/worker/free_timeout_test.go index 4f1b6346e749..f27d4dcf3bf1 100644 --- a/core/services/worker/free_timeout_test.go +++ b/core/services/worker/free_timeout_test.go @@ -6,6 +6,8 @@ import ( "os" "strconv" "syscall" + "testing" + "time" process "github.com/mudler/go-processmanager" gogrpc "google.golang.org/grpc" @@ -16,6 +18,19 @@ import ( . "github.com/onsi/gomega" ) +// TestWorkerFixtureProcess turns the current test binary into a portable +// long-running child for the process-stop assertions below. Using the test +// binary avoids assuming Unix utilities live at paths such as /bin/sleep, +// which is not true in Nix environments. +func TestWorkerFixtureProcess(t *testing.T) { + if os.Getenv("LOCALAI_WORKER_FIXTURE_PROCESS") != "1" { + return + } + for { + time.Sleep(time.Hour) + } +} + // pidAlive probes the OS directly for a process ID. The supervisor's own // liveness helpers all go through go-processmanager's pidfile, which Stop // deletes as part of releasing the handle, so they report "not alive" even if @@ -82,10 +97,13 @@ var _ = Describe("Stopping a backend whose Free never returns", func() { // actually dead afterwards, not merely that Stop() returned. It // outlives every timeout below, so if it is gone at the end it is // because the supervisor signalled it. + executable, err := os.Executable() + Expect(err).ToNot(HaveOccurred()) proc = process.New( process.WithTemporaryStateDir(), - process.WithName("/bin/sleep"), - process.WithArgs("300"), + process.WithName(executable), + process.WithArgs("-test.run=^TestWorkerFixtureProcess$"), + process.WithEnvironment(append(os.Environ(), "LOCALAI_WORKER_FIXTURE_PROCESS=1")...), ) Expect(proc.Run()).To(Succeed()) @@ -94,7 +112,8 @@ var _ = Describe("Stopping a backend whose Free never returns", func() { Expect(pidAlive(procPID)).To(BeTrue(), "the fixture process must be running before the stop") s = &backendSupervisor{ - cfg: &Config{}, + cfg: &Config{}, + backendFreeTimeout: 20 * time.Millisecond, processes: map[string]*backendProcess{ "wedged-model#0": { proc: proc, diff --git a/core/services/worker/lifecycle.go b/core/services/worker/lifecycle.go index 904a7a7f5db9..731475587ced 100644 --- a/core/services/worker/lifecycle.go +++ b/core/services/worker/lifecycle.go @@ -319,7 +319,7 @@ func (s *backendSupervisor) handleModelUnload(data []byte, reply func([]byte)) { // Best-effort bounded gRPC Free(). A model.unload request must not // occupy the NATS reply handler forever when a backend is wedged. client := grpc.NewClientWithToken(targetAddr, false, nil, false, s.cfg.RegistrationToken) - freeCtx, cancel := context.WithTimeout(context.Background(), workerBackendFreeTimeout) + freeCtx, cancel := context.WithTimeout(context.Background(), s.freeTimeout()) if err := client.Free(freeCtx); err != nil { xlog.Warn("Free() failed during model.unload", "error", err, "addr", targetAddr) } diff --git a/core/services/worker/supervisor.go b/core/services/worker/supervisor.go index dab96e4ab235..7313556b47ed 100644 --- a/core/services/worker/supervisor.go +++ b/core/services/worker/supervisor.go @@ -119,6 +119,11 @@ type backendSupervisor struct { // the same not-yet-cached backend) are serialized here so the gallery // download path doesn't race itself on the same directory. backendLocks map[string]*sync.Mutex + + // backendFreeTimeout bounds the best-effort Free call before process + // termination. Zero uses workerBackendFreeTimeout; tests use a shorter + // deadline to exercise a wedged backend without waiting five seconds. + backendFreeTimeout time.Duration } // defaultPortQuarantine is how long a released gRPC port waits before it can be @@ -141,6 +146,13 @@ type backendSupervisor struct { // rows; raising this value is not a substitute for it. const defaultPortQuarantine = 15 * time.Second +func (s *backendSupervisor) freeTimeout() time.Duration { + if s.backendFreeTimeout > 0 { + return s.backendFreeTimeout + } + return workerBackendFreeTimeout +} + // quarantinedPort is a released port that must not be re-bound until `until`. type quarantinedPort struct { port int @@ -723,8 +735,9 @@ func (s *backendSupervisor) stopBackendExact(key string, force bool) error { if !force { client := grpc.NewClientWithToken(bp.addr, false, nil, false, s.cfg.RegistrationToken) - freeCtx, cancel := context.WithTimeout(context.Background(), workerBackendFreeTimeout) - xlog.Debug("Calling bounded Free() before stopping backend", "backend", key, "timeout", workerBackendFreeTimeout) + freeTimeout := s.freeTimeout() + freeCtx, cancel := context.WithTimeout(context.Background(), freeTimeout) + xlog.Debug("Calling bounded Free() before stopping backend", "backend", key, "timeout", freeTimeout) if err := client.Free(freeCtx); err != nil { xlog.Warn("Free() failed (best-effort)", "backend", key, "error", err) } diff --git a/core/startup/model_preload.go b/core/startup/model_preload.go index 4f3bb16832d5..bd3739737a6b 100644 --- a/core/startup/model_preload.go +++ b/core/startup/model_preload.go @@ -75,13 +75,18 @@ func InstallModelsWithOptions(ctx context.Context, galleryService *galleryop.Gal } var status *galleryop.OpStatus - // wait for op to finish + poll := time.NewTicker(50 * time.Millisecond) + defer poll.Stop() for { status = galleryService.GetStatus(uuid.String()) if status != nil && status.Processed { break } - time.Sleep(1 * time.Second) + select { + case <-ctx.Done(): + return ctx.Err() + case <-poll.C: + } } if status.Error != nil { diff --git a/core/startup/model_preload_test.go b/core/startup/model_preload_test.go index 525f183cfa88..ad662a5fa1b3 100644 --- a/core/startup/model_preload_test.go +++ b/core/startup/model_preload_test.go @@ -3,6 +3,8 @@ package startup_test import ( "context" "fmt" + "net/http" + "net/http/httptest" "os" "path/filepath" @@ -39,7 +41,11 @@ var _ = Describe("Preload test", func() { Context("Preloading from strings", func() { It("loads from embedded full-urls", func() { - url := "https://raw.githubusercontent.com/mudler/LocalAI-examples/main/configurations/phi-2.yaml" + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte("name: phi-2\nbackend: llama-cpp\nparameters:\n model: phi-2.gguf\n")) + })) + defer server.Close() + url := server.URL + "/phi-2.yaml" fileName := fmt.Sprintf("%s.yaml", "phi-2") galleryService := galleryop.NewGalleryService(&config.ApplicationConfig{ @@ -59,7 +65,11 @@ var _ = Describe("Preload test", func() { Expect(string(content)).To(ContainSubstring("name: phi-2")) }) It("downloads from urls", func() { - url := "huggingface://TheBloke/TinyLlama-1.1B-Chat-v0.3-GGUF/tinyllama-1.1b-chat-v0.3.Q2_K.gguf" + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte("tiny local GGUF fixture")) + })) + defer server.Close() + url := server.URL + "/tinyllama-1.1b-chat-v0.3.Q2_K.gguf" fileName := fmt.Sprintf("%s.gguf", "tinyllama-1.1b-chat-v0.3.Q2_K") galleryService := galleryop.NewGalleryService(&config.ApplicationConfig{ diff --git a/docs/content/advanced/model-configuration.md b/docs/content/advanced/model-configuration.md index 3e9c1bd16c8f..a162a6225319 100644 --- a/docs/content/advanced/model-configuration.md +++ b/docs/content/advanced/model-configuration.md @@ -200,7 +200,7 @@ These settings apply to most LLM backends (llama.cpp, vLLM, etc.): | Field | Type | Default | Description | |-------|------|---------|-------------| -| `threads` | int | `processor count` | Number of threads for parallel computation | +| `threads` | int | `processor count` | Number of threads for parallel computation. A per-model value overrides the server-wide `--threads`/`LOCALAI_THREADS` setting | | `context_size` | int | `512` | Maximum context size in tokens. Set to `-1` to auto-use the model's full trained context from GGUF metadata (raw max, no VRAM capping; a warning is logged if it may not fit detected VRAM). | | `f16` | bool | `false` | Enable 16-bit floating point precision (GPU acceleration) | | `gpu_layers` | int | `0` | Number of layers to offload to GPU (0 = CPU only) | @@ -911,6 +911,8 @@ Define pipelines for audio-to-audio processing and the [Realtime API]({{%relref | `pipeline.llm` | string | LLM model name | | `pipeline.transcription` | string | Transcription model name | | `pipeline.vad` | string | Voice activity detection model name | +| `pipeline.turn_detection` | object | Realtime turn-detection defaults. Keys: `type` (`server_vad`/`semantic_vad`), `eagerness` (`low`/`medium`/`high`/`auto`), `retranscribe`, `vad_window_sec` (widen the per-tick VAD scan window; values below the automatic floor are ignored). See [Realtime turn detection]({{%relref "features/openai-realtime" %}}) | +| `pipeline.classifier` | object | Realtime classifier mode: prefill-scored option selection instead of generation. Keys: `enabled`, `threshold`, `normalization` (`raw`/`mean`), `history_items`, `fallback` (`mode`: `none`/`reply`/`generate`, `reply`), `options` (list of `id`, `description`, `reply`, `tool` `{name, arguments}`), `address` (wake-word gate: `names`, `mode`: `ignore`/`reply`, `reply`), `model` (optional separate scoring config). See [Realtime classifier mode]({{%relref "features/openai-realtime#classifier-mode-localai-extension" %}}) | ## gRPC Configuration diff --git a/docs/content/development/offline-tests.md b/docs/content/development/offline-tests.md new file mode 100644 index 000000000000..104a04f4d00d --- /dev/null +++ b/docs/content/development/offline-tests.md @@ -0,0 +1,77 @@ +--- +title: "Offline test resources" +--- + +LocalAI tests separate resource acquisition from test execution. Resources are +declared by resource set in `test-resources/manifests/`; files and packed container +images are content-addressed by SHA-256 under +`.cache/test-resources/blobs/sha256/`. + +Prepare the resources before running a target: + +```sh +make prepare-offline-test-cache TEST_RESOURCE_SET=default +``` + +Preparation verifies every cached blob and fails closed. It never substitutes +a live request for a missing or corrupt entry. Maintainers can populate a +cache from pinned declarations only by explicitly enabling online mode: + +```sh +LOCALAI_TEST_RESOURCES_ONLINE=1 make update-offline-test-cache TEST_RESOURCE_SET=default +``` + +The update command records declared responses, files, and digest-pinned images, +then writes a deterministic, zstd level-1 bundle at +`.cache/test-resources/bundles/.tar.zst`. Its SHA-256 is written to the +lock file. The test workflow transfers the bundle as a workflow artifact and +verifies it after deleting the recording cache; the scheduled refresh workflow +also publishes verified bundles to GHCR as OCI artifacts. + +HTTP declarations may include `request_headers`. `Range` participates in the +cache key, and authorization values participate only through a SHA-256 value; +credentials are never written verbatim to the cache index. Redirect responses +are recorded without following them, so every hop needed by a test must be +declared explicitly. + +File and HTTP declarations may list HTTPS `mirrors`. Recording tries the +canonical URL twice, then each mirror twice, and reports the duration of every +attempt. Every candidate must produce the same declared SHA-256; mirrors are +alternate transports, not alternate content. + +A digest mismatch is never accepted automatically. The updater prints the +observed failure for every source and directs maintainers to compare upstream +checksums, signatures, release notes, and redirects, then check the GitHub +Advisory Database and OSV before approving a new digest. Repeated mismatches can +mean a legitimate upstream release, a corrupt mirror, or a supply-chain event. + +Ordinary test recipes execute through `scripts/run-test-offline.sh`. Its +supervised replay proxy terminates HTTP and HTTPS and returns an immediate +error containing the method and URL for undeclared requests. Linux CI also +runs the command in a cgroup with public IPv4 and IPv6 rejected; macOS relies +on replay, declared resources, guarded Go transports, and static lint because +kernel-level subprocess enforcement is Linux-only. + +Testcontainer images must be registry-digest pinned and loaded during +preparation. Container helpers check that an image exists before startup and +attach services to internal-only Docker networks, preventing testcontainers +from silently pulling a missing tag. + +The default Linux and macOS suites use separate resource sets because +Docker archives are platform-specific. Backend and hardware resources remain +separate targets so ordinary contributors do not acquire large model fixtures +that their test command does not use. + +Coverage runs print a wall-clock summary for each test root and list every +Ginkgo spec or hook taking at least three seconds, including its source +location. Set `COVERAGE_SLOW_SPEC_THRESHOLD=` to tune the reporting +threshold. This measures the whole spec or hook, so it exposes time spent in +sleeps, polling, channel waits, cleanup, and resource contention without +replacing Go's global clock or changing test semantics. The same timings are +written to `coverage/timings.tsv` for CI artifacts and comparisons. The report +shows the slowest 25 entries per root by default; set +`COVERAGE_SLOW_SPEC_LIMIT=` to change the cap. + +Real third-party compatibility checks belong in separately named +`external-probe-*` scheduled workflows and must not be part of deterministic +test or coverage gates. diff --git a/docs/content/features/openai-realtime.md b/docs/content/features/openai-realtime.md index 9eb913a6b7e4..959451c1177b 100644 --- a/docs/content/features/openai-realtime.md +++ b/docs/content/features/openai-realtime.md @@ -111,10 +111,13 @@ pipeline: type: semantic_vad # default for sessions on this model (server_vad if unset) eagerness: medium # low | medium | high | auto (auto == medium) retranscribe: false # see below + # vad_window_sec: 6 # widen the per-tick VAD scan window (see below) ``` A client `session.update` still overrides `type` and `eagerness` per session. +**VAD scan window**: each turn-detection tick the VAD rescans only the most recent slice of buffered audio — sized automatically from the commit silence threshold (the `server_vad` silence window, or the semantic eagerness fallback) plus a warm-up margin, so long turns cost the same per tick as short ones. `vad_window_sec` widens the window if needed; values below the automatic floor are ignored. Buffered turn audio is retained for at most 90 s — a turn that genuinely never pauses (continuous speech or a noise source the VAD keeps classifying as speech) keeps only its most recent 90 s for the commit-time batch transcription (the semantic live stream is unaffected — it already consumed the audio incrementally). + **Eagerness** sets the fallback silence window used when no end-of-utterance token was seen (the model missed it, or the user genuinely trails off): `low` waits 8 s, `medium`/`auto` 4 s, `high` 2 s - the same max-timeout semantics OpenAI documents. After the token is seen, the turn commits on the next VAD tick (~300 ms). **Live captions**: while the user speaks, `semantic_vad` streams `conversation.item.input_audio_transcription.delta` events under the item id the commit will later reuse, so clients can render the words as they are recognized. The `completed` event at commit carries the authoritative transcript and replaces the partial text (with `retranscribe: true` it may differ from the captions); a turn discarded before commit emits `conversation.item.input_audio_transcription.failed` so clients can retract its captions. @@ -163,6 +166,90 @@ On CPU, set `summary_model` to a small, fast model so compaction never competes Clients can also manage history directly via the now-supported `conversation.item.delete`, `conversation.item.truncate`, and `input_audio_buffer.clear` realtime events. +### Classifier mode (LocalAI extension) + +On hardware that can afford prompt processing but not token generation — a Raspberry Pi running a small LLM, for example — a realtime session can replace autoregressive generation with **prefill-only classification**: you register a fixed list of options, each user turn is scored against them with the Score primitive (a single forward pass, no decode), and the winning option's canned reply is spoken and/or its canned tool call is emitted. On llama-cpp, scoring runs through the same server slot the LLM uses, so the conversation prefix stays KV-cached across turns, and all options are scored together in one batched decode: the shared prefix (prompt plus the options' common token prefix) is processed once, then each option's unique tail rides a forked sequence in a single forward pass. A warm turn costs roughly one pass over the new words plus one small batch over the option tails, independent of the option count. + +Enable it in the pipeline config: + +```yaml +name: drone-pi +pipeline: + vad: silero-vad-sherpa + transcription: parakeet-cpp-realtime_eou_120m-v1 + llm: lfm2.5-1.2b-instruct # scores AND (if asked) generates + tts: vits-piper-en_US-amy-sherpa + classifier: + enabled: true + threshold: 0.85 # softmax floor the winner must clear (see note below) + fallback: + mode: reply # none | reply | generate + reply: "Say again?" + options: + - id: up + description: the user asks the drone to move or fly up/higher + reply: Going up. + tool: + name: move + arguments: {direction: up} + - id: greeting + description: the user greets the assistant + reply: Hello, ready to fly. +``` + +Or per session / per response from the client (the field is additive — OpenAI clients simply never send it): + +```json +{"type": "session.update", "session": {"type": "realtime", "localai_classifier": { + "enabled": true, "threshold": 0.85, + "options": [{"id": "up", "description": "...", "reply": "Going up.", + "tool": {"name": "move", "arguments": {"direction": "up"}}}], + "fallback": {"mode": "reply", "reply": "Say again?"} +}}} +``` + +Like `tools`, the option list is replaced wholesale on each update. A `response.create` may carry its own `localai_classifier` to override the session for one response — `{"enabled": false}` runs normal generation once. + +How a classified response behaves: + +- The winner's `reply` is emitted through the ordinary response events (spoken via TTS, or `response.output_text.*` in text-only mode), and its `tool` (if any) is emitted as a standard `function_call` item with exactly the configured arguments — the client executes it and reports back with `conversation.item.create` as usual. In classifier mode you typically should **not** send a follow-up `response.create` after the tool output: the canned reply already acknowledged the command, and the follow-up would classify a tool-output turn. +- Every classified response also emits a `localai.classifier.result` event carrying the full softmax distribution, the chosen option id (empty when the fallback applied), the threshold, and the scoring latency — useful for visualizing confidence in a client UI. +- A committed turn whose transcript is empty (the VAD fired on noise and the ASR heard no words) is never scored — an empty prompt produces a confidently arbitrary winner. The fallback applies directly: the result event carries an empty `scores` list, and `generate` mode falls through to generation. +- **Wake-word gating**: set `address: {names: ["drone"], mode: ignore}` and the assistant only acts on turns that mention one of the names ("Drone go up", not just "go up"). The check is a deterministic case-insensitive whole-word match on the latest transcript — deliberately not model-based: scoring cannot detect a missing name (a 1.2B scorer rates "go up" as addressed with p≈1.0 even with a dedicated addressing stage), while a literal match is exact and free. Unaddressed turns skip scoring entirely (ambient conversation costs nothing) and emit a result event with empty `scores` and `fallback: "not_addressed"`; `mode: ignore` completes the response silently, `mode: reply` speaks `reply`. +- When no option clears `threshold`, `fallback.mode` decides: `none` completes the response with no output, `reply` speaks the canned fallback reply, and `generate` falls through to normal autoregressive generation (slow on weak hardware, but always available since the same model config serves both paths). Set the threshold high: with the default `raw` normalization a confident in-list pick lands near 1.0, while an out-of-list request spreads its probability across the options — measured on a 1.2B scorer, in-list utterances scored ≥0.97 and out-of-list ones peaked around 0.8, so a floor of ~0.85 separates them. A low threshold (say 0.35) practically never falls back. Also keep each option's `description` narrowly scoped: a catch-all clause like "…or asks for help" turns that option into a magnet for every request the model cannot map, defeating the fallback. +- Agentic follow-up turns (after a server-side assistant tool executes) always use generation — the option list describes user intents, not tool outputs. + +Knobs that matter for latency and accuracy: keep option `description`s short (they all go into the scoring system prompt) and the option count small. By default only the latest user message is scored — earlier turns echo option names (the canned replies especially) and empirically make small scoring models re-choose the previous option regardless of the new command. `history_items: N` opts the trailing N conversation messages back in (role-labeled); only do that with a scorer large enough to weigh the context. `normalization: mean` divides each option's joint log-prob by its token count — useful when option ids have very different lengths. The scoring model needs a Go-side chat template (`template.chat` / `template.chat_message`); without one the scoring prompt falls back to a generic ChatML envelope, which may be off-distribution for the model. Use `classifier.model` to score on a different config than the pipeline LLM (rarely needed). + +### Argument slots (hybrid classify-then-complete) + +A canned tool call can leave holes for the model to fill: declare `slots` on the option's tool and reference them as `"{{name}}"` in the arguments template. When the option wins, LocalAI runs a short **grammar-constrained completion** that continues the exact scoring prompt (so the llama.cpp prompt cache is already warm) with the chosen route JSON re-opened at the first slot — only the value tokens are free; everything else is pinned by the grammar. Number slots substitute unquoted, enum/string slots inside their quotes. + +```yaml +tool: + name: move_drone + arguments: + direction: forward + distance: "{{distance}}" + units: "{{units}}" + slots: + - name: distance + type: number # number | enum | string + - name: units + type: enum + values: [m, meters, ft, feet] + default: m # used if inference fails outright + hint: assume m when the user gives no units +``` + +The option's spoken `reply` can reference the same placeholders — `reply: "Going forward {{distance}} {{units}}."` — and the filled values are spliced in as plain text before the reply is emitted, so what the assistant says confirms what it inferred. Reply placeholders are optional (one that names no slot stays literal). + +Slot declarations (and hints) are appended to the option's description in the shared system prompt, so they also inform scoring and cost no extra per-turn tokens. The `localai.classifier.result` event carries the final `arguments` and a `fill_latency_ms`. On an inference failure the slots' defaults apply (and template the reply); if any slot lacks a default the response fails (or falls through to generation with `fallback.mode: generate`). Slot filling requires `completion` in the scoring model's `known_usecases` alongside `score`. + +Registering an option list (pipeline seed or `session.update`) prewarms the scoring prompt in the background: one throwaway score prefills the option-list prompt and plants a state checkpoint exactly at the per-turn probe boundary. Every scoring call also declares that boundary (the probe-invariant prompt prefix) to the backend, which checkpoints there on each prefill — on hybrid/recurrent models (which cannot rewind their state arbitrarily, only restore checkpoints) this is what keeps every turn at probe-size cost instead of a full option-list re-prefill. A client that swaps option lists at runtime (e.g. voice-switched command modes) pays nothing on the first turn after a swap: the rewarm hides behind the acknowledgement reply, and it runs on every registration deliberately — if the list's slot was evicted, the rewarm is exactly the re-prefill the next turn would otherwise pay in the foreground. Swapping between several lists? Give the scoring model a slot per list and make the prefix routing selective: `options: [parallel:3, sps:0.5]` (the default slot-similarity threshold of 0.1 funnels different lists onto one slot — they share enough prompt structure to clear it). + +The concrete scoring model must declare `score` in `known_usecases`. A single llama.cpp model can serve ordinary inference and classification concurrently by declaring multiple use cases, for example `known_usecases: [chat, completion, score]`; LocalAI reserves the scoring slots only when `score` is present. Scoring also requires the unified KV cache, which is enabled by default, so a score-enabled model cannot set `kv_unified:false`. + ## Transports The Realtime API supports two transports: **WebSocket** and **WebRTC**. diff --git a/docs/content/features/text-generation.md b/docs/content/features/text-generation.md index f7cc155a80f5..ee76958ac96f 100644 --- a/docs/content/features/text-generation.md +++ b/docs/content/features/text-generation.md @@ -531,7 +531,7 @@ The `llama.cpp` backend supports additional configuration options that can be sp | `warmup` | boolean | Enable warmup run after model loading. Default: `true`. | `warmup:false` | | `no_op_offload` | boolean | Disable offloading host tensor operations to device. Default: `false`. | `no_op_offload:true` | | `device` or `devices` | string | Select the llama.cpp backend devices to use. Repeat the option or pass a comma-separated list; unlisted devices are excluded. Use the names reported by `llama-server --list-devices` / `--list-devices`. | `devices:CUDA1,CUDA2,CUDA3` | -| `kv_unified` or `unified_kv` | boolean | Use a single unified KV buffer shared across all sequences. Default: `true` (LocalAI override; upstream defaults to `false` but auto-enables it when slot count is auto). **Required for `cache_idle_slots` to work**: without it the server force-disables idle-slot saving at init, and the prompt cache is never written across requests. | `kv_unified:false` | +| `kv_unified` or `unified_kv` | boolean | Use a single unified KV buffer shared across all sequences. Default: `true` (LocalAI override; upstream defaults to `false` but auto-enables it when slot count is auto). **Required for `cache_idle_slots` and scoring**: without it the server force-disables idle-slot saving at init, and score-enabled models are rejected at load time. | `kv_unified:false` | | `cache_idle_slots` or `idle_slots_cache` | boolean | On a new task, save the previous slot's KV state into the prompt cache (and clear the slot) so a later request with the same prefix can warm-load it. Default: `true`. Auto-disabled by the server if `kv_unified=false` or `cache_ram=0`. | `cache_idle_slots:false` | | `n_ctx_checkpoints` or `ctx_checkpoints` | integer | Maximum number of context checkpoints per slot (used for partial-prefix recovery, e.g. SWA). Default: `32`. | `ctx_checkpoints:16` | | `checkpoint_min_step` or `checkpoint_min_spacing` (aliases: `checkpoint_every_nt`, `checkpoint_every_n_tokens`) | integer | Minimum spacing in tokens between context checkpoints. `0` disables the minimum-spacing gate. Default: `256`. (Renamed upstream from `checkpoint_every_nt`; semantics shifted from a fixed cadence to a minimum spacing.) | `checkpoint_min_step:1024` | diff --git a/docs/design/realtime-state-machines.md b/docs/design/realtime-state-machines.md index d88313760b46..700929783d9c 100644 --- a/docs/design/realtime-state-machines.md +++ b/docs/design/realtime-state-machines.md @@ -88,6 +88,11 @@ cancelling → done(completed|cancelled) | failed. - Terminal events are **not exactly-once**: failed paths `return` with no `response.done`; cancelled paths emit `done{Cancelled}`; the completed terminal is unconditional at the tail of `emitToolCallItems`. +- **Classifier mode** (`realtime_classifier.go`) is a response-*body* variant, not a + new machine: at turn 0 it may replace the Predict call with a prefill-only Score + and canned emission, but it runs inside the same respcoord-issued response, maps + onto the existing `outcomeCompleted/Cancelled/Failed`, and leaves terminal + emission with `triggerResponse`. No coordinator states or transitions were added. ### M4. Conversation / compaction diff --git a/gallery/sherpa-onnx-vad.yaml b/gallery/sherpa-onnx-vad.yaml index 72c226e02e00..0ed940c7c901 100644 --- a/gallery/sherpa-onnx-vad.yaml +++ b/gallery/sherpa-onnx-vad.yaml @@ -4,6 +4,11 @@ name: "sherpa-onnx-vad" config_file: | backend: sherpa-onnx type: vad + # Silero is a ~2MB recurrent model with no exploitable graph parallelism: + # measured per-call latency is identical at 1 and 10 ORT threads, while + # every extra pool thread just spin-waits between the realtime loop's + # frequent tiny inferences. + threads: 1 options: # Silero VAD. Defaults mirror upstream sherpa-onnx. Override for # faster turn-taking (lower min_silence) or different sample rate diff --git a/internal/testfixtures/images.go b/internal/testfixtures/images.go new file mode 100644 index 000000000000..fe16d766d836 --- /dev/null +++ b/internal/testfixtures/images.go @@ -0,0 +1,55 @@ +// SPDX-License-Identifier: MIT + +// Package testfixtures centralizes immutable resources shared by test suites. +package testfixtures + +import ( + "context" + "errors" + "fmt" + "net" + "os" + + "github.com/moby/moby/client" + "github.com/testcontainers/testcontainers-go" +) + +const ( + Postgres16 = "docker.io/library/postgres@sha256:33f923b05f64ca54ac4401c01126a6b92afe839a0aa0a52bc5aeb5cc958e5f20" + Postgres16Alpine = "docker.io/library/postgres@sha256:57c72fd2a128e416c7fcc499958864df5301e940bca0a56f58fddf30ffc07777" + NATS2Alpine = "docker.io/library/nats@sha256:c11af972c99ae542de8925e6a7d9c533aa1eb039660420d2074beed6089b3bf0" +) + +// RequireImage fails before testcontainers can fall back to a registry pull. +func RequireImage(ctx context.Context, reference, target string) error { + docker, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation()) + if err != nil { + return err + } + defer func() { _ = docker.Close() }() + if _, err := docker.ImageInspect(ctx, reference); err != nil { + return fmt.Errorf("required offline test image %s is not loaded; run `make prepare-offline-test-cache TEST_RESOURCE_SET=%s`: %w", reference, target, err) + } + return nil +} + +func DockerNetwork() (string, error) { + name := os.Getenv("LOCALAI_TEST_DOCKER_NETWORK") //nolint:forbidigo + if name == "" { + return "", errors.New("offline test Docker network is not configured; run the test through scripts/run-test-offline.sh") + } + return name, nil +} + +// ContainerEndpoint returns an address reachable from the Linux test host +// without publishing a port from the internal-only Docker network. +func ContainerEndpoint(ctx context.Context, container testcontainers.Container, port string) (string, error) { + ip, err := container.ContainerIP(ctx) + if err != nil { + return "", err + } + if ip == "" { + return "", errors.New("offline test container has no private network address") + } + return net.JoinHostPort(ip, port), nil +} diff --git a/internal/testresources/bundle.go b/internal/testresources/bundle.go new file mode 100644 index 000000000000..eeffa3f1fa84 --- /dev/null +++ b/internal/testresources/bundle.go @@ -0,0 +1,164 @@ +// SPDX-License-Identifier: MIT + +package testresources + +import ( + "archive/tar" + "bytes" + "crypto/sha256" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "sort" + "strings" + "time" + + "github.com/klauspost/compress/zstd" +) + +func PackBundle(cacheDir, output string, manifest Manifest) (string, error) { + index, err := LoadHTTPIndex(cacheDir) + if err != nil { + return "", err + } + targetIndex := map[string]HTTPEntry{} + digests := map[string]bool{} + for _, resource := range manifest.HTTP { + key := RequestKey(resource.Method, resource.URL, resource.Headers()) + entry, ok := index[key] + if !ok { + return "", fmt.Errorf("cannot pack missing HTTP entry %s", key) + } + targetIndex[key], digests[resource.SHA256] = entry, true + } + for _, resource := range manifest.Files { + digests[resource.SHA256] = true + } + for _, resource := range manifest.Images { + digests[resource.SHA256] = true + } + if err := os.MkdirAll(filepath.Dir(output), 0o755); err != nil { + return "", err + } + tmp, err := os.CreateTemp(filepath.Dir(output), "bundle-*.tmp") + if err != nil { + return "", err + } + name := tmp.Name() + defer func() { _ = os.Remove(name) }() + hash := sha256.New() + zstdWriter, err := zstd.NewWriter(io.MultiWriter(tmp, hash), + zstd.WithEncoderLevel(zstd.SpeedFastest), + zstd.WithEncoderConcurrency(1), + zstd.WithEncoderCRC(true), + ) + if err != nil { + _ = tmp.Close() + return "", err + } + tw := tar.NewWriter(zstdWriter) + indexData, err := json.Marshal(targetIndex) + if err == nil { + err = writeTarBytes(tw, "http-index.json", indexData) + } + ordered := make([]string, 0, len(digests)) + for digest := range digests { + ordered = append(ordered, digest) + } + sort.Strings(ordered) + for _, digest := range ordered { + if err != nil { + break + } + path, verifyErr := VerifyBlob(cacheDir, digest) + if verifyErr != nil { + err = verifyErr + break + } + var data []byte + data, err = os.ReadFile(path) + if err == nil { + err = writeTarBytes(tw, filepath.ToSlash(filepath.Join("blobs", "sha256", digest)), data) + } + } + err = errors.Join(err, tw.Close(), zstdWriter.Close(), tmp.Close()) + if err != nil { + return "", err + } + if err := os.Rename(name, output); err != nil { + return "", err + } + return fmt.Sprintf("%x", hash.Sum(nil)), nil +} + +func RestoreBundle(cacheDir, bundle, expected string) error { + data, err := os.ReadFile(bundle) + if err != nil { + return err + } + actual := fmt.Sprintf("%x", sha256.Sum256(data)) + if actual != expected { + return fmt.Errorf("test resource bundle checksum mismatch: expected %s, got %s", expected, actual) + } + var bundleReader io.Reader = bytes.NewReader(data) + if len(data) >= 4 && bytes.Equal(data[:4], []byte{0x28, 0xb5, 0x2f, 0xfd}) { + zstdReader, err := zstd.NewReader(bundleReader, zstd.WithDecoderConcurrency(1)) + if err != nil { + return err + } + defer zstdReader.Close() + bundleReader = zstdReader + } + tr := tar.NewReader(bundleReader) + recorded := map[string]HTTPEntry{} + for { + header, err := tr.Next() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + return err + } + name := filepath.Clean(filepath.FromSlash(header.Name)) + if filepath.IsAbs(name) || name == ".." || strings.HasPrefix(name, ".."+string(filepath.Separator)) { + return fmt.Errorf("unsafe bundle path %q", header.Name) + } + body, err := io.ReadAll(tr) + if err != nil { + return err + } + if name == "http-index.json" { + if err := json.Unmarshal(body, &recorded); err != nil { + return err + } + continue + } + destination := filepath.Join(cacheDir, name) + if err := os.MkdirAll(filepath.Dir(destination), 0o755); err != nil { + return err + } + if err := os.WriteFile(destination, body, 0o644); err != nil { + return err + } + } + index, err := LoadHTTPIndex(cacheDir) + if err != nil { + return err + } + for key, entry := range recorded { + index[key] = entry + } + return WriteHTTPIndex(cacheDir, index) +} + +func writeTarBytes(tw *tar.Writer, name string, data []byte) error { + header := &tar.Header{Name: name, Mode: 0o644, Size: int64(len(data)), ModTime: time.Unix(0, 0).UTC()} + if err := tw.WriteHeader(header); err != nil { + return err + } + _, err := tw.Write(data) + return err +} diff --git a/internal/testresources/httpcache.go b/internal/testresources/httpcache.go new file mode 100644 index 000000000000..c32f515995af --- /dev/null +++ b/internal/testresources/httpcache.go @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: MIT + +package testresources + +import ( + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "os" + "path/filepath" +) + +var hopHeaders = map[string]bool{ + "Connection": true, "Proxy-Connection": true, "Keep-Alive": true, + "Transfer-Encoding": true, "Content-Length": true, "Te": true, + "Trailer": true, "Upgrade": true, "Proxy-Authenticate": true, + "Proxy-Authorization": true, +} + +func LoadHTTPIndex(cacheDir string) (map[string]HTTPEntry, error) { + index := map[string]HTTPEntry{} + data, err := os.ReadFile(filepath.Join(cacheDir, "index.json")) + if errors.Is(err, os.ErrNotExist) { + return index, nil + } + if err != nil { + return nil, fmt.Errorf("read HTTP cache index: %w", err) + } + if err := json.Unmarshal(data, &index); err != nil { + return nil, fmt.Errorf("parse HTTP cache index: %w", err) + } + return index, nil +} + +func WriteHTTPIndex(cacheDir string, index map[string]HTTPEntry) error { + data, err := json.MarshalIndent(index, "", " ") + if err != nil { + return err + } + data = append(data, '\n') + if err := os.MkdirAll(cacheDir, 0o755); err != nil { + return err + } + tmp, err := os.CreateTemp(cacheDir, "index-*.tmp") + if err != nil { + return err + } + name := tmp.Name() + defer func() { _ = os.Remove(name) }() + if _, err := tmp.Write(data); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Close(); err != nil { + return err + } + return os.Rename(name, filepath.Join(cacheDir, "index.json")) +} + +func SanitizeHeaders(header http.Header) http.Header { + out := header.Clone() + for name := range hopHeaders { + out.Del(name) + } + return out +} + +func ReplayResponse(w http.ResponseWriter, cacheDir string, entry HTTPEntry) error { + path, err := VerifyBlob(cacheDir, entry.Digest) + if err != nil { + return err + } + for name, values := range entry.Header { + for _, value := range values { + w.Header().Add(name, value) + } + } + w.Header().Set("Content-Length", fmt.Sprint(entry.Size)) + w.WriteHeader(entry.Status) + if entry.Size == 0 { + return nil + } + body, err := os.Open(path) + if err != nil { + return err + } + defer func() { _ = body.Close() }() + _, err = io.Copy(w, body) + return err +} diff --git a/internal/testresources/resources.go b/internal/testresources/resources.go new file mode 100644 index 000000000000..a5c5afd8ad64 --- /dev/null +++ b/internal/testresources/resources.go @@ -0,0 +1,207 @@ +// SPDX-License-Identifier: MIT + +package testresources + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" +) + +const ManifestVersion = 1 + +type Manifest struct { + Version int `json:"version"` + Target string `json:"target"` + HTTP []HTTP `json:"http,omitempty"` + Files []File `json:"files,omitempty"` + Images []OCIImage `json:"images,omitempty"` +} + +type HTTP struct { + Method string `json:"method"` + URL string `json:"url"` + Mirrors []string `json:"mirrors,omitempty"` + SHA256 string `json:"sha256"` + RequestHeaders map[string]string `json:"request_headers,omitempty"` +} + +type HTTPEntry struct { + Digest string `json:"digest"` + Size int64 `json:"size"` + Status int `json:"status"` + Header http.Header `json:"header"` +} + +type File struct { + URL string `json:"url"` + Mirrors []string `json:"mirrors,omitempty"` + SHA256 string `json:"sha256"` + Destination string `json:"destination,omitempty"` + Environment string `json:"environment,omitempty"` +} + +type OCIImage struct { + Reference string `json:"reference"` + SHA256 string `json:"sha256"` +} + +type Lock struct { + Version int `json:"version"` + Bundles map[string]string `json:"bundles"` +} + +func LoadManifest(path string) (Manifest, error) { + var manifest Manifest + if err := decode(path, &manifest); err != nil { + return manifest, err + } + if err := manifest.Validate(); err != nil { + return manifest, fmt.Errorf("%s: %w", path, err) + } + return manifest, nil +} + +func LoadLock(path string) (Lock, error) { + var lock Lock + if err := decode(path, &lock); err != nil { + return lock, err + } + if lock.Version != ManifestVersion { + return lock, fmt.Errorf("%s: unsupported version %d", path, lock.Version) + } + return lock, nil +} + +func WriteLock(path string, lock Lock) error { + data, err := json.MarshalIndent(lock, "", " ") + if err != nil { + return err + } + data = append(data, '\n') + return os.WriteFile(path, data, 0o644) +} + +func decode(path string, value any) error { + data, err := os.ReadFile(path) + if err != nil { + return err + } + decoder := json.NewDecoder(strings.NewReader(string(data))) + decoder.DisallowUnknownFields() + if err := decoder.Decode(value); err != nil { + return err + } + if decoder.Decode(&struct{}{}) != io.EOF { + return errors.New("manifest has trailing JSON data") + } + return nil +} + +func (m Manifest) Validate() error { + if m.Version != ManifestVersion { + return fmt.Errorf("unsupported version %d", m.Version) + } + if strings.TrimSpace(m.Target) == "" { + return errors.New("target is required") + } + for _, resource := range m.HTTP { + if resource.Method == "" || resource.URL == "" || !validDigest(resource.SHA256) { + return fmt.Errorf("HTTP resources require method, URL, and lowercase sha256: %s %s", resource.Method, resource.URL) + } + if err := validateMirrors(resource.Mirrors); err != nil { + return fmt.Errorf("HTTP resource %s: %w", resource.URL, err) + } + } + for _, resource := range m.Files { + if resource.URL == "" || !validDigest(resource.SHA256) || (resource.Destination == "" && resource.Environment == "") { + return fmt.Errorf("file resources require URL, sha256, and destination or environment: %s", resource.URL) + } + if filepath.IsAbs(resource.Destination) || strings.HasPrefix(filepath.Clean(resource.Destination), "..") { + return fmt.Errorf("file destination must stay inside the resource directory: %s", resource.Destination) + } + if err := validateMirrors(resource.Mirrors); err != nil { + return fmt.Errorf("file resource %s: %w", resource.URL, err) + } + } + for _, resource := range m.Images { + if !strings.Contains(resource.Reference, "@sha256:") || !validDigest(resource.SHA256) { + return fmt.Errorf("OCI image must be digest-pinned and have a packed sha256: %s", resource.Reference) + } + } + return nil +} + +func validateMirrors(mirrors []string) error { + seen := map[string]bool{} + for _, mirror := range mirrors { + if !strings.HasPrefix(mirror, "https://") { + return fmt.Errorf("mirror must use HTTPS: %s", mirror) + } + if seen[mirror] { + return fmt.Errorf("duplicate mirror: %s", mirror) + } + seen[mirror] = true + } + return nil +} + +func BlobPath(cacheDir, digest string) string { + return filepath.Join(cacheDir, "blobs", "sha256", digest) +} + +func RequestKey(method, rawURL string, headers ...http.Header) string { + key := strings.ToUpper(method) + " " + rawURL + if len(headers) == 0 { + return key + } + for _, name := range []string{"Authorization", "Range"} { + value := headers[0].Get(name) + if value == "" { + continue + } + if name == "Authorization" { + digest := sha256.Sum256([]byte(value)) + value = "sha256:" + hex.EncodeToString(digest[:]) + } + key += "\n" + strings.ToLower(name) + ":" + value + } + return key +} + +func (resource HTTP) Headers() http.Header { + header := make(http.Header, len(resource.RequestHeaders)) + for name, value := range resource.RequestHeaders { + header.Set(name, value) + } + return header +} + +func VerifyBlob(cacheDir, digest string) (string, error) { + path := BlobPath(cacheDir, digest) + data, err := os.ReadFile(path) + if err != nil { + return "", fmt.Errorf("missing CAS blob %s: %w", digest, err) + } + sum := sha256.Sum256(data) + actual := hex.EncodeToString(sum[:]) + if actual != digest { + return "", fmt.Errorf("corrupt CAS blob %s: got sha256:%s", digest, actual) + } + return path, nil +} + +func validDigest(value string) bool { + if len(value) != sha256.Size*2 || strings.ToLower(value) != value { + return false + } + _, err := hex.DecodeString(value) + return err == nil +} diff --git a/internal/testresources/resources_suite_test.go b/internal/testresources/resources_suite_test.go new file mode 100644 index 000000000000..ead665b25526 --- /dev/null +++ b/internal/testresources/resources_suite_test.go @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: MIT + +package testresources_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestResources(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Test resources suite") +} diff --git a/internal/testresources/resources_test.go b/internal/testresources/resources_test.go new file mode 100644 index 000000000000..ae6203b5edf0 --- /dev/null +++ b/internal/testresources/resources_test.go @@ -0,0 +1,126 @@ +// SPDX-License-Identifier: MIT + +package testresources_test + +import ( + "crypto/sha256" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/internal/testresources" +) + +var _ = Describe("Declared test resources", func() { + It("rejects mutable and unpinned resources", func() { + manifest := testresources.Manifest{ + Version: testresources.ManifestVersion, + Target: "backend", + Images: []testresources.OCIImage{{Reference: "postgres:latest", SHA256: fmt.Sprintf("%064d", 0)}}, + } + Expect(manifest.Validate()).To(MatchError(ContainSubstring("digest-pinned"))) + }) + + It("requires HTTPS and unique mirrors", func() { + digest := fmt.Sprintf("%064d", 0) + manifest := testresources.Manifest{Version: 1, Target: "fixture", Files: []testresources.File{{ + URL: "https://primary.invalid/file", Mirrors: []string{"http://mirror.invalid/file"}, + SHA256: digest, Destination: "file", + }}} + Expect(manifest.Validate()).To(MatchError(ContainSubstring("mirror must use HTTPS"))) + manifest.Files[0].Mirrors = []string{"https://mirror.invalid/file", "https://mirror.invalid/file"} + Expect(manifest.Validate()).To(MatchError(ContainSubstring("duplicate mirror"))) + }) + + It("fails before tests when a CAS blob is missing or corrupt", func() { + cache := GinkgoT().TempDir() + digest := fmt.Sprintf("%064d", 0) + _, err := testresources.VerifyBlob(cache, digest) + Expect(err).To(MatchError(ContainSubstring("missing CAS blob"))) + + path := testresources.BlobPath(cache, digest) + Expect(os.MkdirAll(filepath.Dir(path), 0o755)).To(Succeed()) + Expect(os.WriteFile(path, []byte("corrupt"), 0o644)).To(Succeed()) + _, err = testresources.VerifyBlob(cache, digest) + Expect(err).To(MatchError(ContainSubstring("corrupt CAS blob"))) + }) + + It("accepts and verifies a content-addressed blob", func() { + cache := GinkgoT().TempDir() + content := []byte("offline fixture") + digest := fmt.Sprintf("%x", sha256.Sum256(content)) + path := testresources.BlobPath(cache, digest) + Expect(os.MkdirAll(filepath.Dir(path), 0o755)).To(Succeed()) + Expect(os.WriteFile(path, content, 0o644)).To(Succeed()) + Expect(testresources.VerifyBlob(cache, digest)).To(Equal(path)) + }) + + It("persists response metadata and replays a verified body", func() { + cache := GinkgoT().TempDir() + content := []byte("cached response") + digest := fmt.Sprintf("%x", sha256.Sum256(content)) + path := testresources.BlobPath(cache, digest) + Expect(os.MkdirAll(filepath.Dir(path), 0o755)).To(Succeed()) + Expect(os.WriteFile(path, content, 0o644)).To(Succeed()) + index := map[string]testresources.HTTPEntry{ + "GET https://example.invalid/data": { + Digest: digest, Size: int64(len(content)), Status: http.StatusPartialContent, + Header: http.Header{"Content-Range": {"bytes 0-14/15"}}, + }, + } + Expect(testresources.WriteHTTPIndex(cache, index)).To(Succeed()) + loaded, err := testresources.LoadHTTPIndex(cache) + Expect(err).NotTo(HaveOccurred()) + recorder := httptest.NewRecorder() + Expect(testresources.ReplayResponse(recorder, cache, loaded["GET https://example.invalid/data"])).To(Succeed()) + Expect(recorder.Code).To(Equal(http.StatusPartialContent)) + Expect(recorder.Body.Bytes()).To(Equal(content)) + Expect(recorder.Header().Get("Content-Range")).To(Equal("bytes 0-14/15")) + }) + + It("sanitizes connection-specific response headers", func() { + header := http.Header{"Transfer-Encoding": {"chunked"}, "Authorization": {"secret"}, "X-Fixture": {"yes"}} + clean := testresources.SanitizeHeaders(header) + Expect(clean).NotTo(HaveKey("Transfer-Encoding")) + Expect(clean).To(HaveKeyWithValue("Authorization", []string{"secret"})) + Expect(clean).To(HaveKeyWithValue("X-Fixture", []string{"yes"})) + }) + + It("keys range and authorization variants without storing credentials", func() { + header := http.Header{"Authorization": {"Bearer secret"}, "Range": {"bytes=4-"}} + key := testresources.RequestKey(http.MethodGet, "https://example.invalid/model", header) + Expect(key).To(ContainSubstring("range:bytes=4-")) + Expect(key).To(ContainSubstring("authorization:sha256:")) + Expect(key).NotTo(ContainSubstring("Bearer secret")) + Expect(key).NotTo(Equal(testresources.RequestKey(http.MethodGet, "https://example.invalid/model"))) + }) + + It("packs deterministically and restores a target cache", func() { + cache := GinkgoT().TempDir() + content := []byte("bundle fixture") + digest := fmt.Sprintf("%x", sha256.Sum256(content)) + path := testresources.BlobPath(cache, digest) + Expect(os.MkdirAll(filepath.Dir(path), 0o755)).To(Succeed()) + Expect(os.WriteFile(path, content, 0o644)).To(Succeed()) + manifest := testresources.Manifest{Version: 1, Target: "fixture", Files: []testresources.File{{URL: "https://example.invalid/file", SHA256: digest, Destination: "file"}}} + first := filepath.Join(GinkgoT().TempDir(), "first.tar.zst") + second := filepath.Join(GinkgoT().TempDir(), "second.tar.zst") + firstDigest, err := testresources.PackBundle(cache, first, manifest) + Expect(err).NotTo(HaveOccurred()) + secondDigest, err := testresources.PackBundle(cache, second, manifest) + Expect(err).NotTo(HaveOccurred()) + Expect(secondDigest).To(Equal(firstDigest)) + compressed, err := os.ReadFile(first) + Expect(err).NotTo(HaveOccurred()) + Expect(compressed[:4]).To(Equal([]byte{0x28, 0xb5, 0x2f, 0xfd})) + + restored := GinkgoT().TempDir() + Expect(testresources.RestoreBundle(restored, first, firstDigest)).To(Succeed()) + Expect(os.ReadFile(testresources.BlobPath(restored, digest))).To(Equal(content)) + }) +}) diff --git a/pkg/downloader/cancel_test.go b/pkg/downloader/cancel_test.go index 76f8a2df5fb0..57f9bba95baf 100644 --- a/pkg/downloader/cancel_test.go +++ b/pkg/downloader/cancel_test.go @@ -59,9 +59,7 @@ var _ = Describe("Download cancellation", func() { } BeforeEach(func() { - dir, err := os.Getwd() - Expect(err).ToNot(HaveOccurred()) - filePath = dir + "/cancel_model" + filePath = GinkgoT().TempDir() + "/cancel_model" }) AfterEach(func() { @@ -112,7 +110,7 @@ var _ = Describe("Download cancellation", func() { Expect(err).To(HaveOccurred()) Expect(errors.Is(err, context.Canceled)).To(BeTrue()) - Expect(filePath + ".partial").ToNot(BeAnExistingFile(), + Expect(filePath+".partial").ToNot(BeAnExistingFile(), "a deliberate user cancel must not leave a dangling .partial behind") }) diff --git a/pkg/downloader/stall_test.go b/pkg/downloader/stall_test.go index 8e6a003c69e8..c2e096131c71 100644 --- a/pkg/downloader/stall_test.go +++ b/pkg/downloader/stall_test.go @@ -17,9 +17,7 @@ var _ = Describe("Download stall timeout", func() { var savedTimeout time.Duration BeforeEach(func() { - dir, err := os.Getwd() - Expect(err).ToNot(HaveOccurred()) - filePath = dir + "/stall_model" + filePath = GinkgoT().TempDir() + "/stall_model" savedTimeout = DownloadStallTimeout }) diff --git a/pkg/downloader/uri_test.go b/pkg/downloader/uri_test.go index 9cb667b57864..3dc961065cee 100644 --- a/pkg/downloader/uri_test.go +++ b/pkg/downloader/uri_test.go @@ -22,31 +22,15 @@ var _ = Describe("Gallery API tests", func() { Context("URI", func() { It("parses github with a branch", func() { uri := URI("github:go-skynet/model-gallery/gpt4all-j.yaml") - Expect( - uri.ReadWithCallback("", func(url string, i []byte) error { - Expect(url).To(Equal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml")) - return nil - }), - ).ToNot(HaveOccurred()) + Expect(uri.ResolveURL()).To(Equal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml")) // test-network: fixture }) It("parses github without a branch", func() { uri := URI("github:go-skynet/model-gallery/gpt4all-j.yaml@main") - - Expect( - uri.ReadWithCallback("", func(url string, i []byte) error { - Expect(url).To(Equal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml")) - return nil - }), - ).ToNot(HaveOccurred()) + Expect(uri.ResolveURL()).To(Equal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml")) // test-network: fixture }) It("parses github with urls", func() { uri := URI("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml") - Expect( - uri.ReadWithCallback("", func(url string, i []byte) error { - Expect(url).To(Equal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml")) - return nil - }), - ).ToNot(HaveOccurred()) + Expect(uri.ResolveURL()).To(Equal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml")) // test-network: fixture }) }) @@ -263,9 +247,7 @@ var _ = Describe("Download Test", func() { _, err = _mockDataSha.Write(mockData) Expect(err).ToNot(HaveOccurred()) mockDataSha = fmt.Sprintf("%x", _mockDataSha.Sum(nil)) - dir, err := os.Getwd() - filePath = dir + "/my_supercool_model" - Expect(err).NotTo(HaveOccurred()) + filePath = GinkgoT().TempDir() + "/my_supercool_model" }) Context("URI DownloadFile", func() { diff --git a/pkg/httpclient/client.go b/pkg/httpclient/client.go index 99a25c8121d2..2172f36fd402 100644 --- a/pkg/httpclient/client.go +++ b/pkg/httpclient/client.go @@ -28,8 +28,11 @@ import ( "net" "net/http" "net/url" + "os" "strings" "time" + + "github.com/mudler/LocalAI/pkg/testnetwork" ) const ( @@ -105,12 +108,20 @@ func sameOrigin(a, b *url.URL) bool { // (e.g. a credential-injecting RoundTripper) should base it on this rather than // http.DefaultTransport so the TLS floor and timeouts are preserved. func HardenedTransport() *http.Transport { + dialContext := (&net.Dialer{ + Timeout: dialTimeout, + KeepAlive: dialKeepAlive, + }).DialContext + // This is set only by the test-resource supervisor before it starts the + // child process; production configuration does not cross this boundary. + if os.Getenv("LOCALAI_TEST_OFFLINE") == "1" { //nolint:forbidigo + guard := testnetwork.LocalGuard() + guard.Dial = dialContext + dialContext = guard.DialContext + } return &http.Transport{ - Proxy: http.ProxyFromEnvironment, - DialContext: (&net.Dialer{ - Timeout: dialTimeout, - KeepAlive: dialKeepAlive, - }).DialContext, + Proxy: http.ProxyFromEnvironment, + DialContext: dialContext, ForceAttemptHTTP2: true, MaxIdleConns: maxIdleConns, IdleConnTimeout: idleConnTimeout, diff --git a/pkg/huggingface-api/client.go b/pkg/huggingface-api/client.go index 1d1c7ae3ce8d..79943b88df84 100644 --- a/pkg/huggingface-api/client.go +++ b/pkg/huggingface-api/client.go @@ -96,21 +96,49 @@ type Client struct { maxRetries int retryBackoff time.Duration maxBackoff time.Duration - sleepFn func(time.Duration) + clock Clock +} + +// Clock is the small portion of wall-clock time used by retry handling. +// Supplying a fake clock lets tests verify backoff behavior without sleeping. +type Clock interface { + Now() time.Time + Sleep(time.Duration) +} + +type realClock struct{} + +func (realClock) Now() time.Time { return time.Now() } +func (realClock) Sleep(d time.Duration) { time.Sleep(d) } + +// ClientOption configures a Hugging Face API client. +type ClientOption func(*Client) + +// WithClock replaces the clock used for retry delays. +func WithClock(clock Clock) ClientOption { + return func(client *Client) { + if clock != nil { + client.clock = clock + } + } } var ErrRateLimited = errors.New("huggingface API rate limited") // NewClient creates a new Hugging Face API client -func NewClient() *Client { - return &Client{ +func NewClient(options ...ClientOption) *Client { + client := &Client{ baseURL: "https://huggingface.co/api/models", client: httpclient.New(httpclient.WithFollowRedirects()), maxRetries: 5, retryBackoff: 1 * time.Second, maxBackoff: 30 * time.Second, - sleepFn: time.Sleep, + clock: realClock{}, + } + for _, option := range options { + option(client) } + return client } func (c *Client) newRequest(ctx context.Context, method, rawURL, token string) (*http.Request, error) { @@ -143,7 +171,7 @@ func (c *Client) SearchModels(params SearchParams) ([]Model, error) { resp, err := c.client.Do(req) if err != nil { if attempt < c.maxRetries { - c.sleepFn(c.exponentialBackoff(attempt)) + c.clock.Sleep(c.exponentialBackoff(attempt)) continue } return nil, fmt.Errorf("failed to make request: %w", err) @@ -154,7 +182,7 @@ func (c *Client) SearchModels(params SearchParams) ([]Model, error) { return nil, fmt.Errorf("failed to close response body: %w", err) } if c.isRetryableStatus(resp.StatusCode) && attempt < c.maxRetries { - c.sleepFn(c.retryDelay(resp, attempt)) + c.clock.Sleep(c.retryDelay(resp, attempt)) continue } if resp.StatusCode == http.StatusTooManyRequests { @@ -199,7 +227,7 @@ func (c *Client) retryDelay(resp *http.Response, attempt int) time.Duration { return delay } if at, err := http.ParseTime(retryAfter); err == nil { - delay := time.Until(at) + delay := at.Sub(c.clock.Now()) if delay > 0 { if delay > c.maxBackoff { return c.maxBackoff diff --git a/pkg/huggingface-api/client_test.go b/pkg/huggingface-api/client_test.go index feac4dba3c95..591c26aaf042 100644 --- a/pkg/huggingface-api/client_test.go +++ b/pkg/huggingface-api/client_test.go @@ -14,14 +14,28 @@ import ( hfapi "github.com/mudler/LocalAI/pkg/huggingface-api" ) +type fakeClock struct { + now time.Time + sleeps []time.Duration +} + +func (c *fakeClock) Now() time.Time { return c.now } + +func (c *fakeClock) Sleep(d time.Duration) { + c.sleeps = append(c.sleeps, d) + c.now = c.now.Add(d) +} + var _ = Describe("HuggingFace API Client", func() { var ( client *hfapi.Client server *httptest.Server + clock *fakeClock ) BeforeEach(func() { - client = hfapi.NewClient() + clock = &fakeClock{now: time.Date(2026, time.January, 1, 0, 0, 0, 0, time.UTC)} + client = hfapi.NewClient(hfapi.WithClock(clock)) }) AfterEach(func() { @@ -211,14 +225,35 @@ var _ = Describe("HuggingFace API Client", func() { Search: "GGUF", } - start := time.Now() models, err := client.SearchModels(params) - elapsed := time.Since(start) Expect(err).ToNot(HaveOccurred()) Expect(models).To(HaveLen(0)) Expect(attempts).To(Equal(2)) - Expect(elapsed).To(BeNumerically(">=", 900*time.Millisecond)) + Expect(clock.sleeps).To(Equal([]time.Duration{time.Second})) + }) + + It("should calculate HTTP-date Retry-After using the injected clock", func() { + attempts := 0 + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attempts++ + if attempts == 1 { + w.Header().Set("Retry-After", clock.now.Add(2*time.Second).Format(http.TimeFormat)) + w.WriteHeader(http.StatusTooManyRequests) + return + } + w.Header().Set("Content-Type", "application/json") + _, err := w.Write([]byte("[]")) + Expect(err).ToNot(HaveOccurred()) + })) + client.SetBaseURL(server.URL) + + models, err := client.SearchModels(hfapi.SearchParams{Search: "GGUF"}) + + Expect(err).ToNot(HaveOccurred()) + Expect(models).To(BeEmpty()) + Expect(attempts).To(Equal(2)) + Expect(clock.sleeps).To(Equal([]time.Duration{2 * time.Second})) }) It("should fail fast on non-retryable 4xx responses", func() { @@ -267,6 +302,9 @@ var _ = Describe("HuggingFace API Client", func() { Expect(errors.Is(err, hfapi.ErrRateLimited)).To(BeTrue()) Expect(err.Error()).To(ContainSubstring("Status code: 429")) Expect(models).To(BeNil()) + Expect(clock.sleeps).To(Equal([]time.Duration{ + time.Second, time.Second, time.Second, time.Second, + })) }) }) @@ -336,8 +374,12 @@ var _ = Describe("HuggingFace API Client", func() { Context("when handling network errors", func() { It("should handle connection failures gracefully", func() { - // Use an invalid URL to simulate connection failure - client.SetBaseURL("http://invalid-url-that-does-not-exist") + // A closed loopback listener produces a deterministic connection + // failure without relying on DNS or public network access. + closedServer := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) + closedURL := closedServer.URL + closedServer.Close() + client.SetBaseURL(closedURL) params := hfapi.SearchParams{ Sort: "lastModified", @@ -351,11 +393,26 @@ var _ = Describe("HuggingFace API Client", func() { Expect(err).To(HaveOccurred()) Expect(err.Error()).To(ContainSubstring("failed to make request")) Expect(models).To(BeNil()) + Expect(clock.sleeps).To(Equal([]time.Duration{ + time.Second, 2 * time.Second, 4 * time.Second, 8 * time.Second, + })) }) }) - Context("when getting file SHA on remote model", func() { + Context("when getting file SHA from repository metadata", func() { It("should get file SHA successfully", func() { + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, err := w.Write([]byte(`[{ + "type":"file", + "path":"localai-functioncall-qwen2.5-7b-v0.5-q4_k_m.gguf", + "size":42, + "oid":"pointer-oid", + "lfs":{"oid":"4e7b7fe1d54b881f1ef90799219dc6cc285d29db24f559c8998d1addb35713d4","size":42,"pointerSize":128} + }]`)) + Expect(err).NotTo(HaveOccurred()) + })) + client.SetBaseURL(server.URL + "/api/models") sha, err := client.GetFileSHA( "mudler/LocalAI-functioncall-qwen2.5-7b-v0.5-Q4_K_M-GGUF", "localai-functioncall-qwen2.5-7b-v0.5-q4_k_m.gguf") Expect(err).ToNot(HaveOccurred()) @@ -886,14 +943,33 @@ var _ = Describe("HuggingFace API Client", func() { }) }) - Context("integration test with real HuggingFace API", func() { - It("should recursively list all files including subfolders from real repository", func() { - // This test makes actual API calls to HuggingFace - // Skip if running in CI or if network is not available - realClient := hfapi.NewClient() + Context("repository API compatibility fixtures", func() { + It("should recursively list all files including subfolders", func() { + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + var response string + switch { + case strings.HasSuffix(r.URL.Path, "/tree/main"): + response = `[ + {"type":"file","path":"README.md","size":100,"oid":"readme-oid"}, + {"type":"directory","path":"Q4_K_M","size":0,"oid":"directory-oid"} + ]` + case strings.HasSuffix(r.URL.Path, "/tree/main/Q4_K_M"): + response = `[ + {"type":"file","path":"Q4_K_M/model-00001-of-00002.gguf","size":1000,"oid":"model-oid"} + ]` + default: + w.WriteHeader(http.StatusNotFound) + return + } + _, err := w.Write([]byte(response)) + Expect(err).NotTo(HaveOccurred()) + })) + fixtureClient := hfapi.NewClient() + fixtureClient.SetBaseURL(server.URL + "/api/models") repoID := "bartowski/Qwen_Qwen3-Next-80B-A3B-Instruct-GGUF" - files, err := realClient.ListFiles(repoID) + files, err := fixtureClient.ListFiles(repoID) Expect(err).ToNot(HaveOccurred()) Expect(files).ToNot(BeEmpty(), "should return at least some files") @@ -956,12 +1032,19 @@ var _ = Describe("HuggingFace API Client", func() { }) It("should populate PipelineTag and LibraryName on ModelDetails", func() { - // Sentence-transformers/all-MiniLM-L6-v2 is a public, stable repo: - // pipeline_tag: sentence-similarity, library_name: sentence-transformers. - // This exercises the /api/models/{repo} metadata fetch layered on top - // of ListFiles in GetModelDetails. - realClient := hfapi.NewClient() - details, err := realClient.GetModelDetails("sentence-transformers/all-MiniLM-L6-v2") + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if strings.Contains(r.URL.Path, "/tree/main") { + _, err := w.Write([]byte(`[{"type":"file","path":"config.json","size":100,"oid":"config-oid"}]`)) + Expect(err).NotTo(HaveOccurred()) + return + } + _, err := w.Write([]byte(`{"pipeline_tag":"sentence-similarity","library_name":"sentence-transformers"}`)) + Expect(err).NotTo(HaveOccurred()) + })) + fixtureClient := hfapi.NewClient() + fixtureClient.SetBaseURL(server.URL + "/api/models") + details, err := fixtureClient.GetModelDetails("sentence-transformers/all-MiniLM-L6-v2") Expect(err).ToNot(HaveOccurred()) Expect(details).ToNot(BeNil()) Expect(details.PipelineTag).To(Equal("sentence-similarity")) diff --git a/pkg/model/loader_test.go b/pkg/model/loader_test.go index 1a882943127f..da5b2f037282 100644 --- a/pkg/model/loader_test.go +++ b/pkg/model/loader_test.go @@ -65,8 +65,7 @@ var _ = Describe("ModelLoader", func() { BeforeEach(func() { // Setup the model loader with a test directory - modelPath = "/tmp/test_model_path" - os.Mkdir(modelPath, 0755) + modelPath = GinkgoT().TempDir() systemState, err := system.GetSystemState( system.WithModelPath(modelPath), @@ -75,11 +74,6 @@ var _ = Describe("ModelLoader", func() { modelLoader = model.NewModelLoader(systemState) }) - AfterEach(func() { - // Cleanup test directory - os.RemoveAll(modelPath) - }) - Context("NewModelLoader", func() { It("should create a new ModelLoader with an empty model map", func() { Expect(modelLoader).ToNot(BeNil()) diff --git a/pkg/oci/blob.go b/pkg/oci/blob.go index e034c41622a9..63aa44d5f122 100644 --- a/pkg/oci/blob.go +++ b/pkg/oci/blob.go @@ -16,6 +16,10 @@ import ( ) func FetchImageBlob(ctx context.Context, r, reference, dst string, statusReader func(ocispec.Descriptor) io.Writer) error { + return fetchImageBlob(ctx, r, reference, dst, statusReader, false) +} + +func fetchImageBlob(ctx context.Context, r, reference, dst string, statusReader func(ocispec.Descriptor) io.Writer, plainHTTP bool) error { // 0. Create a file store for the output fs, err := os.Create(dst) if err != nil { @@ -29,6 +33,7 @@ func FetchImageBlob(ctx context.Context, r, reference, dst string, statusReader return fmt.Errorf("failed to create repository: %v", err) } repo.SkipReferrersGC = true + repo.PlainHTTP = plainHTTP // Identify LocalAI to the registry. This mirrors oras' auth.DefaultClient // (same retry policy) but advertises a LocalAI User-Agent instead of the diff --git a/pkg/oci/blob_test.go b/pkg/oci/blob_test.go index cef29a972228..76b7d49b6e2f 100644 --- a/pkg/oci/blob_test.go +++ b/pkg/oci/blob_test.go @@ -1,10 +1,14 @@ -package oci_test +package oci import ( "context" + "crypto/sha256" + "fmt" + "net/http" + "net/http/httptest" "os" + "strings" - . "github.com/mudler/LocalAI/pkg/oci" // Update with your module path . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) @@ -12,11 +16,22 @@ import ( var _ = Describe("OCI", func() { Context("pulling images", func() { It("should fetch blobs correctly", func() { + payload := []byte("local OCI blob fixture") + digest := fmt.Sprintf("sha256:%x", sha256.Sum256(payload)) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Docker-Distribution-API-Version", "registry/2.0") + w.Header().Set("Content-Length", fmt.Sprint(len(payload))) + if r.Method != http.MethodHead { + _, _ = w.Write(payload) + } + })) + defer server.Close() f, err := os.CreateTemp("", "ollama") Expect(err).NotTo(HaveOccurred()) defer os.RemoveAll(f.Name()) - err = FetchImageBlob(context.TODO(), "registry.ollama.ai/library/gemma", "sha256:c1864a5eb19305c40519da12cc543519e48a0697ecd30e15d5ac228644957d12", f.Name(), nil) + err = fetchImageBlob(context.Background(), strings.TrimPrefix(server.URL, "http://")+"/library/gemma", digest, f.Name(), nil, true) Expect(err).NotTo(HaveOccurred()) + Expect(os.ReadFile(f.Name())).To(Equal(payload)) }) }) }) diff --git a/pkg/oci/cosignverify/verify.go b/pkg/oci/cosignverify/verify.go index a644a9ded83a..abb2db8f55b4 100644 --- a/pkg/oci/cosignverify/verify.go +++ b/pkg/oci/cosignverify/verify.go @@ -33,6 +33,8 @@ import ( "github.com/sigstore/sigstore-go/pkg/root" "github.com/sigstore/sigstore-go/pkg/tuf" "github.com/sigstore/sigstore-go/pkg/verify" + + "github.com/mudler/LocalAI/pkg/httpclient" ) // Policy is the verification policy a backend image must satisfy. @@ -288,7 +290,7 @@ func enforceNotBefore(result *verify.VerificationResult, cutoff time.Time) error func (v *Verifier) remoteOptions(ctx context.Context) []remote.Option { t := v.transport if t == nil { - t = http.DefaultTransport + t = httpclient.HardenedTransport() } // Match the retry policy used elsewhere in pkg/oci so transient // registry hiccups don't fail verification. diff --git a/pkg/oci/image_test.go b/pkg/oci/image_test.go index 447bc90f61ca..300cfc683ea1 100644 --- a/pkg/oci/image_test.go +++ b/pkg/oci/image_test.go @@ -1,10 +1,15 @@ package oci_test import ( + "archive/tar" + "bytes" "context" "os" - "runtime" + "path/filepath" + "github.com/google/go-containerregistry/pkg/v1/empty" + "github.com/google/go-containerregistry/pkg/v1/mutate" + "github.com/google/go-containerregistry/pkg/v1/tarball" "github.com/mudler/LocalAI/pkg/oci" . "github.com/mudler/LocalAI/pkg/oci" // Update with your module path . "github.com/onsi/ginkgo/v2" @@ -15,25 +20,30 @@ var _ = Describe("OCI", func() { Context("when template is loaded successfully", func() { It("should evaluate the template correctly", func() { - if runtime.GOOS == "darwin" { - Skip("Skipping test on darwin") - } - imageName := "alpine" - img, err := GetImage(imageName, "", nil, nil) + var layerTar bytes.Buffer + writer := tar.NewWriter(&layerTar) + content := []byte("offline OCI fixture\n") + Expect(writer.WriteHeader(&tar.Header{Name: "fixture.txt", Mode: 0o644, Size: int64(len(content))})).To(Succeed()) + _, err := writer.Write(content) Expect(err).NotTo(HaveOccurred()) + Expect(writer.Close()).To(Succeed()) - size, err := GetOCIImageSize(imageName, "", nil, nil) + layer, err := tarball.LayerFromReader(bytes.NewReader(layerTar.Bytes())) Expect(err).NotTo(HaveOccurred()) - - Expect(size).ToNot(Equal(int64(0))) + img, err := mutate.AppendLayers(empty.Image, layer) + Expect(err).NotTo(HaveOccurred()) + size, err := layer.Size() + Expect(err).NotTo(HaveOccurred()) + Expect(size).To(BeNumerically(">", 0)) // Create tempdir dir, err := os.MkdirTemp("", "example") Expect(err).NotTo(HaveOccurred()) - defer os.RemoveAll(dir) + DeferCleanup(os.RemoveAll, dir) - err = ExtractOCIImage(context.TODO(), img, imageName, dir, nil) + err = ExtractOCIImage(context.TODO(), img, "fixture:offline", dir, nil) Expect(err).NotTo(HaveOccurred()) + Expect(os.ReadFile(filepath.Join(dir, "fixture.txt"))).To(Equal(content)) }) }) }) diff --git a/pkg/oci/ollama.go b/pkg/oci/ollama.go index f0a874013a16..42bf64bfca0f 100644 --- a/pkg/oci/ollama.go +++ b/pkg/oci/ollama.go @@ -35,6 +35,10 @@ type LayerDetail struct { } func OllamaModelManifest(image string) (*Manifest, error) { + return ollamaModelManifest("https", "registry.ollama.ai", image) +} + +func ollamaModelManifest(scheme, registry, image string) (*Manifest, error) { // parse the repository and tag from `image`. `image` should be for e.g. gemma:2b, or foobar/gemma:2b // if there is a : in the image, then split it @@ -42,7 +46,7 @@ func OllamaModelManifest(image string) (*Manifest, error) { tag, repository, image := ParseImageParts(image) // get e.g. https://registry.ollama.ai/v2/library/llama3/manifests/latest - req, err := http.NewRequest("GET", "https://registry.ollama.ai/v2/"+repository+"/"+image+"/manifests/"+tag, nil) + req, err := http.NewRequest("GET", scheme+"://"+registry+"/v2/"+repository+"/"+image+"/manifests/"+tag, nil) if err != nil { return nil, err } @@ -65,7 +69,11 @@ func OllamaModelManifest(image string) (*Manifest, error) { } func OllamaModelBlob(image string) (string, error) { - manifest, err := OllamaModelManifest(image) + return ollamaModelBlob("https", "registry.ollama.ai", image) +} + +func ollamaModelBlob(scheme, registry, image string) (string, error) { + manifest, err := ollamaModelManifest(scheme, registry, image) if err != nil { return "", err } @@ -81,12 +89,16 @@ func OllamaModelBlob(image string) (string, error) { } func OllamaFetchModel(ctx context.Context, image string, output string, statusWriter func(ocispec.Descriptor) io.Writer) error { + return ollamaFetchModel(ctx, "https", "registry.ollama.ai", image, output, statusWriter) +} + +func ollamaFetchModel(ctx context.Context, scheme, registry, image string, output string, statusWriter func(ocispec.Descriptor) io.Writer) error { _, repository, imageNoTag := ParseImageParts(image) - blobID, err := OllamaModelBlob(image) + blobID, err := ollamaModelBlob(scheme, registry, image) if err != nil { return err } - return FetchImageBlob(ctx, fmt.Sprintf("registry.ollama.ai/%s/%s", repository, imageNoTag), blobID, output, statusWriter) + return fetchImageBlob(ctx, fmt.Sprintf("%s/%s/%s", registry, repository, imageNoTag), blobID, output, statusWriter, scheme == "http") } diff --git a/pkg/oci/ollama_test.go b/pkg/oci/ollama_test.go index fbda69e6b40e..bed92a19c01b 100644 --- a/pkg/oci/ollama_test.go +++ b/pkg/oci/ollama_test.go @@ -1,10 +1,15 @@ -package oci_test +package oci import ( "context" + "crypto/sha256" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" "os" + "strings" - . "github.com/mudler/LocalAI/pkg/oci" // Update with your module path . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) @@ -12,11 +17,26 @@ import ( var _ = Describe("OCI", func() { Context("ollama", func() { It("pulls model files", func() { + payload := []byte("local Ollama model fixture") + digest := fmt.Sprintf("sha256:%x", sha256.Sum256(payload)) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Docker-Distribution-API-Version", "registry/2.0") + if strings.Contains(r.URL.Path, "/manifests/") { + _ = json.NewEncoder(w).Encode(Manifest{SchemaVersion: 2, Layers: []LayerDetail{{Digest: digest, MediaType: "application/vnd.ollama.image.model", Size: len(payload)}}}) + return + } + w.Header().Set("Content-Length", fmt.Sprint(len(payload))) + if r.Method != http.MethodHead { + _, _ = w.Write(payload) + } + })) + defer server.Close() f, err := os.CreateTemp("", "ollama") Expect(err).NotTo(HaveOccurred()) defer os.RemoveAll(f.Name()) - err = OllamaFetchModel(context.TODO(), "gemma:2b", f.Name(), nil) + err = ollamaFetchModel(context.Background(), "http", strings.TrimPrefix(server.URL, "http://"), "gemma:2b", f.Name(), nil) Expect(err).NotTo(HaveOccurred()) + Expect(os.ReadFile(f.Name())).To(Equal(payload)) }) }) }) diff --git a/pkg/testnetwork/guard.go b/pkg/testnetwork/guard.go new file mode 100644 index 000000000000..17db55c0e0ea --- /dev/null +++ b/pkg/testnetwork/guard.go @@ -0,0 +1,57 @@ +// SPDX-License-Identifier: MIT + +// Package testnetwork provides an explicit outbound-network guard for tests. +package testnetwork + +import ( + "context" + "fmt" + "net" + "net/netip" + "strings" +) + +type Guard struct { + Dialer net.Dialer + Dial func(context.Context, string, string) (net.Conn, error) + Allowed []netip.Prefix +} + +func LocalGuard() *Guard { + prefixes := []string{"127.0.0.0/8", "::1/128", "10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16"} + guard := &Guard{} + for _, prefix := range prefixes { + guard.Allowed = append(guard.Allowed, netip.MustParsePrefix(prefix)) + } + return guard +} + +func (g *Guard) DialContext(ctx context.Context, network, address string) (net.Conn, error) { + originalAddress := address + host, _, err := net.SplitHostPort(address) + if err != nil { + return nil, fmt.Errorf("test network guard: invalid address %q: %w", address, err) + } + addresses, err := net.DefaultResolver.LookupNetIP(ctx, "ip", strings.Trim(host, "[]")) + if err != nil { + return nil, fmt.Errorf("test network guard: resolve %q: %w", host, err) + } + for _, resolved := range addresses { + if !g.allowed(resolved.Unmap()) { + return nil, fmt.Errorf("test network guard: public dial blocked: %s (%s)", host, resolved) + } + } + if g.Dial != nil { + return g.Dial(ctx, network, originalAddress) + } + return g.Dialer.DialContext(ctx, network, originalAddress) +} + +func (g *Guard) allowed(address netip.Addr) bool { + for _, prefix := range g.Allowed { + if prefix.Contains(address) { + return true + } + } + return false +} diff --git a/pkg/testnetwork/guard_suite_test.go b/pkg/testnetwork/guard_suite_test.go new file mode 100644 index 000000000000..9259f5a64860 --- /dev/null +++ b/pkg/testnetwork/guard_suite_test.go @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: MIT + +package testnetwork_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestNetworkGuard(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Test network guard suite") +} diff --git a/pkg/testnetwork/guard_test.go b/pkg/testnetwork/guard_test.go new file mode 100644 index 000000000000..5f658eb19a49 --- /dev/null +++ b/pkg/testnetwork/guard_test.go @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: MIT + +package testnetwork_test + +import ( + "context" + "errors" + "net" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/pkg/testnetwork" +) + +var _ = Describe("Guard", func() { + It("blocks a public IP before dialing", func() { + _, err := testnetwork.LocalGuard().DialContext(context.Background(), "tcp", "203.0.113.1:443") + Expect(err).To(MatchError(ContainSubstring("public dial blocked"))) + }) + + It("allows loopback fixtures", func() { + guard := testnetwork.LocalGuard() + called := false + guard.Dial = func(_ context.Context, network, address string) (net.Conn, error) { + called = true + Expect(network).To(Equal("tcp")) + Expect(address).To(Equal("127.0.0.1:8080")) + return nil, errors.New("fixture dial sentinel") + } + _, err := guard.DialContext(context.Background(), "tcp", "127.0.0.1:8080") + Expect(err).To(MatchError("fixture dial sentinel")) + Expect(called).To(BeTrue()) + }) +}) diff --git a/scripts/run-coverage.sh b/scripts/run-coverage.sh index 430ca1f11fe9..771fa4242bb3 100755 --- a/scripts/run-coverage.sh +++ b/scripts/run-coverage.sh @@ -21,6 +21,13 @@ # "!real-models" (those specs need a downloaded model). # COVERAGE_EXCLUDE_RE egrep pattern of profile lines to drop before merging, # e.g. generated protobuf (grpc/proto/.*\.pb\.go). +# COVERAGE_PROCS parallel Ginkgo processes; 0 lets Ginkgo detect CPUs. +# COVERAGE_SUITE_TIMEOUT maximum duration of each recursive root (default 5m). +# COVERAGE_PROGRESS_AFTER emit diagnostics when a spec is slow (default 30s). +# +# Verbose Ginkgo output is retained in OUTPUT_DIR/logs. The previous run's log +# for each root is kept with a .previous suffix, so a noisy failure remains +# available without flooding the commit-hook output. # # Why one ginkgo invocation per root: passing several recursive roots to a # single ginkgo run only merges ONE root's coverprofile into --output-dir @@ -41,11 +48,43 @@ shift 3 unit_roots="$*" # space-free tokens (./pkg ./core) mkdir -p "$out_dir" +lock_dir="$out_dir/.run-coverage.lock" +if ! mkdir "$lock_dir" 2>/dev/null; then + echo "run-coverage: another coverage run is using $out_dir" >&2 + echo "run-coverage: wait for it to finish; if none is running, remove stale lock $lock_dir" >&2 + exit 2 +fi +cleanup() { + for root in $unit_roots ${COVERAGE_E2E_ROOTS:-}; do + # --keep-separate-coverprofiles leaves Go's package test binaries behind. + # They are deterministic runner artifacts, never source inputs. + find "$root" -type f -name '*.test' -delete 2>/dev/null || : + done + rmdir "$lock_dir" 2>/dev/null || : +} +trap cleanup EXIT +trap 'exit 130' HUP INT TERM + +log_dir="$out_dir/logs" +mkdir -p "$log_dir" # Clear per-root profiles from a previous run: the merge collects them by glob, # so a stale profile (e.g. from a root that failed to rebuild this run) must not # leak into the merged result. rm -f "$out_dir"/cover-*.out +rm -f "$out_dir"/*_cover-*.out +rm -f "$merged" fail=0 +timings="$out_dir/timings.tsv" +: > "$timings" +run_started="$(date +%s)" + +procs="${COVERAGE_PROCS:-0}" +suite_timeout="${COVERAGE_SUITE_TIMEOUT:-5m}" +progress_after="${COVERAGE_PROGRESS_AFTER:-30s}" +parallel_flags="-p --keep-going --timeout=$suite_timeout --poll-progress-after=$progress_after --poll-progress-interval=10s" +if [ "$procs" -gt 0 ] 2>/dev/null; then + parallel_flags="$parallel_flags --procs=$procs --compilers=$procs" +fi # Common optional flags go into "$@"; unquoted ${VAR:+...} would word-split a # --tags value that contains a space. The unit roots were captured above, so @@ -59,26 +98,133 @@ profile_name() { printf 'cover-%s.out' "$(printf '%s' "$1" | sed 's#[./][./]*#_#g; s#^_##; s#_$##')" } +log_name() { + printf '%s.log' "$(printf '%s' "$1" | sed 's#[./][./]*#_#g; s#^_##; s#_$##')" +} + +rotate_log() { + log="$1" + if [ -f "$log" ]; then + mv -f "$log" "$log.previous" + fi +} + +# Ginkgo's recursive-run merger can fail after every suite has passed when a +# large --coverpkg run produces many profiles. Keep its profiles separate and +# merge them here using the same block-summing rule as the cross-root merge. +consolidate_root_profiles() { + base="$1" + set -- "$out_dir"/*_"$base" + if [ ! -e "$1" ]; then + echo "run-coverage: no per-package profiles produced for $base" >&2 + return 1 + fi + tmp="$out_dir/.${base}.tmp" + { + echo "mode: atomic" + awk ' + /^mode:/ { next } + { stmts[$1] = $2; cnt[$1] += $3 } + END { for (k in stmts) print k, stmts[k], cnt[k] } + ' "$@" + } > "$tmp" + mv "$tmp" "$out_dir/$base" + rm -f "$@" +} + +report_failure() { + root="$1" + log="$2" + echo "run-coverage: FAIL — tests under coverage failed for $root" >&2 + echo "run-coverage: full output: $log" >&2 + echo "run-coverage: relevant tail:" >&2 + # Keep the terminal useful even when Ginkgo emits thousands of verbose lines. + # The complete log remains available when this short extract is insufficient. + summary="$(grep -E 'Summarizing|\[FAIL(ED)?\]|FAIL!|--- FAIL:|Test Suite Failed|could not finalize|Status code: 429|HTTP 429|rate limit|timed out|panic:|fork/exec|no such file or directory|Expected.*(but got|success)' "$log" \ + | tail -n 30)" + if [ -n "$summary" ]; then + printf '%s\n' "$summary" >&2 + else + tail -n 30 "$log" >&2 + fi +} + +record_timing() { + root="$1" + started="$2" + elapsed="$(( $(date +%s) - started ))" + printf '%s\t%s\n' "$root" "$elapsed" >> "$timings" + echo "run-coverage: TIMING — $root ${elapsed}s" +} + +print_timing_summary() { + echo "run-coverage: wall-clock summary" + sort -t "$(printf '\t')" -k2,2nr "$timings" | awk -F '\t' '{ printf " %5ss %s\n", $2, $1 }' + echo " $(( $(date +%s) - run_started ))s total" + echo "run-coverage: slowest specs/hooks taking at least ${COVERAGE_SLOW_SPEC_THRESHOLD:-3}s (up to ${COVERAGE_SLOW_SPEC_LIMIT:-25} per root)" + found=0 + while IFS="$(printf '\t')" read -r root elapsed; do + log="$log_dir/$(log_name "$root")" + entries="$(scripts/summarize-ginkgo-waits.sh "${COVERAGE_SLOW_SPEC_THRESHOLD:-3}" "$root" "$log" "${COVERAGE_SLOW_SPEC_LIMIT:-25}")" + if [ -n "$entries" ]; then + printf '%s\n' "$entries" + found=1 + fi + done < "$timings" + [ "$found" -eq 1 ] || echo " none" +} + # Unit/suite roots: recursive. for root in $unit_roots; do base="$(profile_name "$root")" - go run github.com/onsi/ginkgo/v2/ginkgo --flake-attempts "$flakes" -v -r "$@" \ - --cover --covermode=atomic --coverprofile="$base" --output-dir="$out_dir" "$root" || fail=1 + log="$log_dir/$(log_name "$root")" + rotate_log "$log" + echo "run-coverage: testing $root (full output: $log)" + started="$(date +%s)" + # parallel_flags is intentionally word-split: it contains CLI arguments only. + # shellcheck disable=SC2086 + go run github.com/onsi/ginkgo/v2/ginkgo $parallel_flags --keep-separate-coverprofiles --flake-attempts "$flakes" -v -r "$@" \ + --cover --covermode=atomic --coverprofile="$base" --output-dir="$out_dir" "$root" >"$log" 2>&1 \ + && consolidate_root_profiles "$base" \ + && echo "run-coverage: PASS — $root" \ + || { fail=1; report_failure "$root" "$log"; } + record_timing "$root" "$started" done # In-process integration roots: NON-recursive + optional label filter. for root in ${COVERAGE_E2E_ROOTS:-}; do base="$(profile_name "$root")" + log="$log_dir/$(log_name "$root")" + rotate_log "$log" + echo "run-coverage: testing $root (full output: $log)" + started="$(date +%s)" if [ -n "${COVERAGE_E2E_LABELS:-}" ]; then - go run github.com/onsi/ginkgo/v2/ginkgo --flake-attempts "$flakes" -v "$@" \ + # shellcheck disable=SC2086 + go run github.com/onsi/ginkgo/v2/ginkgo $parallel_flags --keep-separate-coverprofiles --flake-attempts "$flakes" -v "$@" \ --label-filter="$COVERAGE_E2E_LABELS" \ - --cover --covermode=atomic --coverprofile="$base" --output-dir="$out_dir" "$root" || fail=1 + --cover --covermode=atomic --coverprofile="$base" --output-dir="$out_dir" "$root" >"$log" 2>&1 \ + && consolidate_root_profiles "$base" \ + && echo "run-coverage: PASS — $root" \ + || { fail=1; report_failure "$root" "$log"; } else - go run github.com/onsi/ginkgo/v2/ginkgo --flake-attempts "$flakes" -v "$@" \ - --cover --covermode=atomic --coverprofile="$base" --output-dir="$out_dir" "$root" || fail=1 + # shellcheck disable=SC2086 + go run github.com/onsi/ginkgo/v2/ginkgo $parallel_flags --keep-separate-coverprofiles --flake-attempts "$flakes" -v "$@" \ + --cover --covermode=atomic --coverprofile="$base" --output-dir="$out_dir" "$root" >"$log" 2>&1 \ + && consolidate_root_profiles "$base" \ + && echo "run-coverage: PASS — $root" \ + || { fail=1; report_failure "$root" "$log"; } fi + record_timing "$root" "$started" done +print_timing_summary + +if [ "$fail" -ne 0 ]; then + echo "run-coverage: FAILED — one or more test suites failed; no merged profile was produced." >&2 + echo "run-coverage: the coverage percentage ratchet was not run." >&2 + exit "$fail" +fi + # Collect the per-root profiles by glob (space-safe, no list to track). set -- "$out_dir"/cover-*.out if [ ! -e "$1" ]; then @@ -98,4 +244,4 @@ fi ' "$@" } > "$merged" -exit "$fail" +echo "run-coverage: all test suites passed; merged profile: $merged" diff --git a/scripts/run-test-linux-offline.sh b/scripts/run-test-linux-offline.sh new file mode 100755 index 000000000000..c2be4931dc00 --- /dev/null +++ b/scripts/run-test-linux-offline.sh @@ -0,0 +1,42 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: MIT +set -euo pipefail + +if [[ $(uname -s) != Linux ]]; then + echo 'kernel-level test egress enforcement is Linux-only' >&2 + exit 2 +fi +if [[ $# -lt 2 ]]; then + echo "usage: $0 TARGET COMMAND [ARG...]" >&2 + exit 2 +fi + +root=$(cd "$(dirname "$0")/.." && pwd) +group="localai-test-$$" +cgroup="/sys/fs/cgroup/$group" +parent_cgroup="/sys/fs/cgroup$(awk -F: '$1 == "0" {print $3}' /proc/self/cgroup)" + +sudo mkdir "$cgroup" +cleanup() { + echo $$ | sudo tee "$parent_cgroup/cgroup.procs" >/dev/null 2>&1 || true + sudo iptables -D OUTPUT -m cgroup --path "$group" -j REJECT 2>/dev/null || true + sudo iptables -D OUTPUT -m cgroup --path "$group" -d 192.168.0.0/16 -j ACCEPT 2>/dev/null || true + sudo iptables -D OUTPUT -m cgroup --path "$group" -d 172.16.0.0/12 -j ACCEPT 2>/dev/null || true + sudo iptables -D OUTPUT -m cgroup --path "$group" -d 10.0.0.0/8 -j ACCEPT 2>/dev/null || true + sudo iptables -D OUTPUT -m cgroup --path "$group" -d 127.0.0.0/8 -j ACCEPT 2>/dev/null || true + sudo ip6tables -D OUTPUT -m cgroup --path "$group" -d ::1/128 -j ACCEPT 2>/dev/null || true + sudo ip6tables -D OUTPUT -m cgroup --path "$group" -j REJECT 2>/dev/null || true + sudo rmdir "$cgroup" 2>/dev/null || true +} +trap cleanup EXIT INT TERM + +sudo iptables -I OUTPUT 1 -m cgroup --path "$group" -j REJECT +sudo iptables -I OUTPUT 1 -m cgroup --path "$group" -d 192.168.0.0/16 -j ACCEPT +sudo iptables -I OUTPUT 1 -m cgroup --path "$group" -d 172.16.0.0/12 -j ACCEPT +sudo iptables -I OUTPUT 1 -m cgroup --path "$group" -d 10.0.0.0/8 -j ACCEPT +sudo iptables -I OUTPUT 1 -m cgroup --path "$group" -d 127.0.0.0/8 -j ACCEPT +sudo ip6tables -I OUTPUT 1 -m cgroup --path "$group" -j REJECT +sudo ip6tables -I OUTPUT 1 -m cgroup --path "$group" -d ::1/128 -j ACCEPT +echo $$ | sudo tee "$cgroup/cgroup.procs" >/dev/null + +LOCALAI_TEST_KERNEL_ACTIVE=1 "$root/scripts/run-test-offline.sh" "$@" diff --git a/scripts/run-test-offline.sh b/scripts/run-test-offline.sh new file mode 100755 index 000000000000..614ba71beb1d --- /dev/null +++ b/scripts/run-test-offline.sh @@ -0,0 +1,19 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: MIT +set -euo pipefail + +if [[ $# -lt 2 ]]; then + echo "usage: $0 TARGET COMMAND [ARG...]" >&2 + exit 2 +fi + +target=$1 +shift +root=$(cd "$(dirname "$0")/.." && pwd) + +if [[ ${LOCALAI_TEST_KERNEL_ENFORCE:-0} == 1 && ${LOCALAI_TEST_KERNEL_ACTIVE:-0} != 1 ]]; then + exec "$root/scripts/run-test-linux-offline.sh" "$target" "$@" +fi + +exec go run "$root/cmd/test-resources" run "$target" \ + "$root/test-resources/manifests" "${TEST_RESOURCE_CACHE:-$root/.cache/test-resources}" -- "$@" diff --git a/scripts/summarize-ginkgo-waits.sh b/scripts/summarize-ginkgo-waits.sh new file mode 100755 index 000000000000..545e55340253 --- /dev/null +++ b/scripts/summarize-ginkgo-waits.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env sh +# summarize-ginkgo-waits.sh THRESHOLD_SECONDS ROOT LOG [LIMIT] +set -eu + +threshold="${1:?missing threshold in seconds}" +root="${2:?missing test root}" +log="${3:?missing Ginkgo log}" +limit="${4:-25}" +case "$limit" in + ''|*[!0-9]*) echo "limit must be a non-negative integer" >&2; exit 2 ;; +esac + +# Strip terminal colour sequences, then report slow specs and hooks. This +# catches sleeps, polling, channel waits, teardown and any other idle time +# without unsafe attempts to replace Go's process-wide clock primitives. +sed 's/\[[0-9;]*[[:alpha:]]//g' "$log" | awk -v threshold="$threshold" -v root="$root" ' + /\[[0-9]+([.][0-9]+)? seconds\]/ { + line = $0 + sub(/^.*\[/, "", line) + sub(/ seconds\].*$/, "", line) + seconds = line + 0 + if (seconds < threshold) next + description = "(description unavailable)" + location = "(location unavailable)" + if (getline > 0) description = $0 + if (getline > 0) location = $0 + printf "%010.3f\t%-18s\t%s\t%s\n", seconds, root, description, location + } +' | sort -t "$(printf '\t')" -k1,1nr | sed -n "1,${limit}p" | awk -F '\t' '{ printf " %7.3fs %s %s %s\n", $1 + 0, $2, $3, $4 }' diff --git a/scripts/test-network-lint.sh b/scripts/test-network-lint.sh new file mode 100755 index 000000000000..e081f7993175 --- /dev/null +++ b/scripts/test-network-lint.sh @@ -0,0 +1,39 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: MIT +set -euo pipefail + +# The full-tree fingerprint makes this effective on a clean CI checkout (where +# a worktree-only diff would always be empty). Most existing direct clients are +# loopback fixtures; changing the inventory requires an intentional baseline +# update after review. +expected_inventory=2885a428cdab55eea357dae3ec47b3d44f9999b59542d06cd3d792cf491c76b3 +inventory=$( + { + rg --no-heading --no-line-number --glob '*_test.go' \ + '(http\.(Get|Post|Head)\(|http\.Default(Client|Transport)|net\.Dial\(|exec\.Command\([^,]+,[[:space:]]*"(curl|wget)")' \ + pkg core tests backend || true + rg --no-heading --no-line-number --glob '*.sh' '(curl|wget)[[:space:]]' tests backend || true + } | LC_ALL=C sort +) +if command -v sha256sum >/dev/null 2>&1; then + actual_inventory=$(printf '%s\n' "$inventory" | sha256sum | awk '{print $1}') +else + actual_inventory=$(printf '%s\n' "$inventory" | shasum -a 256 | awk '{print $1}') +fi +if [[ $actual_inventory != "$expected_inventory" ]]; then + echo 'Test network mechanism inventory changed; remove the direct access or review and update the lint baseline:' >&2 + echo "$inventory" >&2 + exit 1 +fi + +# Also give contributors a focused diagnostic for newly introduced remote +# literals and direct mechanisms instead of only reporting the fingerprint. +base=${TEST_NETWORK_LINT_BASE:-HEAD} +violations=$(git diff --unified=0 "$base" -- api pkg core tests backend | \ + rg '^\+[^+].*(http\.(Get|Post|Head)\(|http\.Default(Client|Transport)|net\.Dial\(|exec\.Command\([^,]+,[[:space:]]*"(curl|wget)"|https?://)' | \ + rg -v 'test-network: fixture' || true) +if [[ -n "$violations" ]]; then + echo 'Direct test network access is forbidden; use a fixture or guarded transport:' >&2 + echo "$violations" >&2 + exit 1 +fi diff --git a/test-resources/manifests/aio.json b/test-resources/manifests/aio.json new file mode 100644 index 000000000000..c21ff9a97375 --- /dev/null +++ b/test-resources/manifests/aio.json @@ -0,0 +1,12 @@ +{ + "version": 1, + "target": "aio", + "files": [ + { + "url": "https://cdn.openai.com/whisper/draft-20220913a/micro-machines.wav", + "sha256": "37de21902b32aa2fc147ccbfdcc0566cc7061fffb2c0b10874f05147c0b9de0f", + "destination": "audio/micro-machines.wav", + "environment": "AIO_AUDIO_FIXTURE" + } + ] +} diff --git a/test-resources/manifests/backend.json b/test-resources/manifests/backend.json new file mode 100644 index 000000000000..178b8b9c577b --- /dev/null +++ b/test-resources/manifests/backend.json @@ -0,0 +1 @@ +{"version":1,"target":"backend"} diff --git a/test-resources/manifests/default-darwin.json b/test-resources/manifests/default-darwin.json new file mode 100644 index 000000000000..1dc5aed0389e --- /dev/null +++ b/test-resources/manifests/default-darwin.json @@ -0,0 +1,7 @@ +{ + "version": 1, + "target": "default-darwin", + "http": [], + "images": [], + "files": [] +} diff --git a/test-resources/manifests/default.json b/test-resources/manifests/default.json new file mode 100644 index 000000000000..b4bf99004d8e --- /dev/null +++ b/test-resources/manifests/default.json @@ -0,0 +1,10 @@ +{ + "version": 1, + "target": "default", + "images": [ + { + "reference": "docker.io/library/postgres@sha256:33f923b05f64ca54ac4401c01126a6b92afe839a0aa0a52bc5aeb5cc958e5f20", + "sha256": "bd98262690143a2a05167b3e924cde07e82ea269e45d5610e00a45f50e76d45a" + } + ] +} diff --git a/test-resources/manifests/distributed-e2e.json b/test-resources/manifests/distributed-e2e.json new file mode 100644 index 000000000000..b9c353d9c292 --- /dev/null +++ b/test-resources/manifests/distributed-e2e.json @@ -0,0 +1,14 @@ +{ + "version": 1, + "target": "distributed-e2e", + "images": [ + { + "reference": "docker.io/library/postgres@sha256:57c72fd2a128e416c7fcc499958864df5301e940bca0a56f58fddf30ffc07777", + "sha256": "fee330cbd34786da2b211fe2e4b7424d7bf974466cd3db0e376b2e4c457a339c" + }, + { + "reference": "docker.io/library/nats@sha256:c11af972c99ae542de8925e6a7d9c533aa1eb039660420d2074beed6089b3bf0", + "sha256": "96e0f53430696eadaf1d7e250914ab78779cadc40df715a168ebe342e039b6f2" + } + ] +} diff --git a/test-resources/manifests/external-probes.json b/test-resources/manifests/external-probes.json new file mode 100644 index 000000000000..af7a85e52f01 --- /dev/null +++ b/test-resources/manifests/external-probes.json @@ -0,0 +1 @@ +{"version":1,"target":"external-probes"} diff --git a/test-resources/manifests/hardware.json b/test-resources/manifests/hardware.json new file mode 100644 index 000000000000..cddf8b0c33f6 --- /dev/null +++ b/test-resources/manifests/hardware.json @@ -0,0 +1 @@ +{"version":1,"target":"hardware"} diff --git a/test-resources/manifests/lock.json b/test-resources/manifests/lock.json new file mode 100644 index 000000000000..e68826e0ad33 --- /dev/null +++ b/test-resources/manifests/lock.json @@ -0,0 +1,12 @@ +{ + "version": 1, + "bundles": { + "aio": "sha256:03b05eedb51c853b0f05f2f3f592e2edc210e337388d3ff5a766680c0a166e55", + "backend": "embedded", + "default": "sha256:185ebd3cfb994b9c1d1d1d9fad0a5d95c9de698b45e010276db388a9d66474bd", + "default-darwin": "embedded", + "distributed-e2e": "sha256:797dad68952914bf6612a42486086aa45eb17112067bec9955e2f5cf65a030dc", + "external-probes": "embedded", + "hardware": "embedded" + } +} diff --git a/tests/e2e-aio/e2e_suite_test.go b/tests/e2e-aio/e2e_suite_test.go index f82b7c5e7b2f..581ad654efc6 100644 --- a/tests/e2e-aio/e2e_suite_test.go +++ b/tests/e2e-aio/e2e_suite_test.go @@ -6,14 +6,13 @@ import ( "os" "runtime" "testing" - "time" + "github.com/mudler/LocalAI/internal/testfixtures" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "github.com/openai/openai-go/v3" "github.com/openai/openai-go/v3/option" "github.com/testcontainers/testcontainers-go" - "github.com/testcontainers/testcontainers-go/wait" ) var container testcontainers.Container @@ -39,10 +38,10 @@ var _ = BeforeSuite(func() { if apiEndpoint == "" { startDockerImage() - apiPort, err := container.MappedPort(context.Background(), defaultApiPort) + apiAddress, err := testfixtures.ContainerEndpoint(context.Background(), container, defaultApiPort) Expect(err).To(Not(HaveOccurred())) - apiEndpoint = "http://localhost:" + apiPort.Port() + "/v1" // So that other tests can reference this value safely. + apiEndpoint = "http://" + apiAddress + "/v1" // test-network: fixture } else { GinkgoWriter.Printf("docker apiEndpoint set from env: %q\n", apiEndpoint) } @@ -122,15 +121,16 @@ func startDockerImage() { Target: "/backends", }, }, - WaitingFor: wait.ForAll( - wait.ForListeningPort(defaultApiPort).WithStartupTimeout(10*time.Minute), - wait.ForHTTP("/v1/models").WithPort(defaultApiPort).WithStartupTimeout(10*time.Minute), - ), } GinkgoWriter.Printf("Launching Docker Container %s:%s\n", containerImage, containerImageTag) ctx := context.Background() + imageReference := fmt.Sprintf("%s:%s", containerImage, containerImageTag) + Expect(testfixtures.RequireImage(ctx, imageReference, "aio")).To(Succeed()) + testNetwork, err := testfixtures.DockerNetwork() + Expect(err).NotTo(HaveOccurred()) + req.Networks = []string{testNetwork} c, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{ ContainerRequest: req, Started: true, diff --git a/tests/e2e-aio/e2e_test.go b/tests/e2e-aio/e2e_test.go index 6472b5d63cff..3382885933ff 100644 --- a/tests/e2e-aio/e2e_test.go +++ b/tests/e2e-aio/e2e_test.go @@ -3,11 +3,13 @@ package e2e_test import ( "bytes" "context" + "encoding/base64" "encoding/json" "fmt" "io" "net/http" "os" + "path/filepath" "github.com/mudler/LocalAI/core/schema" . "github.com/onsi/ginkgo/v2" @@ -257,6 +259,9 @@ var _ = Describe("E2E test", func() { Context("vision", func() { It("correctly", func() { + image, err := os.ReadFile(filepath.Join("..", "..", "core", "http", "static", "logo.png")) + Expect(err).NotTo(HaveOccurred()) + imageURI := "data:image/png;base64," + base64.StdEncoding.EncodeToString(image) model := "gpt-4o" resp, err := client.Chat.Completions.New(context.TODO(), openai.ChatCompletionNewParams{ @@ -276,7 +281,7 @@ var _ = Describe("E2E test", func() { { OfImageURL: &openai.ChatCompletionContentPartImageParam{ ImageURL: openai.ChatCompletionContentPartImageImageURLParam{ - URL: "https://picsum.photos/id/22/4434/3729", + URL: imageURI, Detail: "low", }, }, @@ -289,7 +294,7 @@ var _ = Describe("E2E test", func() { }) Expect(err).ToNot(HaveOccurred()) Expect(len(resp.Choices)).To(Equal(1), fmt.Sprint(resp)) - Expect(resp.Choices[0].Message.Content).To(Or(ContainSubstring("man"), ContainSubstring("road")), fmt.Sprint(resp.Choices[0].Message.Content)) + Expect(resp.Choices[0].Message.Content).NotTo(BeEmpty()) }) }) @@ -310,11 +315,7 @@ var _ = Describe("E2E test", func() { Context("audio to text", func() { It("correctly", func() { - downloadURL := "https://cdn.openai.com/whisper/draft-20220913a/micro-machines.wav" - file, err := downloadHttpFile(downloadURL) - Expect(err).ToNot(HaveOccurred()) - - fileHandle, err := os.Open(file) + fileHandle, err := os.Open(preparedAudioFixture()) Expect(err).ToNot(HaveOccurred()) defer fileHandle.Close() @@ -328,11 +329,7 @@ var _ = Describe("E2E test", func() { }) It("with VTT format", func() { - downloadURL := "https://cdn.openai.com/whisper/draft-20220913a/micro-machines.wav" - file, err := downloadHttpFile(downloadURL) - Expect(err).ToNot(HaveOccurred()) - - fileHandle, err := os.Open(file) + fileHandle, err := os.Open(preparedAudioFixture()) Expect(err).ToNot(HaveOccurred()) defer fileHandle.Close() @@ -443,25 +440,10 @@ var _ = Describe("E2E test", func() { }) }) -func downloadHttpFile(url string) (string, error) { - resp, err := http.Get(url) - if err != nil { - return "", err - } - defer resp.Body.Close() - - tmpfile, err := os.CreateTemp("", "example") - if err != nil { - return "", err - } - defer tmpfile.Close() - - _, err = io.Copy(tmpfile, resp.Body) - if err != nil { - return "", err - } - - return tmpfile.Name(), nil +func preparedAudioFixture() string { + path := os.Getenv("AIO_AUDIO_FIXTURE") + Expect(path).NotTo(BeEmpty(), "run `make prepare-offline-test-cache TEST_RESOURCE_SET=aio` before the AIO suite") + return path } func requestRerank(modelName, query string, documents []string, topN *int, apiEndpoint string) (*http.Response, []byte) { diff --git a/tests/e2e/distributed/nats_jwt_helpers_test.go b/tests/e2e/distributed/nats_jwt_helpers_test.go index 80060ef6a801..e0723b378a92 100644 --- a/tests/e2e/distributed/nats_jwt_helpers_test.go +++ b/tests/e2e/distributed/nats_jwt_helpers_test.go @@ -8,6 +8,7 @@ import ( "time" "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/internal/testfixtures" "github.com/mudler/LocalAI/pkg/natsauth" "github.com/nats-io/jwt/v2" "github.com/nats-io/nkeys" @@ -17,6 +18,8 @@ import ( "github.com/testcontainers/testcontainers-go" tcnats "github.com/testcontainers/testcontainers-go/modules/nats" + tcnetwork "github.com/testcontainers/testcontainers-go/network" + "github.com/testcontainers/testcontainers-go/wait" ) // JWTTestInfra holds a NATS server configured with JWT auth and minted worker credentials. @@ -34,6 +37,9 @@ func SetupJWTInfra() *JWTTestInfra { GinkgoHelper() infra := &JWTTestInfra{TestInfra: &TestInfra{Ctx: context.Background()}} + Expect(testfixtures.RequireImage(infra.Ctx, testfixtures.NATS2Alpine, "distributed-e2e")).To(Succeed()) + testNetwork, err := testfixtures.DockerNetwork() + Expect(err).NotTo(HaveOccurred()) operatorJWT, accountJWT, accountSeed, err := jwtResolverMaterial() Expect(err).ToNot(HaveOccurred()) @@ -51,15 +57,18 @@ resolver_preload: { var natsContainer *tcnats.NATSContainer // Override default testcontainers -js: JetStream fails without a system account in JWT mode. - natsContainer, err = tcnats.Run(infra.Ctx, "nats:2-alpine", + natsContainer, err = tcnats.Run(infra.Ctx, testfixtures.NATS2Alpine, tcnats.WithConfigFile(bytes.NewBufferString(conf)), testcontainers.WithCmd("-c", "/etc/nats.conf"), + tcnetwork.WithNetworkName([]string{"nats"}, testNetwork), + testcontainers.WithWaitStrategy(wait.ForLog("Server is ready")), ) Expect(err).ToNot(HaveOccurred()) infra.NATSContainer = natsContainer - infra.NatsURL, err = infra.NATSContainer.ConnectionString(infra.Ctx) + natsEndpoint, err := testfixtures.ContainerEndpoint(infra.Ctx, infra.NATSContainer, "4222") Expect(err).ToNot(HaveOccurred()) + infra.NatsURL = "nats://" + natsEndpoint infra.NodeID = "550e8400-e29b-41d4-a716-446655440000" cfg := natsauth.Config{AccountSeed: infra.AccountSeed, WorkerJWTTTL: time.Hour} @@ -153,4 +162,4 @@ func accountPublicKeyFromSeed(accountSeed string) string { func nodeSubjectPrefix(nodeID string) string { tok := strings.NewReplacer(".", "-", "*", "-", ">", "-", " ", "-", "\t", "-", "\n", "-").Replace(nodeID) return "nodes." + tok -} \ No newline at end of file +} diff --git a/tests/e2e/distributed/testhelpers_test.go b/tests/e2e/distributed/testhelpers_test.go index 68cf537e30bd..d54c0897a1aa 100644 --- a/tests/e2e/distributed/testhelpers_test.go +++ b/tests/e2e/distributed/testhelpers_test.go @@ -2,9 +2,11 @@ package distributed_test import ( "context" + "fmt" "time" "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/internal/testfixtures" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -12,6 +14,7 @@ import ( "github.com/testcontainers/testcontainers-go" tcnats "github.com/testcontainers/testcontainers-go/modules/nats" tcpostgres "github.com/testcontainers/testcontainers-go/modules/postgres" + tcnetwork "github.com/testcontainers/testcontainers-go/network" "github.com/testcontainers/testcontainers-go/wait" ) @@ -32,9 +35,13 @@ func SetupInfra(dbName string) *TestInfra { infra := &TestInfra{Ctx: context.Background()} var err error + Expect(testfixtures.RequireImage(infra.Ctx, testfixtures.Postgres16Alpine, "distributed-e2e")).To(Succeed()) + Expect(testfixtures.RequireImage(infra.Ctx, testfixtures.NATS2Alpine, "distributed-e2e")).To(Succeed()) + testNetwork, err := testfixtures.DockerNetwork() + Expect(err).NotTo(HaveOccurred()) // Start PostgreSQL container - infra.PGContainer, err = tcpostgres.Run(infra.Ctx, "postgres:16-alpine", + infra.PGContainer, err = tcpostgres.Run(infra.Ctx, testfixtures.Postgres16Alpine, tcpostgres.WithDatabase(dbName), tcpostgres.WithUsername("test"), tcpostgres.WithPassword("test"), @@ -43,18 +50,23 @@ func SetupInfra(dbName string) *TestInfra { WithOccurrence(2). WithStartupTimeout(30*time.Second), ), + tcnetwork.WithNetworkName([]string{"postgres"}, testNetwork), ) Expect(err).ToNot(HaveOccurred()) - infra.PGURL, err = infra.PGContainer.ConnectionString(infra.Ctx, "sslmode=disable") + pgEndpoint, err := testfixtures.ContainerEndpoint(infra.Ctx, infra.PGContainer, "5432") Expect(err).ToNot(HaveOccurred()) + infra.PGURL = fmt.Sprintf("postgres://test:test@%s/%s?sslmode=disable", pgEndpoint, dbName) // Start NATS container - infra.NATSContainer, err = tcnats.Run(infra.Ctx, "nats:2-alpine") + infra.NATSContainer, err = tcnats.Run(infra.Ctx, testfixtures.NATS2Alpine, + tcnetwork.WithNetworkName([]string{"nats"}, testNetwork), + testcontainers.WithWaitStrategy(wait.ForLog("Server is ready"))) Expect(err).ToNot(HaveOccurred()) - infra.NatsURL, err = infra.NATSContainer.ConnectionString(infra.Ctx) + natsEndpoint, err := testfixtures.ContainerEndpoint(infra.Ctx, infra.NATSContainer, "4222") Expect(err).ToNot(HaveOccurred()) + infra.NatsURL = "nats://" + natsEndpoint // Connect messaging client infra.NC, err = messaging.New(infra.NatsURL) @@ -83,12 +95,18 @@ func SetupNATSOnly() *TestInfra { infra := &TestInfra{Ctx: context.Background()} var err error + Expect(testfixtures.RequireImage(infra.Ctx, testfixtures.NATS2Alpine, "distributed-e2e")).To(Succeed()) + testNetwork, err := testfixtures.DockerNetwork() + Expect(err).NotTo(HaveOccurred()) - infra.NATSContainer, err = tcnats.Run(infra.Ctx, "nats:2-alpine") + infra.NATSContainer, err = tcnats.Run(infra.Ctx, testfixtures.NATS2Alpine, + tcnetwork.WithNetworkName([]string{"nats"}, testNetwork), + testcontainers.WithWaitStrategy(wait.ForLog("Server is ready"))) Expect(err).ToNot(HaveOccurred()) - infra.NatsURL, err = infra.NATSContainer.ConnectionString(infra.Ctx) + natsEndpoint, err := testfixtures.ContainerEndpoint(infra.Ctx, infra.NATSContainer, "4222") Expect(err).ToNot(HaveOccurred()) + infra.NatsURL = "nats://" + natsEndpoint infra.NC, err = messaging.New(infra.NatsURL) Expect(err).ToNot(HaveOccurred()) diff --git a/tests/e2e/e2e_suite_test.go b/tests/e2e/e2e_suite_test.go index 38e49f1cc9f3..31d8d0d9cc66 100644 --- a/tests/e2e/e2e_suite_test.go +++ b/tests/e2e/e2e_suite_test.go @@ -180,6 +180,11 @@ var _ = BeforeSuite(func() { "model": name + ".bin", }, } + if name == "mock-llm" { + // Realtime classifier tests exercise concurrent generation and + // scoring on the same model, matching the production slot setup. + cfg["known_usecases"] = []string{"chat", "score"} + } data, err := yaml.Marshal(cfg) Expect(err).ToNot(HaveOccurred()) Expect(os.WriteFile(filepath.Join(modelsPath, name+".yaml"), data, 0644)).To(Succeed()) @@ -245,6 +250,42 @@ var _ = BeforeSuite(func() { Expect(err).ToNot(HaveOccurred()) Expect(os.WriteFile(filepath.Join(modelsPath, "realtime-pipeline.yaml"), pipelineData, 0644)).To(Succeed()) + // Classifier-mode pipeline (LocalAI extension): responses are + // prefill-scored against the option list via the mock backend's + // ROUTE_HINT-driven Score instead of being generated. Threshold 0.6: + // a hinted option scores ≈0.99, no hint leaves a uniform distribution + // (0.5 for two options) and triggers the fallback reply. + classifierPipelineCfg := map[string]any{ + "name": "realtime-pipeline-classifier", + "pipeline": map[string]any{ + "vad": "mock-vad", + "transcription": "mock-stt", + "llm": "mock-llm", + "tts": "mock-tts", + "classifier": map[string]any{ + "enabled": true, + "threshold": 0.6, + "fallback": map[string]any{"mode": "reply", "reply": "Say again?"}, + "options": []map[string]any{ + { + "id": "up", + "description": "the user asks the drone to fly up", + "reply": "Going up.", + "tool": map[string]any{"name": "move", "arguments": map[string]any{"direction": "up"}}, + }, + { + "id": "greeting", + "description": "the user greets the assistant", + "reply": "Hello.", + }, + }, + }, + }, + } + classifierPipelineData, err := yaml.Marshal(classifierPipelineCfg) + Expect(err).ToNot(HaveOccurred()) + Expect(os.WriteFile(filepath.Join(modelsPath, "realtime-pipeline-classifier.yaml"), classifierPipelineData, 0644)).To(Succeed()) + // Speaker-recognition model (mock-backend) + a voice-recognition-gated // pipeline for the realtime gate e2e. The reference WAV carries a positive // DC bias so the mock embeds it to one orthogonal "speaker"; the test then diff --git a/tests/e2e/realtime_classifier_ws_test.go b/tests/e2e/realtime_classifier_ws_test.go new file mode 100644 index 000000000000..6ee196f4aeb7 --- /dev/null +++ b/tests/e2e/realtime_classifier_ws_test.go @@ -0,0 +1,241 @@ +package e2e_test + +import ( + "time" + + "github.com/gorilla/websocket" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// Classifier mode (LocalAI extension): the pipeline scores each user turn +// against a fixed option list via the Score primitive and emits the winning +// option's canned reply / tool call instead of generating. The mock backend +// scores by ROUTE_HINT= markers in the prompt (see mock-backend Score), +// so these specs steer the outcome deterministically through user text. +var _ = Describe("Realtime classifier mode", Label("Realtime"), func() { + const model = "realtime-pipeline-classifier" + + openSession := func() *websocket.Conn { + c := connectWS(model) + created := readServerEvent(c, 30*time.Second) + Expect(created["type"]).To(Equal("session.created")) + sendClientEvent(c, disableVADEvent()) + drainUntil(c, "session.updated", 10*time.Second) + return c + } + + sendUserText := func(c *websocket.Conn, text string) { + sendClientEvent(c, map[string]any{ + "type": "conversation.item.create", + "item": map[string]any{ + "type": "message", + "role": "user", + "content": []map[string]any{ + {"type": "input_text", "text": text}, + }, + }, + }) + drainUntil(c, "conversation.item.added", 10*time.Second) + } + + textResponseCreate := map[string]any{ + "type": "response.create", + "response": map[string]any{ + "output_modalities": []string{"text"}, + }, + } + + It("echoes the classifier config in session.created", func() { + c := connectWS(model) + defer func() { _ = c.Close() }() + + created := readServerEvent(c, 30*time.Second) + Expect(created["type"]).To(Equal("session.created")) + session, ok := created["session"].(map[string]any) + Expect(ok).To(BeTrue()) + classifier, ok := session["localai_classifier"].(map[string]any) + Expect(ok).To(BeTrue(), "session.created should carry localai_classifier") + options, ok := classifier["options"].([]any) + Expect(ok).To(BeTrue()) + Expect(options).To(HaveLen(2)) + }) + + It("picks the hinted option and emits its canned reply and tool call", func() { + c := openSession() + defer func() { _ = c.Close() }() + + sendUserText(c, "ROUTE_HINT=up please fly higher") + sendClientEvent(c, textResponseCreate) + + result := drainUntil(c, "localai.classifier.result", 30*time.Second) + Expect(result["chosen_id"]).To(Equal("up")) + Expect(result["fallback"]).To(BeNil()) + scores, ok := result["scores"].([]any) + Expect(ok).To(BeTrue()) + Expect(scores).To(HaveLen(2)) + top, ok := scores[0].(map[string]any) + Expect(ok).To(BeTrue()) + Expect(top["id"]).To(Equal("up")) + Expect(top["score"]).To(BeNumerically(">", 0.9)) + + textDone := drainUntil(c, "response.output_text.done", 30*time.Second) + Expect(textDone["text"]).To(Equal("Going up.")) + + fcDone := drainUntil(c, "response.function_call_arguments.done", 30*time.Second) + Expect(fcDone["arguments"]).To(MatchJSON(`{"direction":"up"}`)) + + done := drainUntil(c, "response.done", 30*time.Second) + Expect(done).ToNot(BeNil()) + }) + + It("applies the fallback reply when no option clears the threshold", func() { + c := openSession() + defer func() { _ = c.Close() }() + + sendUserText(c, "mumble mumble nothing matches") + sendClientEvent(c, textResponseCreate) + + result := drainUntil(c, "localai.classifier.result", 30*time.Second) + Expect(result["chosen_id"]).To(BeNil()) + Expect(result["fallback"]).To(Equal("reply")) + + textDone := drainUntil(c, "response.output_text.done", 30*time.Second) + Expect(textDone["text"]).To(Equal("Say again?")) + + done := drainUntil(c, "response.done", 30*time.Second) + Expect(done).ToNot(BeNil()) + }) + + It("honors a client-pushed option list via session.update", func() { + // The drone demo app replaces the whole option list at runtime; + // mirror its exact sequence: session.update with + // localai_classifier, then a turn that hints the NEW option. + c := openSession() + defer func() { _ = c.Close() }() + + sendClientEvent(c, map[string]any{ + "type": "session.update", + "session": map[string]any{ + "type": "realtime", + "localai_classifier": map[string]any{ + "enabled": true, + "threshold": 0.6, + "fallback": map[string]any{"mode": "reply", "reply": "Say again?"}, + "options": []map[string]any{ + { + "id": "spin", + "description": "the user asks the drone to spin around", + "reply": "Spinning.", + "tool": map[string]any{"name": "move_drone", "arguments": map[string]any{"direction": "spin"}}, + }, + }, + }, + }, + }) + updated := drainUntil(c, "session.updated", 10*time.Second) + session, ok := updated["session"].(map[string]any) + Expect(ok).To(BeTrue()) + classifier, ok := session["localai_classifier"].(map[string]any) + Expect(ok).To(BeTrue(), "session.updated should echo the replaced classifier") + options, ok := classifier["options"].([]any) + Expect(ok).To(BeTrue()) + Expect(options).To(HaveLen(1)) + + sendUserText(c, "ROUTE_HINT=spin do a barrel roll") + sendClientEvent(c, textResponseCreate) + + result := drainUntil(c, "localai.classifier.result", 30*time.Second) + Expect(result["chosen_id"]).To(Equal("spin")) + + fcDone := drainUntil(c, "response.function_call_arguments.done", 30*time.Second) + Expect(fcDone["name"]).To(Equal("move_drone")) + Expect(fcDone["arguments"]).To(MatchJSON(`{"direction":"spin"}`)) + drainUntil(c, "response.done", 30*time.Second) + }) + + It("silently ignores turns that do not address the assistant by name", func() { + c := openSession() + defer func() { _ = c.Close() }() + + sendClientEvent(c, map[string]any{ + "type": "session.update", + "session": map[string]any{ + "type": "realtime", + "localai_classifier": map[string]any{ + "enabled": true, + "threshold": 0.6, + "address": map[string]any{"names": []string{"drone"}, "mode": "ignore"}, + "options": []map[string]any{ + { + "id": "up", + "description": "the user asks the drone to fly higher", + "reply": "Going up.", + }, + }, + }, + }, + }) + drainUntil(c, "session.updated", 10*time.Second) + + // No name mentioned: dropped before scoring, response completes + // with no output. + sendUserText(c, "ROUTE_HINT=up please fly higher") + sendClientEvent(c, textResponseCreate) + result := drainUntil(c, "localai.classifier.result", 30*time.Second) + Expect(result["fallback"]).To(Equal("not_addressed")) + Expect(result["chosen_id"]).To(BeNil()) + scores, ok := result["scores"].([]any) + Expect(ok).To(BeTrue()) + Expect(scores).To(BeEmpty()) + done := drainUntil(c, "response.done", 30*time.Second) + resp, ok := done["response"].(map[string]any) + Expect(ok).To(BeTrue()) + // Empty output marshals as JSON null. + Expect(resp["output"]).To(SatisfyAny(BeNil(), BeEmpty())) + + // Addressed: scored and acted on as usual. + sendUserText(c, "drone ROUTE_HINT=up please fly higher") + sendClientEvent(c, textResponseCreate) + result = drainUntil(c, "localai.classifier.result", 30*time.Second) + Expect(result["chosen_id"]).To(Equal("up")) + textDone := drainUntil(c, "response.output_text.done", 30*time.Second) + Expect(textDone["text"]).To(Equal("Going up.")) + drainUntil(c, "response.done", 30*time.Second) + }) + + It("runs normal generation when the response disables the classifier", func() { + c := openSession() + defer func() { _ = c.Close() }() + + sendUserText(c, "ROUTE_HINT=up please fly higher") + sendClientEvent(c, map[string]any{ + "type": "response.create", + "response": map[string]any{ + "output_modalities": []string{"text"}, + "localai_classifier": map[string]any{"enabled": false}, + }, + }) + + // The classifier must not run: the next terminal must arrive with + // generated (mock-llm) output and no classifier result event. + deadline := time.Now().Add(60 * time.Second) + sawClassifier := false + var text string + for time.Now().Before(deadline) { + evt := readServerEvent(c, time.Until(deadline)) + switch evt["type"] { + case "localai.classifier.result": + sawClassifier = true + case "response.output_text.done": + text, _ = evt["text"].(string) + case "response.done": + Expect(sawClassifier).To(BeFalse(), "classifier must not run when disabled per response") + Expect(text).ToNot(BeEmpty()) + Expect(text).ToNot(Equal("Going up.")) + return + } + } + Fail("timed out waiting for response.done") + }) +})