Compare commits

...
Author SHA1 Message Date
Parth Sareen 9ba5a04914 launch: claude app (#15937) 2026-05-02 19:19:57 -07:00
Bruce MacDonald 938ca6e274 app: source featured models from experimental recommendations endpoint (#15909)
Replace the hardcoded FEATURED_MODELS list with the
/api/experimental/model-recommendations endpoint so the picker stays in
sync with server-driven recommendations. Inline the merge into useModels
(recommendations first, then the rest of /api/tags) and drop the
standalone mergeModels util.
2026-05-01 11:10:20 -07:00
Pratham Agarwal 8f39fff70b fix: resolve OpenClaw gateway launch timeout on Windows by enforcing IPv4 loopback (#15726) 2026-04-30 22:20:08 -04:00
Daniel Hiltgen 4fe5609563 metal: harden for ggml initialization failures (#15755)
* metal: harden for ggml initialization failures

ggml_metal_device_init performs a probe to verify the tensor API compiles.  On
some systems this passes, even though kernel coverage isn't complete, which
results in a later crash when compiling the real kernels.  This change adds a
single retry if any of the error strings match this failure mode to disable the
tensor API.  It also hardens an error case in the Go initDevices to detect
device initialization failures and panic instead of crashing later on a nil
array entry.

Fixes #15734

* review comments

* review comments
2026-04-30 16:28:03 -07:00
Bruce MacDonald 917324bb4d app: remove ollama update url env var used for testing (#15905) 2026-04-30 13:14:08 -07:00
Parth Sareen c7c2837c96 renderers: update gemma4 renderer (#15886) 2026-04-29 18:40:23 -07:00
Parth Sareen b6447caebc launch: use vram bytes for model recommendations (#15885) 2026-04-29 18:40:14 -07:00
Eva H bad32c7244 launch/docs: fix title for pool (#15883) 2026-04-29 17:18:44 -04:00
Eva H ab2e005bf7 app: align the app launch page with ollama launch (#15753) 2026-04-29 14:45:19 -04:00
Parth Sareen 321cc8a2ba server/launch: add model recommendations cache endpoint (#15868) 2026-04-28 17:09:04 -07:00
Daniel HiltgenandEva Ho 87288ced4f New models (#15861)
* mlx: add laguna model support

* convert: support fp8 safetensors import

Decode HF F8_E4M3 safetensors with block scale companions into GGUF-supported tensor types, and record which output tensors came from FP8 source weights.

Use that source-precision metadata during create quantization: default FP8-sourced GGUFs to Q8_0, keep non-FP8 tensors at their original precision for Q8_0, and promote non-FP8 quantizable tensors to Q8_0 for Q4_K requests.

* ggml: add laguna model support

* server: preserve generate logprobs with builtin parsers

Generate requests were dropping logprob-only chunks whenever a builtin parser buffered visible content. Chat already handled this case, but generate only forwarded chunks with visible response, thinking, or tool-call output.

Keep generate chunks that carry logprobs even when the builtin parser has not flushed visible content yet, and add a regression test that exercises the behavior with a generic thinking parser.

* review comments - perf improvements

* ggml: implement nemotron 3 nano omni

* add poolside integration

* update poolside doc

* adapt to new cache setup

* fix test

* fix test

---------

Co-authored-by: Eva Ho <hoyyeva@gmail.com>
2026-04-28 11:50:12 -07:00
Jesse Gross 2bbe2405fe mlxrunner: decouple models from attention cache storage layout
Models build their own attention masks and read K/V directly from
the cache's buffers, which ties them to the cache's storage layout.
That blocks multi-sequence batching — right-padded rows need a
query-padding mask composed onto every model — and rules out
variants like paged attention where K/V isn't one contiguous tensor.

Caches now hand back a per-layer KVHistory holding post-update K, V,
and a MaskApplier that merges the cache's storage restrictions into
the model's logical mask. Models describe their mask in logical
terms; SDPA composes model, padding, and applier contributions and
dispatches to the kernel's causal or no-mask fast path when it can.
KVHistory still exposes K, V, and the composed mask for manual
attention paths (e.g. CUDA prefill at head_dim > 128).

Performance for single-sequence inference is unchanged.
2026-04-27 20:04:46 -07:00
Jesse Gross bd21678b16 mlxrunner: apply RoPE at per-row positions
Switch RoPE from the scalar-offset kernel (mlx_fast_rope) to the
array-offset one (mlx_fast_rope_dynamic) so each batch row can start
at its own position. The pipeline tracks the current position locally
and passes it to the model through Batch.SeqOffsets; each model
materializes that slice into an int32 array for the RoPE call.

Single-sequence behavior is unchanged; this is the wiring needed
before the runner can batch independent sequences.
2026-04-27 20:04:46 -07:00
Jesse Gross 088dfd89a8 mlxrunner: wrap model forward inputs in a Batch struct
Gives a single extension point for per-call context (positions,
sequence IDs, masks) as multi-sequence batching grows, without having
to churn every model's Forward signature again.
2026-04-27 20:04:46 -07:00
Eva H 3cab8a7b02 app/server: fix desktop app startup killing active ollama launch sessions (#15657) 2026-04-27 22:52:53 -04:00
Daniel Hiltgen 03aee88186 mlx: Support NVIDIA TensorRT Model Optimizer import (#15566)
* mlx: Support NVIDIA TensorRT Model Optimizer import

* x/create: support FP8 safetensors import

Decode HF F8_E4M3 safetensors with block scale companions into MLX-importable tensor blobs, including compressed-tensors weight_scale metadata, packed NVFP4 layouts, and mixed-precision tensor headers.

Use that source-precision metadata during create quantization: default FP8-sourced imports to mxfp8, allow source FP8 to target MLX low-bit formats, preserve source-quantized NVFP4 layouts, selectively keep or promote tensors based on their source precision, and detect quantized dtype from mixed-precision safetensors manifests.

* review comments
2026-04-27 18:28:10 -07:00
Daniel Hiltgen ec9b4e9e47 tokenizer: fix multi-regex BPE offset handling (#15844)
Use the current fragment offset when emitting unmatched spans during multi-regex BPE splitting. This avoids duplicating earlier prompt text and inflating token counts for multi-stage BPE tokenizers.
2026-04-27 14:14:27 -07:00
Jesse Gross 4656a07e56 mlxrunner: batch the sampler across multiple sequences
Register sequences with Add/Remove; each Sample call takes any subset of
registered slots and samples one token per row, appending to each slot's
ring-buffer history. When all slots share Options and penalty rings are
full, one fused transform pass runs over the whole batch via a persistent
pooled history tensor; otherwise calls fall back to per-slot serial
processing indexed against the same pool.

Performance is unchanged for a single sequence, which is all that is
exposed for now.
2026-04-25 09:53:53 -07:00
Jesse Gross 30f86cb9dd mlxrunner: track sampler history in a fixed-size ring buffer
AppendToken used to concatenate the new token onto the history tensor
and slice it back to RepeatLastN every decode step, churning the graph
shape and reallocating a fresh tensor each call. The stateful penalties
don't care about order within the window, so a fixed-capacity ring with
one SliceUpdate per append keeps the tensor shape constant across
steps.
2026-04-25 09:53:53 -07:00
Parth Sareen ea01af6f76 openai: map responses reasoning effort to think (#15789) 2026-04-24 02:49:36 -07:00
Parth Sareen c2ebb4d57c api: accept "max" as a think value (#15787) 2026-04-24 01:49:39 -07:00
Parth Sareen 590109c835 launch: harden OpenClaw onboarding flow (#15777) 2026-04-23 16:47:20 -07:00
Eva H b4442c6d17 launch: resave managed integration config when live config drifts (#15776) 2026-04-23 19:32:36 -04:00
Eva H 85ff8e4a21 launch: keep launch recommended models in a fixed canonical order (#15750) 2026-04-23 16:33:00 -04:00
Parth Sareen 160660e572 launch: use bundled OpenClaw ollama web search (#15757) 2026-04-22 16:34:19 -07:00
madflowandParth Sareen 3b43b9bc4b docs: update structured outputs doc for cloud (#15733)
---------

Co-authored-by: Parth Sareen <parth.sareen@ollama.com>
2026-04-22 00:42:39 -07:00
Parth Sareen 21883571b7 launch: replace kimi-k2.5 with k2.6 as top recommended model (#15737) 2026-04-21 15:13:20 -07:00
Jesse Gross ce99f24731 mlxrunner: tokenize prompts in request handler goroutines
Move tokenization out of the single GPU processing goroutine and
into each request's HTTP handler goroutine. This allows the next
request's prompt to be tokenized on the CPU while the current
request is executing on the GPU.
2026-04-21 14:38:49 -07:00
Jesse Gross 04f5f0cdb4 mlx: improve thread safety of array management
Use atomic.Int32 for Array.pinned and a sync.Mutex for the global
arrays slice so MLX arrays can be created and pinned from multiple
goroutines without racing on those structures. Convert Array value
receivers to pointer receivers and struct fields from Array to
*Array to avoid copying the atomic.

This does not fully achieve thread safety even when building
completely independent graphs. The tracing flag and traceScratch
slice in compile.go are unprotected, so concurrent Compile calls
will race. MLX itself is not fully thread-safe either although
it is working to improve.
2026-04-21 14:38:49 -07:00
Matteo Celani fb36a01ffe app/ui: fix model picker showing stale model after switching chats (#15280)
* app/ui: fix model picker showing stale model after switching chats

Optimistic messages created during streaming were storing the full
Model object instead of the model name string. When switching back
to a chat with cached streaming data, the restore effect read an
object where it expected a string, causing the model picker to fail
matching and remain stuck on the previous chat's model.

* app/ui: fix two more instances of Model object passed as model name

Fix the same bug at lines 523 and 536 in the assistant_with_tools
event handler, where selectedModel (object) was used instead of
selectedModel.model (string).
2026-04-21 15:08:06 -04:00
Michael Verrilli 0c65ed33bc cmd: populate model capabilities in launchInteractiveModel (#15712)
launchInteractiveModel was introduced in PR #14609 without the
client.Show() capability-detection block that RunHandler uses.
This left opts.MultiModal always false in the TUI path, causing
image/audio file paths to always be treated as unknown commands
instead of being loaded as multimodal attachments.

Mirror the Show() call, pull-on-404 fallback, cloud auth handling,
and MultiModal/Think population from RunHandler into
launchInteractiveModel.

Fixes #15711
2026-04-21 14:37:36 -04:00
Jesse Gross 22d6c817f8 mlxrunner: fuse top-P and top-K into a single sort pass
When both filters are active, avoid paying for a full sort in top-P
and a partial sort in top-K. Single-filter paths are unchanged.
Improves generation throughput on gemma4:e4b by 1.5%.
2026-04-20 17:43:00 -07:00
Jesse Gross ca01373b28 mlxrunner: use MaxAxis in the min-P sampler
One reduction op instead of Argmax + TakeAlongAxis.
2026-04-20 17:43:00 -07:00
Jesse Gross 24e038d56a mlxrunner: add logprobs support
Match the ollamarunner and OpenAI semantics: raw, full-vocab log-softmax
with the top-K ranked by probability. Skipped on the GPU when the request
doesn't ask for logprobs so decode doesn't pay for it otherwise.
2026-04-20 17:43:00 -07:00
Parth Sareen 5d1021603a server: apply format when think=false for gemma4 (#15678) 2026-04-20 17:42:29 -07:00
Parth Sareen 8e05d734b9 launch: add kimi cli integration with installer flow (#15723) 2026-04-20 15:33:32 -07:00
Jesse Gross 05e0f21bec mlx: fuse sigmoid router head in glm4_moe_lite
DeepSeek-V2-style aux-loss-free routing computes sigmoid(gates) once but
needs it twice: the raw sigmoid output is gathered after top-k, while the
post-bias negation is the argpartition key. Fuse into a single multi-output
Compiled kernel returning both, saving two launches on the routing path
per token. Exposed as a general SigmoidRouter since the same pattern is
shared across DeepSeek-V2 descendants.

Improves glm4.7 generation performance by approximately 1%.
2026-04-20 15:02:14 -07:00
Daniel Hiltgen ff23dd343f mlx: apply repeat penalties in sampler (#15631) 2026-04-18 07:49:38 -07:00
Parth Sareen 123b300af6 docs: update hermes (#15655) 2026-04-17 14:20:59 -07:00
Parth Sareen 57653b8e42 cmd/launch: show WSL guidance on Windows instead of handing off (#15637) 2026-04-16 17:18:04 -07:00
Parth Sareen a50ce61c54 launch: skip unchanged managed-single rewrite (#15633) 2026-04-16 16:20:42 -07:00
Daniel Hiltgen 2bb7ea00d2 create: avoid gc race with create (#15628)
If you have a long running create, and start another ollama server with the
same model dir, the GC algorithm deletes the pending blobs and breaks the
create.  This adds a 1h grace period to avoid deleting in-flight creation
operations.
2026-04-16 13:29:16 -07:00
Daniel Hiltgen 55fa80d07a mlx: additional gemma4 cache fixes (#15607)
Harden additional corner cases
2026-04-16 13:07:19 -07:00
Daniel Hiltgen b9cb535407 mlx: fix gemma4 cache to use logical view (#15617) 2026-04-16 11:54:30 -07:00
Daniel Hiltgen 031baef094 mlx: fix imagegen lookup (#15588)
* mlx: fix imagegen lookup

Fixes #15533 - imagegen had fallen out of sync with the new layout
for multiple mlx libraries on Metal.

* review comments
2026-04-16 10:39:00 -07:00
7d271e6dc9 cmd/launch: add Copilot CLI integration (#15583)
---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: ParthSareen <parth.sareen@ollama.com>
2026-04-15 17:22:53 -07:00
Devon Rifkin c88dae2d6b Merge pull request #15612 from ollama/drifkin/gemma4-split-templates
gemma4: render differently based on model size
2026-04-15 17:15:35 -07:00
Devon Rifkin 9e3618d663 make empty block conditional 2026-04-15 15:35:25 -07:00
Devon Rifkin e585ecd11f gemma4: render differently based on model size
Following up on #15560, this change now has e2b/e4b render differently
from 26b/31b.

For backwards compatibility, we take the existing renderer name `gemma4`
and make it do dynamic resolution based on the model name/size, but the
intended use is for the models to be republished with the renderer
variant specified explicitly: `gemma4-small` or `gemma4-large`.
2026-04-15 14:37:16 -07:00
187 changed files with 25286 additions and 3192 deletions

No files matched your search

+2 -2
View File
@@ -55,7 +55,7 @@ The official [Ollama Docker image](https://hub.docker.com/r/ollama/ollama) `olla
ollama
```
You'll be prompted to run a model or connect Ollama to your existing agents or applications such as `claude`, `codex`, `openclaw` and more.
You'll be prompted to run a model or connect Ollama to your existing agents or applications such as `Claude Code`, `OpenClaw`, `OpenCode` , `Codex`, `Copilot`, and more.
### Coding
@@ -65,7 +65,7 @@ To launch a specific integration:
ollama launch claude
```
Supported integrations include [Claude Code](https://docs.ollama.com/integrations/claude-code), [Codex](https://docs.ollama.com/integrations/codex), [Droid](https://docs.ollama.com/integrations/droid), and [OpenCode](https://docs.ollama.com/integrations/opencode).
Supported integrations include [Claude Code](https://docs.ollama.com/integrations/claude-code), [Codex](https://docs.ollama.com/integrations/codex), [Copilot CLI](https://docs.ollama.com/integrations/copilot-cli), [Droid](https://docs.ollama.com/integrations/droid), and [OpenCode](https://docs.ollama.com/integrations/opencode).
### AI assistant
+10
View File
@@ -368,6 +368,16 @@ func (c *Client) List(ctx context.Context) (*ListResponse, error) {
return &lr, nil
}
// ModelRecommendationsExperimental lists model recommendations from the local
// server's experimental recommendations endpoint.
func (c *Client) ModelRecommendationsExperimental(ctx context.Context) (*ModelRecommendationsResponse, error) {
var resp ModelRecommendationsResponse
if err := c.do(ctx, http.MethodGet, "/api/experimental/model-recommendations", nil, &resp); err != nil {
return nil, err
}
return &resp, nil
}
// ListRunning lists running models.
func (c *Client) ListRunning(ctx context.Context) (*ProcessResponse, error) {
var lr ProcessResponse
+21 -7
View File
@@ -802,6 +802,20 @@ type ListResponse struct {
Models []ListModelResponse `json:"models"`
}
// ModelRecommendationsResponse is the response from [Client.ModelRecommendationsExperimental].
type ModelRecommendationsResponse struct {
Recommendations []ModelRecommendation `json:"recommendations"`
}
// ModelRecommendation is a single recommendation entry in [ModelRecommendationsResponse].
type ModelRecommendation struct {
Model string `json:"model"`
Description string `json:"description"`
ContextLength int `json:"context_length,omitempty"`
MaxOutputTokens int `json:"max_output_tokens,omitempty"`
VRAMBytes int64 `json:"vram_bytes,omitempty"`
}
// ProcessResponse is the response from [Client.Process].
type ProcessResponse struct {
Models []ProcessModelResponse `json:"models"`
@@ -1080,7 +1094,7 @@ func DefaultOptions() Options {
}
}
// ThinkValue represents a value that can be a boolean or a string ("high", "medium", "low")
// ThinkValue represents a value that can be a boolean or a string ("high", "medium", "low", "max")
type ThinkValue struct {
// Value can be a bool or string
Value interface{}
@@ -1096,7 +1110,7 @@ func (t *ThinkValue) IsValid() bool {
case bool:
return true
case string:
return v == "high" || v == "medium" || v == "low"
return v == "high" || v == "medium" || v == "low" || v == "max"
default:
return false
}
@@ -1130,8 +1144,8 @@ func (t *ThinkValue) Bool() bool {
case bool:
return v
case string:
// Any string value ("high", "medium", "low") means thinking is enabled
return v == "high" || v == "medium" || v == "low"
// Any string value ("high", "medium", "low", "max") means thinking is enabled
return v == "high" || v == "medium" || v == "low" || v == "max"
default:
return false
}
@@ -1169,14 +1183,14 @@ func (t *ThinkValue) UnmarshalJSON(data []byte) error {
var s string
if err := json.Unmarshal(data, &s); err == nil {
// Validate string values
if s != "high" && s != "medium" && s != "low" {
return fmt.Errorf("invalid think value: %q (must be \"high\", \"medium\", \"low\", true, or false)", s)
if s != "high" && s != "medium" && s != "low" && s != "max" {
return fmt.Errorf("invalid think value: %q (must be \"high\", \"medium\", \"low\", \"max\", true, or false)", s)
}
t.Value = s
return nil
}
return fmt.Errorf("think must be a boolean or string (\"high\", \"medium\", \"low\", true, or false)")
return fmt.Errorf("think must be a boolean or string (\"high\", \"medium\", \"low\", \"max\", true, or false)")
}
// MarshalJSON implements json.Marshaler
+5
View File
@@ -495,6 +495,11 @@ func TestThinking_UnmarshalJSON(t *testing.T) {
input: `{ "think": "low" }`,
expectedThinking: &ThinkValue{Value: "low"},
},
{
name: "string_max",
input: `{ "think": "max" }`,
expectedThinking: &ThinkValue{Value: "max"},
},
{
name: "invalid_string",
input: `{ "think": "invalid" }`,
-4
View File
@@ -157,10 +157,6 @@ func main() {
}
}
if u := os.Getenv("OLLAMA_UPDATE_URL"); u != "" {
updater.UpdateCheckURLBase = u
}
// Detect if this is a first start after an upgrade, in
// which case we need to do some cleanup
var skipMove bool
+23
View File
@@ -83,6 +83,29 @@ func resolvePath(name string) string {
return name
}
func ollamaServeArgs(args []string) bool {
if len(args) < 2 {
return false
}
switch strings.Trim(filepath.Base(args[0]), `"`) {
case "ollama", "ollama.exe":
default:
return false
}
for _, rawArg := range args[1:] {
arg := strings.Trim(rawArg, `"`)
if strings.HasPrefix(arg, "-") {
continue
}
return arg == "serve" || arg == "start"
}
return false
}
// cleanup checks the pid file for a running ollama process
// and shuts it down gracefully if it is running
func cleanup() error {
+57
View File
@@ -205,6 +205,63 @@ func TestServerCmdCloudSettingEnv(t *testing.T) {
}
}
func TestOllamaServeArgs(t *testing.T) {
tests := []struct {
name string
args []string
want bool
}{
{
name: "system ollama serve",
args: []string{"ollama", "serve"},
want: true,
},
{
name: "relative path ollama serve",
args: []string{"./ollama", "serve"},
want: true,
},
{
name: "serve after other flags",
args: []string{"./ollama", "--verbose", "serve"},
want: true,
},
{
name: "start alias",
args: []string{"ollama", "start"},
want: true,
},
{
name: "launch command",
args: []string{"ollama", "launch", "opencode"},
want: false,
},
{
name: "run command with model named serve",
args: []string{"ollama", "run", "serve"},
want: false,
},
{
name: "launch command with serve in passthrough args",
args: []string{"ollama", "launch", "codex", "--", "-p", "serve"},
want: false,
},
{
name: "different executable",
args: []string{"go", "run", "serve"},
want: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := ollamaServeArgs(tt.args); got != tt.want {
t.Fatalf("ollamaServeArgs(%v) = %v, want %v", tt.args, got, tt.want)
}
})
}
}
func TestGetInferenceInfo(t *testing.T) {
tests := []struct {
name string
+14 -1
View File
@@ -46,7 +46,17 @@ func terminated(pid int) (bool, error) {
return false, nil
}
// reapServers kills all ollama processes except our own
func ollamaServeProcess(pid int) bool {
output, err := exec.Command("ps", "-p", strconv.Itoa(pid), "-o", "args=").Output()
if err != nil {
slog.Debug("failed to inspect ollama process", "pid", pid, "err", err)
return false
}
return ollamaServeArgs(strings.Fields(strings.TrimSpace(string(output))))
}
// reapServers kills external ollama serve processes except our own.
func reapServers() error {
// Get our own PID to avoid killing ourselves
currentPID := os.Getpid()
@@ -82,6 +92,9 @@ func reapServers() error {
if pid == currentPID {
continue
}
if !ollamaServeProcess(pid) {
continue
}
proc, err := os.FindProcess(pid)
if err != nil {
+26 -1
View File
@@ -101,7 +101,29 @@ func terminated(pid int) (bool, error) {
return true, nil
}
// reapServers kills all ollama processes except our own
func ollamaServeProcess(pid int) bool {
cmd := exec.Command("wmic", "process", "where", fmt.Sprintf("ProcessId=%d", pid), "get", "CommandLine", "/value")
cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true}
output, err := cmd.Output()
if err != nil {
slog.Debug("failed to inspect ollama process", "pid", pid, "err", err)
return false
}
for _, line := range strings.Split(string(output), "\n") {
line = strings.TrimSpace(line)
commandLine, ok := strings.CutPrefix(line, "CommandLine=")
if !ok {
continue
}
return ollamaServeArgs(strings.Fields(strings.ToLower(commandLine)))
}
return false
}
// reapServers kills external ollama serve processes except our own.
func reapServers() error {
// Get current process ID to avoid killing ourselves
currentPID := os.Getpid()
@@ -138,6 +160,9 @@ func reapServers() error {
if pid == currentPID {
continue
}
if !ollamaServeProcess(pid) {
continue
}
cmd := exec.Command("taskkill", "/F", "/PID", pidStr)
if err := cmd.Run(); err != nil {
+10 -7
View File
@@ -1201,13 +1201,16 @@ func (db *database) getSettings() (Settings, error) {
func (db *database) setSettings(s Settings) error {
lastHomeView := strings.ToLower(strings.TrimSpace(s.LastHomeView))
validLaunchView := map[string]struct{}{
"launch": {},
"openclaw": {},
"claude": {},
"codex": {},
"opencode": {},
"droid": {},
"pi": {},
"launch": {},
"openclaw": {},
"claude": {},
"claude-desktop": {},
"hermes": {},
"codex": {},
"copilot": {},
"opencode": {},
"droid": {},
"pi": {},
}
if lastHomeView != "chat" {
if _, ok := validLaunchView[lastHomeView]; !ok {
@@ -0,0 +1 @@
<svg height="1em" style="flex:none;line-height:1" viewBox="0 0 24 24" width="1em" xmlns="http://www.w3.org/2000/svg"><title>Claude Code</title><path clip-rule="evenodd" d="M20.998 10.949H24v3.102h-3v3.028h-1.487V20H18v-2.921h-1.487V20H15v-2.921H9V20H7.488v-2.921H6V20H4.487v-2.921H3V14.05H0V10.95h3V5h17.998v5.949zM6 10.949h1.488V8.102H6v2.847zm10.51 0H18V8.102h-1.49v2.847z" fill="#D97757" fill-rule="evenodd"></path></svg>

After

Width:  |  Height:  |  Size: 424 B

@@ -0,0 +1 @@
<svg xmlns="http://www.w3.org/2000/svg" fill="currentColor" fill-rule="evenodd" style="flex:none;line-height:1" viewBox="0 2.5 24 19"><path d="M19.245 5.364c1.322 1.36 1.877 3.216 2.11 5.817.622 0 1.2.135 1.592.654l.73.964c.21.278.323.61.323.955v2.62c0 .339-.173.669-.453.868C20.239 19.602 16.157 21.5 12 21.5c-4.6 0-9.205-2.583-11.547-4.258-.28-.2-.452-.53-.453-.868v-2.62c0-.345.113-.679.321-.956l.73-.963c.392-.517.974-.654 1.593-.654l.029-.297c.25-2.446.81-4.213 2.082-5.52 2.461-2.54 5.71-2.851 7.146-2.864h.198c1.436.013 4.685.323 7.146 2.864zm-7.244 4.328c-.284 0-.613.016-.962.05-.123.447-.305.85-.57 1.108-1.05 1.023-2.316 1.18-2.994 1.18-.638 0-1.306-.13-1.851-.464-.516.165-1.012.403-1.044.996a65.882 65.882 0 00-.063 2.884l-.002.48c-.002.563-.005 1.126-.013 1.69.002.326.204.63.51.765 2.482 1.102 4.83 1.657 6.99 1.657 2.156 0 4.504-.555 6.985-1.657a.854.854 0 00.51-.766c.03-1.682.006-3.372-.076-5.053-.031-.596-.528-.83-1.046-.996-.546.333-1.212.464-1.85.464-.677 0-1.942-.157-2.993-1.18-.266-.258-.447-.661-.57-1.108-.32-.032-.64-.049-.96-.05zm-2.525 4.013c.539 0 .976.426.976.95v1.753c0 .525-.437.95-.976.95a.964.964 0 01-.976-.95v-1.752c0-.525.437-.951.976-.951zm5 0c.539 0 .976.426.976.95v1.753c0 .525-.437.95-.976.95a.964.964 0 01-.976-.95v-1.752c0-.525.437-.951.976-.951zM7.635 5.087c-1.05.102-1.935.438-2.385.906-.975 1.037-.765 3.668-.21 4.224.405.394 1.17.657 1.995.657h.09c.649-.013 1.785-.176 2.73-1.11.435-.41.705-1.433.675-2.47-.03-.834-.27-1.52-.63-1.813-.39-.336-1.275-.482-2.265-.394zm6.465.394c-.36.292-.6.98-.63 1.813-.03 1.037.24 2.06.675 2.47.968.957 2.136 1.104 2.776 1.11h.044c.825 0 1.59-.263 1.995-.657.555-.556.765-3.187-.21-4.224-.45-.468-1.335-.804-2.385-.906-.99-.088-1.875.058-2.265.394zM12 7.615c-.24 0-.525.015-.84.044.03.16.045.336.06.526l-.001.159a2.94 2.94 0 01-.014.25c.225-.022.425-.027.612-.028h.366c.187 0 .387.006.612.028-.015-.146-.015-.277-.015-.409.015-.19.03-.365.06-.526a9.29 9.29 0 00-.84-.044z" fill="white"/></svg>

After

Width:  |  Height:  |  Size: 1.9 KiB

@@ -0,0 +1 @@
<svg xmlns="http://www.w3.org/2000/svg" fill="currentColor" fill-rule="evenodd" style="flex:none;line-height:1" viewBox="0 2.5 24 19"><path d="M19.245 5.364c1.322 1.36 1.877 3.216 2.11 5.817.622 0 1.2.135 1.592.654l.73.964c.21.278.323.61.323.955v2.62c0 .339-.173.669-.453.868C20.239 19.602 16.157 21.5 12 21.5c-4.6 0-9.205-2.583-11.547-4.258-.28-.2-.452-.53-.453-.868v-2.62c0-.345.113-.679.321-.956l.73-.963c.392-.517.974-.654 1.593-.654l.029-.297c.25-2.446.81-4.213 2.082-5.52 2.461-2.54 5.71-2.851 7.146-2.864h.198c1.436.013 4.685.323 7.146 2.864zm-7.244 4.328c-.284 0-.613.016-.962.05-.123.447-.305.85-.57 1.108-1.05 1.023-2.316 1.18-2.994 1.18-.638 0-1.306-.13-1.851-.464-.516.165-1.012.403-1.044.996a65.882 65.882 0 00-.063 2.884l-.002.48c-.002.563-.005 1.126-.013 1.69.002.326.204.63.51.765 2.482 1.102 4.83 1.657 6.99 1.657 2.156 0 4.504-.555 6.985-1.657a.854.854 0 00.51-.766c.03-1.682.006-3.372-.076-5.053-.031-.596-.528-.83-1.046-.996-.546.333-1.212.464-1.85.464-.677 0-1.942-.157-2.993-1.18-.266-.258-.447-.661-.57-1.108-.32-.032-.64-.049-.96-.05zm-2.525 4.013c.539 0 .976.426.976.95v1.753c0 .525-.437.95-.976.95a.964.964 0 01-.976-.95v-1.752c0-.525.437-.951.976-.951zm5 0c.539 0 .976.426.976.95v1.753c0 .525-.437.95-.976.95a.964.964 0 01-.976-.95v-1.752c0-.525.437-.951.976-.951zM7.635 5.087c-1.05.102-1.935.438-2.385.906-.975 1.037-.765 3.668-.21 4.224.405.394 1.17.657 1.995.657h.09c.649-.013 1.785-.176 2.73-1.11.435-.41.705-1.433.675-2.47-.03-.834-.27-1.52-.63-1.813-.39-.336-1.275-.482-2.265-.394zm6.465.394c-.36.292-.6.98-.63 1.813-.03 1.037.24 2.06.675 2.47.968.957 2.136 1.104 2.776 1.11h.044c.825 0 1.59-.263 1.995-.657.555-.556.765-3.187-.21-4.224-.45-.468-1.335-.804-2.385-.906-.99-.088-1.875.058-2.265.394zM12 7.615c-.24 0-.525.015-.84.044.03.16.045.336.06.526l-.001.159a2.94 2.94 0 01-.014.25c.225-.022.425-.027.612-.028h.366c.187 0 .387.006.612.028-.015-.146-.015-.277-.015-.409.015-.19.03-.365.06-.526a9.29 9.29 0 00-.84-.044z"/></svg>

After

Width:  |  Height:  |  Size: 1.9 KiB

@@ -0,0 +1,181 @@
<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" width="1000" height="1000" viewBox="0 0 1000 1000"><circle cx="500.0" cy="500.0" r="500.0" fill="white"/><g transform="translate(100.0 100.0) scale(0.8333333333333334)"><g transform="translate(0.000000,960.000000) scale(0.100000,-0.100000)"
fill="black" stroke="none">
<path d="M4485 9589 c-248 -27 -432 -60 -730 -130 -458 -108 -798 -230 -1207
-435 -533 -267 -1072 -675 -1358 -1030 -205 -255 -442 -748 -535 -1114 -108
-426 -97 -870 29 -1160 73 -169 236 -369 381 -467 139 -94 425 -206 425 -167
0 3 -26 32 -58 63 -74 71 -147 182 -184 278 -16 41 -27 77 -25 79 2 3 31 -28
63 -68 91 -114 153 -177 231 -235 40 -29 77 -62 83 -73 7 -13 4 -85 -10 -242
-10 -123 -24 -281 -30 -353 -6 -71 -29 -296 -51 -500 -135 -1262 -202 -1568
-378 -1733 -69 -64 -105 -77 -216 -77 -132 0 -188 20 -271 97 -110 103 -165
248 -167 438 -1 123 12 191 60 318 20 51 33 95 29 99 -10 10 -92 -74 -136
-137 -153 -223 -204 -504 -146 -790 19 -91 97 -252 161 -333 112 -141 316
-237 504 -237 180 0 419 118 591 290 81 81 107 115 196 255 50 78 54 56 13
-71 -107 -337 -343 -577 -613 -625 -213 -39 -544 99 -707 295 -93 111 -133
199 -173 381 -34 153 -34 154 -46 135 -15 -23 -12 -194 5 -305 42 -280 149
-488 327 -636 133 -111 287 -195 465 -254 62 -20 115 -40 117 -44 3 -4 12 -43
21 -87 44 -222 199 -385 416 -438 96 -24 265 -21 359 5 138 38 281 148 388
297 39 55 48 63 50 45 4 -27 -38 -185 -78 -294 -80 -217 -176 -374 -312 -515
-54 -55 -98 -104 -98 -107 0 -4 245 -7 545 -7 l544 0 126 66 c69 36 130 64
135 62 6 -2 -10 -28 -34 -59 -46 -59 -48 -69 -10 -69 17 0 32 15 58 59 36 59
51 71 87 71 23 0 25 -27 4 -77 -8 -19 -15 -39 -15 -44 0 -12 447 -11 470 1 10
5 85 87 168 182 257 297 400 415 317 263 -16 -30 -26 -57 -23 -60 11 -12 210
100 383 215 115 76 226 143 375 225 58 32 178 100 266 151 150 87 244 132 244
117 0 -4 -78 -86 -174 -182 -207 -208 -246 -254 -334 -386 -48 -71 -122 -157
-259 -298 -117 -120 -193 -206 -193 -217 0 -19 8 -20 115 -20 101 0 116 2 122
18 19 54 112 249 156 327 124 221 283 436 369 502 87 67 225 152 337 209 103
51 297 117 384 130 38 6 29 -2 -92 -75 -74 -45 -170 -107 -214 -139 -106 -76
-247 -225 -332 -351 -67 -99 -70 -102 -79 -77 -6 14 -13 26 -18 26 -11 0 -57
-75 -94 -155 -63 -136 -124 -376 -103 -401 15 -18 152 -19 166 -2 6 7 18 36
27 63 23 72 66 131 217 301 302 339 430 464 590 577 148 104 233 128 520 148
205 14 255 9 447 -38 118 -29 183 -65 298 -163 229 -195 538 -577 579 -715 24
-83 71 -117 109 -79 9 8 16 27 16 42 0 62 -40 212 -73 277 -158 313 -551 635
-922 755 -127 41 -142 55 -52 46 117 -11 231 -35 321 -65 158 -54 259 -123
405 -277 164 -174 283 -379 322 -556 11 -51 25 -113 32 -137 36 -128 148 -81
117 49 -14 60 -7 106 27 177 36 74 62 94 224 179 254 133 425 260 561 417 272
315 403 732 358 1138 -24 218 -84 401 -179 545 -65 98 -155 203 -166 192 -3
-3 7 -37 23 -75 136 -309 136 -725 2 -1063 -130 -330 -441 -760 -633 -879 -60
-37 -60 -20 2 54 478 577 578 1337 315 2390 -25 99 -86 326 -136 505 -169 611
-386 1539 -494 2109 -161 857 -200 998 -442 1606 -161 405 -321 692 -529 950
-93 114 -322 343 -433 431 -203 160 -497 332 -737 428 -293 118 -702 220 -978
246 -129 12 -406 11 -520 -1z m-267 -214 c301 -47 596 -191 852 -419 85 -76
266 -270 345 -371 294 -376 520 -873 640 -1409 45 -199 45 -222 3 -244 -18 -9
-45 -30 -61 -45 -25 -23 -28 -33 -23 -60 9 -44 21 -54 121 -98 91 -40 225
-124 225 -140 0 -9 -33 5 -211 89 -126 60 -159 94 -167 173 -5 50 14 89 43 89
26 0 65 49 65 81 0 52 -22 110 -48 127 -23 15 -37 14 -186 -7 -172 -25 -210
-36 -240 -69 -23 -27 -31 -99 -14 -131 16 -29 30 -37 85 -51 26 -7 44 -17 48
-29 10 -32 -15 -153 -41 -197 -27 -49 -131 -161 -139 -152 -9 9 36 77 85 126
23 24 48 60 56 79 13 31 13 38 -3 71 -10 20 -29 42 -43 48 -14 7 -37 18 -52
24 -33 15 -42 39 -63 155 -24 133 -90 393 -135 532 -119 367 -287 686 -483
919 -175 208 -434 427 -659 557 -229 131 -498 197 -804 197 -136 0 -185 -4
-313 -26 -87 -14 -31 14 137 70 186 61 289 89 437 117 120 23 376 20 543 -6z
m2913 -962 c50 -53 113 -138 198 -268 76 -116 60 -104 -41 32 -34 46 -63 81
-66 79 -2 -2 2 -30 9 -63 9 -44 9 -86 1 -171 -6 -63 -13 -116 -16 -119 -3 -4
-15 2 -26 12 -32 29 -34 12 -5 -57 50 -118 66 -160 62 -164 -2 -2 -29 34 -61
81 -31 47 -61 83 -66 80 -6 -4 -7 -24 -4 -47 7 -39 6 -40 -14 -27 -12 7 -25
10 -29 5 -17 -17 -4 -106 31 -210 20 -60 35 -111 33 -112 -2 -2 -30 43 -62 99
-60 102 -173 233 -182 210 -2 -7 26 -82 62 -168 37 -85 65 -160 63 -166 -2 -6
-43 66 -91 160 -63 121 -93 171 -106 171 -9 0 -37 -21 -61 -46 -55 -56 -56
-56 -248 20 -78 31 -156 59 -172 63 l-30 6 24 -54 c13 -30 51 -114 85 -188 34
-73 60 -135 58 -137 -2 -3 -21 24 -42 58 -56 91 -85 128 -102 128 -13 0 -14
-8 -9 -42 l7 -42 -83 80 c-109 106 -163 133 -260 126 -35 -3 -38 0 -53 34 -9
21 -13 44 -11 51 7 17 58 16 147 -2 101 -21 104 -17 92 106 -6 52 -7 98 -3
102 4 5 25 -20 47 -55 22 -35 47 -68 54 -75 18 -14 278 -83 316 -83 23 0 42
13 86 59 61 64 64 72 42 111 -19 33 -19 54 0 46 11 -4 23 6 37 30 l21 35 66
-3 66 -3 -2 29 c-1 16 -25 63 -52 105 -28 41 -49 77 -47 78 2 2 47 -42 101
-97 75 -76 105 -100 125 -100 14 0 35 -7 47 -15 20 -14 22 -14 27 2 3 10 6 72
7 138 2 102 -1 128 -19 175 -12 30 -22 56 -22 58 0 9 40 -22 71 -55z m-2642
-139 c29 -35 60 -64 67 -64 8 0 52 27 97 61 l82 60 47 -51 c59 -63 161 -218
218 -327 23 -46 48 -83 55 -83 8 0 26 5 41 11 49 18 73 4 104 -64 33 -70 48
-127 33 -127 -6 0 -32 9 -58 21 -28 12 -56 18 -71 15 -21 -6 -27 1 -55 56 -67
132 -208 340 -244 361 -10 5 -19 -1 -29 -22 -29 -55 -35 -106 -20 -183 8 -40
14 -74 14 -75 0 -1 -12 2 -26 8 l-27 10 7 -62 c6 -61 6 -62 -14 -44 -15 14
-31 17 -72 13 -58 -6 -88 -32 -88 -76 l0 -26 -29 34 c-29 35 -30 35 -122 38
-52 2 -99 8 -106 14 -15 12 -83 132 -83 146 0 6 23 39 50 73 56 70 68 99 52
134 -12 27 -15 25 81 46 65 14 68 21 42 105 -26 85 -18 85 54 -2z m-397 -294
c84 -126 196 -352 237 -475 33 -99 28 -115 -23 -66 -45 44 -55 29 -49 -77 5
-95 -4 -97 -33 -7 -22 67 -52 125 -64 125 -6 0 -10 -39 -11 -92 0 -51 -4 -101
-8 -111 -9 -23 -10 -22 -85 101 -32 50 -62 92 -68 92 -7 0 -9 -22 -4 -70 6
-74 -3 -90 -24 -40 -21 50 -33 54 -87 30 -26 -11 -57 -30 -69 -41 -20 -19 -21
-19 -75 10 -30 17 -91 41 -137 54 -46 13 -89 30 -97 38 -16 16 -45 142 -36
157 3 5 58 33 121 61 l114 51 21 -26 21 -26 49 39 c51 40 69 66 80 116 5 23
10 28 32 25 19 -2 28 -11 37 -38 37 -115 42 114 6 250 -45 171 -47 164 22 90
33 -36 92 -112 130 -170z m-1789 73 c-3 -10 -32 -77 -63 -148 -48 -109 -137
-340 -242 -628 -11 -32 -22 -56 -24 -54 -2 2 -9 32 -14 68 -24 141 -20 131
-47 124 -13 -3 -50 -18 -81 -33 -43 -20 -62 -37 -79 -67 -32 -59 -40 -65 -88
-65 -55 0 -56 6 -16 106 29 75 176 339 194 351 12 7 14 14 -47 -125 -25 -56
-46 -113 -46 -127 0 -25 0 -25 35 -11 109 46 179 120 301 316 135 217 241 360
217 293z m-869 -35 c-4 -7 -33 -49 -64 -93 -140 -200 -268 -431 -368 -665
-114 -268 -153 -314 -74 -86 41 115 44 129 29 137 -28 16 -39 80 -22 128 27
77 98 202 139 246 58 61 355 345 362 345 3 0 2 -6 -2 -12z m1415 -533 l1 -130
-23 39 c-26 47 -41 51 -45 14 -4 -36 -18 -35 -37 1 -8 17 -19 32 -25 36 -5 3
-60 -10 -122 -30 -76 -25 -130 -36 -165 -36 -29 1 -82 -5 -118 -14 -99 -23
-109 -21 -149 29 -20 23 -36 50 -36 58 0 12 55 182 75 231 11 27 38 21 147
-34 76 -38 108 -49 130 -45 21 4 28 2 28 -9 0 -10 11 -15 34 -15 45 0 153 34
177 56 10 9 35 69 55 133 56 180 58 180 65 1 4 -85 7 -213 8 -285z m-1394 272
c-31 -55 -32 -83 -2 -91 47 -12 65 -6 97 34 18 23 34 39 36 38 2 -2 -20 -48
-48 -102 l-50 -99 33 7 c84 17 149 19 149 5 0 -8 -37 -104 -82 -213 -74 -177
-83 -195 -86 -162 -4 50 -26 56 -66 17 -17 -17 -35 -31 -39 -31 -4 0 -7 18 -7
40 0 28 -6 43 -18 51 -15 10 -18 21 -14 65 8 93 -18 69 -89 -80 -35 -74 -65
-133 -67 -131 -3 2 5 35 17 74 11 38 21 74 21 80 0 6 -35 11 -87 13 l-88 3 3
30 c4 42 162 364 187 381 11 8 44 14 76 14 55 0 59 2 93 42 20 23 36 45 36 50
0 4 5 8 10 8 6 0 -1 -20 -15 -43z m1710 6 c120 -17 143 -19 164 -12 10 3 21
-17 37 -67 24 -77 75 -306 69 -312 -2 -2 -21 14 -43 37 l-39 41 -145 0 c-128
0 -148 2 -170 19 -16 13 -29 17 -39 10 -8 -5 -17 -9 -20 -9 -10 0 -66 178 -74
232 -4 25 -4 56 0 68 7 21 11 22 79 16 39 -4 121 -14 181 -23z m-79 -988 c74
-190 129 -303 247 -514 76 -136 79 -145 62 -157 -24 -17 -76 -18 -98 -1 -21
16 -126 237 -165 347 -30 87 -143 516 -141 541 1 19 17 -19 95 -216z m3644 61
c64 -163 185 -534 301 -926 66 -223 147 -490 179 -595 175 -566 250 -858 321
-1240 22 -121 43 -231 46 -245 4 -18 3 -22 -6 -15 -6 6 -22 71 -36 145 -58
313 -100 487 -200 820 -40 135 -94 317 -120 405 -161 550 -413 1387 -465 1540
-55 165 -67 205 -56 194 2 -2 18 -39 36 -83z m-4395 -812 c61 -60 129 -118
149 -128 46 -22 69 -19 238 25 70 18 130 30 133 27 3 -3 -19 -42 -49 -87 -53
-78 -75 -125 -63 -137 14 -14 111 32 205 96 57 38 136 85 176 104 84 40 299
116 327 116 12 0 37 -24 66 -65 25 -36 54 -67 62 -69 9 -2 178 -1 376 2 l360
7 10 70 c10 68 10 69 22 40 7 -16 17 -49 23 -73 14 -53 18 -55 187 -82 178
-28 191 -33 200 -78 12 -57 8 -612 -6 -742 -17 -170 -53 -395 -141 -880 -206
-1142 -248 -1540 -194 -1863 8 -49 12 -97 9 -107 -8 -24 -195 -200 -213 -200
-7 0 -21 14 -30 31 -140 256 -353 528 -467 594 -66 39 -95 39 -361 1 -137 -19
-303 -40 -368 -47 -128 -12 -307 -7 -368 11 -54 16 -140 78 -185 132 -41 51
-348 610 -438 801 -105 220 -178 478 -191 667 -7 97 11 328 25 343 5 5 23 -17
41 -49 18 -32 72 -97 125 -148 54 -53 97 -104 101 -120 12 -50 -1 -119 -34
-170 -22 -33 -32 -62 -32 -88 0 -45 31 -126 70 -182 23 -35 28 -50 23 -82 -3
-24 4 -70 17 -119 18 -68 28 -87 68 -128 59 -60 118 -101 132 -92 6 4 10 18 8
32 -3 22 3 28 48 44 28 11 59 28 70 40 l18 20 -107 -7 c-118 -8 -146 0 -166
43 -17 37 -14 52 18 84 33 33 78 41 66 12 -14 -31 -16 -77 -5 -93 9 -13 14
-11 36 14 21 24 25 37 21 68 -5 37 -4 38 36 49 56 15 192 6 242 -15 l40 -17
-27 -20 c-66 -49 -1 -50 120 -2 80 31 92 40 92 64 0 33 -30 52 -72 45 -33 -5
-51 1 -140 49 -226 121 -267 139 -327 139 -31 1 -68 -3 -83 -8 -24 -7 -32 -3
-62 30 -62 67 -63 76 -30 145 53 111 38 151 -110 308 -88 93 -131 162 -142
230 -11 67 0 84 99 164 142 113 222 232 261 384 46 178 -18 329 -188 441 -74
48 -75 49 -51 60 35 16 31 28 -20 64 -31 22 -41 34 -32 40 21 14 168 59 252
77 67 15 80 21 83 39 2 12 -33 92 -82 187 -100 193 -117 263 -35 144 28 -41
102 -124 164 -185z m-430 -364 c-13 -21 19 -51 100 -97 78 -43 142 -93 119
-93 -5 0 -34 7 -64 15 -69 19 -148 19 -180 0 -24 -14 -24 -14 21 -15 63 0 238
-37 267 -56 24 -16 42 -52 42 -85 0 -17 -8 -14 -57 24 -92 71 -138 91 -213 90
-36 0 -85 -8 -109 -17 -55 -20 -64 -14 -103 71 -37 82 -36 104 4 127 60 36
190 63 173 36z m4328 -154 c76 -37 148 -103 127 -116 -27 -17 -107 -10 -174
15 -91 34 -212 35 -310 1 -55 -19 -76 -22 -100 -14 -17 5 -37 12 -45 14 -17 6
18 45 66 75 76 47 131 59 258 57 109 -3 125 -6 178 -32z m-4293 -282 c3 -28
11 -58 17 -66 7 -8 12 -30 12 -49 -1 -40 10 -72 46 -128 33 -53 32 -67 -5 -80
-78 -27 -118 37 -133 210 -10 105 -9 120 7 144 27 42 50 29 56 -31z m-53 -438
c-4 -9 -11 -16 -17 -16 -11 0 -14 33 -3 44 11 10 26 -11 20 -28z m449 -922
c60 -20 62 -23 44 -34 -23 -15 -104 -12 -126 5 -19 15 -19 15 0 30 24 18 22
19 82 -1z m5843 -211 c82 -179 180 -472 218 -653 34 -160 38 -387 10 -517 -37
-172 -107 -345 -191 -471 -81 -122 -215 -266 -232 -249 -2 2 10 37 27 78 188
447 261 1077 185 1584 -15 98 -31 196 -35 217 -10 44 0 50 18 11z m-2689 -313
c25 -333 24 -319 22 -342 -1 -10 -66 -56 -164 -115 -175 -106 -203 -121 -196
-101 3 7 26 81 53 163 39 124 214 633 241 699 15 37 22 -12 44 -304z m110
-832 c37 -172 37 -173 -56 -257 -80 -73 -143 -103 -230 -109 -105 -7 -132 9
-183 108 -53 101 -53 132 -3 177 21 19 126 88 233 153 182 111 194 117 201 97
3 -12 20 -88 38 -169z m-1046 -650 c67 -151 141 -236 331 -383 138 -106 239
-173 309 -204 42 -18 47 -23 36 -36 -7 -9 -110 -81 -229 -160 -178 -118 -251
-161 -416 -238 -110 -51 -250 -117 -312 -147 -62 -29 -118 -50 -126 -47 -7 3
-24 34 -36 69 -28 77 -74 158 -252 445 -75 122 -140 231 -144 242 -6 19 -2 21
42 21 63 0 171 24 230 50 35 15 99 72 238 209 182 180 271 261 285 261 4 0 24
-37 44 -82z m-2387 -88 c-10 -81 -64 -284 -102 -378 -55 -136 -119 -238 -204
-323 -108 -107 -162 -132 -306 -137 -109 -4 -111 -3 -175 30 -42 22 -76 48
-98 78 l-35 45 114 7 c338 22 478 105 645 383 29 50 70 131 89 180 49 124 67
165 73 165 3 0 2 -22 -1 -50z m1647 -1402 c-54 -78 -300 -358 -315 -358 -16 0
-10 15 29 76 113 173 323 408 330 370 2 -10 -18 -49 -44 -88z"/>
<path d="M3229 5845 c-108 -15 -150 -30 -198 -71 -49 -41 -111 -123 -111 -146
0 -18 5 -20 40 -15 22 3 40 3 40 1 0 -2 -9 -26 -20 -53 -12 -32 -16 -52 -9
-56 9 -6 7 -23 -8 -62 -3 -8 4 -13 20 -13 16 0 26 -7 30 -20 8 -30 33 -24 48
12 15 37 122 148 142 148 12 0 11 -10 0 -56 -17 -66 -10 -118 17 -137 15 -11
19 -21 14 -44 -5 -25 2 -39 44 -90 28 -33 66 -67 85 -76 39 -19 112 -22 152
-7 39 15 108 74 140 121 27 38 28 41 13 72 -15 32 -15 34 8 54 13 12 24 33 24
47 l0 25 38 -17 c58 -26 111 -62 122 -81 6 -13 3 -29 -11 -57 -98 -186 -404
-264 -685 -174 -101 32 -143 37 -162 18 -21 -21 -13 -28 31 -28 49 0 82 -19
91 -53 5 -20 11 -23 31 -19 14 2 31 0 38 -6 17 -14 145 -34 168 -27 11 4 26 1
34 -5 9 -7 23 -9 37 -4 13 5 41 9 63 11 22 1 51 3 65 4 14 1 37 -1 52 -5 21
-6 27 -4 33 13 5 18 14 21 55 21 43 0 49 3 52 23 3 19 10 22 51 25 49 3 53 7
37 32 -11 17 4 30 38 30 15 0 22 6 22 19 0 14 16 27 55 45 40 18 56 31 61 51
3 14 16 32 27 40 18 13 20 18 10 34 -11 17 -8 21 25 35 100 44 116 63 55 68
-33 3 -38 6 -36 26 3 22 0 22 -45 16 -44 -6 -49 -4 -83 29 -58 56 -354 199
-469 227 -116 28 -147 43 -70 36 30 -3 108 -23 173 -45 64 -22 117 -36 117
-31 0 13 -23 24 -160 75 -142 53 -197 60 -331 40z"/>
<path d="M3027 5303 c-3 -5 -2 -15 2 -22 7 -10 10 -10 16 -1 4 6 3 16 -3 22
-5 5 -12 6 -15 1z"/>
<path d="M2180 4462 c0 -11 136 -122 149 -122 19 0 12 46 -11 75 -12 15 -36
34 -54 41 -37 15 -84 19 -84 6z"/>
</g></g></svg>

After

Width:  |  Height:  |  Size: 13 KiB

+25
View File
@@ -406,6 +406,31 @@ export async function* pullModel(
}
}
export interface ModelRecommendation {
model: string;
description: string;
context_length?: number;
max_output_tokens?: number;
vram_bytes?: number;
}
export interface ModelRecommendationsResponse {
recommendations: ModelRecommendation[];
}
export async function getModelRecommendations(): Promise<ModelRecommendation[]> {
const response = await fetch(
`${API_BASE}/api/experimental/model-recommendations`,
);
if (!response.ok) {
throw new Error(
`Failed to fetch model recommendations: ${response.statusText}`,
);
}
const data: ModelRecommendationsResponse = await response.json();
return data.recommendations || [];
}
export async function getInferenceCompute(): Promise<InferenceComputeResponse> {
const response = await fetch(`${API_BASE}/api/v1/inference-compute`);
if (!response.ok) {
+36 -11
View File
@@ -13,6 +13,22 @@ interface LaunchCommand {
}
const LAUNCH_COMMANDS: LaunchCommand[] = [
{
id: "claude-desktop",
name: "Claude Desktop",
command: "ollama launch claude-desktop",
description: "Claude Desktop with Ollama Cloud",
icon: "/launch-icons/claude.svg",
iconClassName: "h-7 w-7",
},
{
id: "claude",
name: "Claude Code",
command: "ollama launch claude",
description: "Anthropic's coding tool with subagents",
icon: "/launch-icons/claude-code.svg",
iconClassName: "h-7 w-7",
},
{
id: "openclaw",
name: "OpenClaw",
@@ -21,13 +37,21 @@ const LAUNCH_COMMANDS: LaunchCommand[] = [
icon: "/launch-icons/openclaw.svg",
},
{
id: "claude",
name: "Claude",
command: "ollama launch claude",
description: "Anthropic's coding tool with subagents",
icon: "/launch-icons/claude.svg",
id: "hermes",
name: "Hermes Agent",
command: "ollama launch hermes",
description: "Self-improving AI agent built by Nous Research",
icon: "/launch-icons/hermes-agent.svg",
iconClassName: "h-7 w-7",
},
{
id: "opencode",
name: "OpenCode",
command: "ollama launch opencode",
description: "Anomaly's open-source coding agent",
icon: "/launch-icons/opencode.svg",
iconClassName: "h-7 w-7 rounded",
},
{
id: "codex",
name: "Codex",
@@ -38,12 +62,13 @@ const LAUNCH_COMMANDS: LaunchCommand[] = [
iconClassName: "h-7 w-7",
},
{
id: "opencode",
name: "OpenCode",
command: "ollama launch opencode",
description: "Anomaly's open-source coding agent",
icon: "/launch-icons/opencode.svg",
iconClassName: "h-7 w-7 rounded",
id: "copilot",
name: "Copilot CLI",
command: "ollama launch copilot",
description: "GitHub's AI coding agent for the terminal",
icon: "/launch-icons/copilot.svg",
darkIcon: "/launch-icons/copilot-dark.svg",
iconClassName: "h-7 w-7",
},
{
id: "droid",
+5 -5
View File
@@ -381,7 +381,7 @@ export const useSendMessage = (chatId: string) => {
role: "assistant",
content: "",
thinking: "",
model: effectiveModel,
model: effectiveModel.model,
}),
);
lastMessage = newMessages[newMessages.length - 1];
@@ -433,7 +433,7 @@ export const useSendMessage = (chatId: string) => {
role: "assistant",
content: "",
thinking: "",
model: effectiveModel,
model: effectiveModel.model,
}),
);
lastMessage = newMessages[newMessages.length - 1];
@@ -520,7 +520,7 @@ export const useSendMessage = (chatId: string) => {
thinkingTimeStart:
lastMessage.thinkingTimeStart || event.thinkingTimeStart,
thinkingTimeEnd: event.thinkingTimeEnd,
model: selectedModel,
model: selectedModel.model,
});
newMessages[newMessages.length - 1] = updatedMessage;
} else {
@@ -533,7 +533,7 @@ export const useSendMessage = (chatId: string) => {
tool_calls: event.toolCalls,
thinkingTimeStart: event.thinkingTimeStart,
thinkingTimeEnd: event.thinkingTimeEnd,
model: selectedModel,
model: selectedModel.model,
}),
);
}
@@ -699,7 +699,7 @@ export const useSendMessage = (chatId: string) => {
queryClient.setQueryData(["chat", newId], {
chat: new Chat({
id: newId,
model: effectiveModel,
model: effectiveModel.model,
messages: [
new Message({
role: "user",
+13
View File
@@ -0,0 +1,13 @@
import { useQuery } from "@tanstack/react-query";
import { getModelRecommendations } from "@/api";
import type { ModelRecommendation } from "@/api";
export function useFeaturedModels() {
return useQuery<ModelRecommendation[], Error>({
queryKey: ["modelRecommendations"],
queryFn: getModelRecommendations,
staleTime: 5 * 60 * 1000,
gcTime: 30 * 60 * 1000,
refetchOnWindowFocus: false,
});
}
+35 -24
View File
@@ -1,51 +1,49 @@
import { useQuery } from "@tanstack/react-query";
import { Model } from "@/gotypes";
import { getModels } from "@/api";
import { mergeModels } from "@/utils/mergeModels";
import { useMemo } from "react";
import { useCloudStatus } from "./useCloudStatus";
import { useFeaturedModels } from "./useFeaturedModels";
export function useModels(searchQuery = "") {
const { cloudDisabled } = useCloudStatus();
const { data: recommendations, isLoading: recommendationsLoading } =
useFeaturedModels();
const localQuery = useQuery<Model[], Error>({
queryKey: ["models", searchQuery],
queryFn: () => getModels(searchQuery),
gcTime: 10 * 60 * 1000, // Keep in cache for 10 minutes
gcTime: 10 * 60 * 1000,
retry: 10,
// exponential backoff, starting at 100ms and capping at 5s
retryDelay: (attemptIndex) => Math.min(100 * 2 ** attemptIndex, 5000),
refetchOnWindowFocus: true,
refetchInterval: 30 * 1000, // Refetch every 30 seconds to keep models updated
refetchInterval: 30 * 1000,
refetchIntervalInBackground: true,
});
const allModels = useMemo(() => {
const models = mergeModels(localQuery.data || [], cloudDisabled);
const local = localQuery.data || [];
const featured = (recommendations || []).map((r) => r.model);
const featuredSet = new Set(featured);
if (searchQuery && searchQuery.trim()) {
const query = searchQuery.toLowerCase().trim();
const filteredModels = models.filter((model) =>
model.model.toLowerCase().includes(query),
);
// Recommended models first (using the local copy when downloaded),
// then everything else from /api/tags in tags order.
const recommended = featured.map(
(name) =>
local.find((m) => m.model === name) || new Model({ model: name }),
);
const rest = local.filter((m) => !featuredSet.has(m.model));
const merged = [...recommended, ...rest];
const seen = new Set<string>();
return filteredModels.filter((model) => {
const currentModel = model.model.toLowerCase();
if (seen.has(currentModel)) {
return false;
}
seen.add(currentModel);
return true;
});
}
return models;
}, [localQuery.data, searchQuery, cloudDisabled]);
const visible = cloudDisabled
? merged.filter((m) => !m.isCloud())
: merged;
return filterBySearch(visible, searchQuery);
}, [localQuery.data, searchQuery, cloudDisabled, recommendations]);
return {
...localQuery,
data: allModels,
isLoading: localQuery.isLoading,
isLoading: localQuery.isLoading || recommendationsLoading,
};
}
@@ -53,3 +51,16 @@ export function useRefetchModels() {
const { refetch } = useModels();
return refetch;
}
function filterBySearch(models: Model[], query: string): Model[] {
const q = query.trim().toLowerCase();
if (!q) return models;
const seen = new Set<string>();
return models.filter((m) => {
const name = m.model.toLowerCase();
if (!name.includes(q) || seen.has(name)) return false;
seen.add(name);
return true;
});
}
+1 -4
View File
@@ -4,7 +4,6 @@ import { useModels } from "./useModels";
import { useChat } from "./useChats";
import { useSettings } from "./useSettings.ts";
import { Model } from "@/gotypes";
import { FEATURED_MODELS } from "@/utils/mergeModels";
import { getTotalVRAM } from "@/utils/vram.ts";
import { getInferenceCompute } from "@/api";
import { useCloudStatus } from "./useCloudStatus";
@@ -92,9 +91,7 @@ export function useSelectedModel(currentChatId?: string, searchQuery?: string) {
(settings.selectedModel &&
new Model({
model: settings.selectedModel,
cloud: FEATURED_MODELS.some(
(f) => f.endsWith("cloud") && f === settings.selectedModel,
),
cloud: settings.selectedModel.endsWith("cloud"),
ollama_host: false,
})) ||
null
-128
View File
@@ -1,128 +0,0 @@
import { describe, it, expect } from "vitest";
import { Model } from "@/gotypes";
import { mergeModels, FEATURED_MODELS } from "@/utils/mergeModels";
import "@/api";
describe("Model merging logic", () => {
it("should handle cloud models with -cloud suffix", () => {
const localModels: Model[] = [
new Model({ model: "gpt-oss:120b-cloud" }),
new Model({ model: "llama3:latest" }),
new Model({ model: "mistral:latest" }),
];
const merged = mergeModels(localModels);
// First verify cloud models are first and in FEATURED_MODELS order
const cloudModels = FEATURED_MODELS.filter((m: string) =>
m.endsWith("cloud"),
);
for (let i = 0; i < cloudModels.length; i++) {
expect(merged[i].model).toBe(cloudModels[i]);
expect(merged[i].isCloud()).toBe(true);
}
// Then verify non-cloud featured models are next and in FEATURED_MODELS order
const nonCloudFeatured = FEATURED_MODELS.filter(
(m: string) => !m.endsWith("cloud"),
);
for (let i = 0; i < nonCloudFeatured.length; i++) {
const model = merged[i + cloudModels.length];
expect(model.model).toBe(nonCloudFeatured[i]);
expect(model.isCloud()).toBe(false);
}
// Verify local models are preserved and come after featured models
const featuredCount = FEATURED_MODELS.length;
expect(merged[featuredCount].model).toBe("llama3:latest");
expect(merged[featuredCount + 1].model).toBe("mistral:latest");
// Length should be exactly featured models plus our local models
expect(merged.length).toBe(FEATURED_MODELS.length + 2);
});
it("should hide cloud models when cloud is disabled", () => {
const localModels: Model[] = [
new Model({ model: "gpt-oss:120b-cloud" }),
new Model({ model: "llama3:latest" }),
new Model({ model: "mistral:latest" }),
];
const merged = mergeModels(localModels, true); // cloud disabled = true
// No cloud models should be present
const cloudModels = merged.filter((m) => m.isCloud());
expect(cloudModels.length).toBe(0);
// Should have non-cloud featured models
const nonCloudFeatured = FEATURED_MODELS.filter(
(m) => !m.endsWith("cloud"),
);
for (let i = 0; i < nonCloudFeatured.length; i++) {
const model = merged[i];
expect(model.model).toBe(nonCloudFeatured[i]);
expect(model.isCloud()).toBe(false);
}
// Local models should be preserved
const featuredCount = nonCloudFeatured.length;
expect(merged[featuredCount].model).toBe("llama3:latest");
expect(merged[featuredCount + 1].model).toBe("mistral:latest");
});
it("should handle empty input", () => {
const merged = mergeModels([]);
// First verify cloud models are first and in FEATURED_MODELS order
const cloudModels = FEATURED_MODELS.filter((m) => m.endsWith("cloud"));
for (let i = 0; i < cloudModels.length; i++) {
expect(merged[i].model).toBe(cloudModels[i]);
expect(merged[i].isCloud()).toBe(true);
}
// Then verify non-cloud featured models are next and in FEATURED_MODELS order
const nonCloudFeatured = FEATURED_MODELS.filter(
(m) => !m.endsWith("cloud"),
);
for (let i = 0; i < nonCloudFeatured.length; i++) {
const model = merged[i + cloudModels.length];
expect(model.model).toBe(nonCloudFeatured[i]);
expect(model.isCloud()).toBe(false);
}
// Length should be exactly FEATURED_MODELS length
expect(merged.length).toBe(FEATURED_MODELS.length);
});
it("should sort models correctly", () => {
const localModels: Model[] = [
new Model({ model: "zephyr:latest" }),
new Model({ model: "alpha:latest" }),
new Model({ model: "gpt-oss:120b-cloud" }),
];
const merged = mergeModels(localModels);
// First verify cloud models are first and in FEATURED_MODELS order
const cloudModels = FEATURED_MODELS.filter((m) => m.endsWith("cloud"));
for (let i = 0; i < cloudModels.length; i++) {
expect(merged[i].model).toBe(cloudModels[i]);
expect(merged[i].isCloud()).toBe(true);
}
// Then verify non-cloud featured models are next and in FEATURED_MODELS order
const nonCloudFeatured = FEATURED_MODELS.filter(
(m) => !m.endsWith("cloud"),
);
for (let i = 0; i < nonCloudFeatured.length; i++) {
const model = merged[i + cloudModels.length];
expect(model.model).toBe(nonCloudFeatured[i]);
expect(model.isCloud()).toBe(false);
}
// Non-featured local models should be at the end in alphabetical order
const featuredCount = FEATURED_MODELS.length;
expect(merged[featuredCount].model).toBe("alpha:latest");
expect(merged[featuredCount + 1].model).toBe("zephyr:latest");
});
});
-102
View File
@@ -1,102 +0,0 @@
import { Model } from "@/gotypes";
// Featured models list (in priority order)
export const FEATURED_MODELS = [
"kimi-k2.5:cloud",
"glm-5:cloud",
"minimax-m2.7:cloud",
"gemma4:31b-cloud",
"qwen3.5:397b-cloud",
"gpt-oss:120b-cloud",
"gpt-oss:20b-cloud",
"deepseek-v3.1:671b-cloud",
"gpt-oss:120b",
"gpt-oss:20b",
"gemma4:31b",
"gemma4:26b",
"gemma4:e4b",
"gemma4:e2b",
"deepseek-r1:8b",
"qwen3-coder:30b",
"qwen3-vl:30b",
"qwen3-vl:8b",
"qwen3-vl:4b",
"qwen3.5:27b",
"qwen3.5:9b",
"qwen3.5:4b",
];
function alphabeticalSort(a: Model, b: Model): number {
return a.model.toLowerCase().localeCompare(b.model.toLowerCase());
}
//Merges models, sorting cloud models first, then other models
export function mergeModels(
localModels: Model[],
hideCloudModels: boolean = false,
): Model[] {
const allModels = (localModels || []).map((model) => model);
// 1. Get cloud models from local models and featured list
const cloudModels = [...allModels.filter((m) => m.isCloud())];
// Add any cloud models from FEATURED_MODELS that aren't in local models
FEATURED_MODELS.filter((f) => f.endsWith("cloud")).forEach((cloudModel) => {
if (!cloudModels.some((m) => m.model === cloudModel)) {
cloudModels.push(new Model({ model: cloudModel }));
}
});
// 2. Get other featured models (non-cloud)
const featuredModels = FEATURED_MODELS.filter(
(f) => !f.endsWith("cloud"),
).map((model) => {
// Check if this model exists in local models
const localMatch = allModels.find(
(m) => m.model.toLowerCase() === model.toLowerCase(),
);
if (localMatch) return localMatch;
return new Model({
model,
});
});
// 3. Get remaining local models that aren't featured and aren't cloud models
const remainingModels = allModels.filter(
(model) =>
!model.isCloud() &&
!FEATURED_MODELS.some(
(f) => f.toLowerCase() === model.model.toLowerCase(),
),
);
cloudModels.sort((a, b) => {
const aIndex = FEATURED_MODELS.indexOf(a.model);
const bIndex = FEATURED_MODELS.indexOf(b.model);
// If both are featured, sort by their position in FEATURED_MODELS
if (aIndex !== -1 && bIndex !== -1) {
return aIndex - bIndex;
}
// If only one is featured, featured model comes first
if (aIndex !== -1 && bIndex === -1) return -1;
if (aIndex === -1 && bIndex !== -1) return 1;
// If neither is featured, sort alphabetically
return a.model.toLowerCase().localeCompare(b.model.toLowerCase());
});
featuredModels.sort(
(a, b) =>
FEATURED_MODELS.indexOf(a.model) - FEATURED_MODELS.indexOf(b.model),
);
remainingModels.sort(alphabeticalSort);
return hideCloudModels
? [...featuredModels, ...remainingModels]
: [...cloudModels, ...featuredModels, ...remainingModels];
}
+1
View File
@@ -302,6 +302,7 @@ func (s *Server) Handler() http.Handler {
mux.Handle("HEAD /api/version", ollamaProxy)
mux.Handle("POST /api/me", ollamaProxy)
mux.Handle("POST /api/signout", ollamaProxy)
mux.Handle("GET /api/experimental/model-recommendations", ollamaProxy)
// React app - catch all non-API routes and serve the React app
mux.Handle("GET /", s.appHandler())
+67 -6
View File
@@ -582,10 +582,10 @@ func RunHandler(cmd *cobra.Command, args []string) error {
opts.Think = &api.ThinkValue{Value: true}
case "false":
opts.Think = &api.ThinkValue{Value: false}
case "high", "medium", "low":
case "high", "medium", "low", "max":
opts.Think = &api.ThinkValue{Value: thinkStr}
default:
return fmt.Errorf("invalid value for --think: %q (must be true, false, high, medium, or low)", thinkStr)
return fmt.Errorf("invalid value for --think: %q (must be true, false, high, medium, low, or max)", thinkStr)
}
} else {
opts.Think = nil
@@ -1975,8 +1975,61 @@ func launchInteractiveModel(cmd *cobra.Command, modelName string) error {
Options: map[string]any{},
ShowConnect: true,
}
// loadOrUnloadModel is cloud-safe here: remote/cloud models skip local preload
// and only validate auth/connectivity before interactive chat starts.
client, err := api.ClientFromEnvironment()
if err != nil {
return err
}
requestedCloud := modelref.HasExplicitCloudSource(modelName)
info, err := func() (*api.ShowResponse, error) {
showReq := &api.ShowRequest{Name: modelName}
info, err := client.Show(cmd.Context(), showReq)
var se api.StatusError
if errors.As(err, &se) && se.StatusCode == http.StatusNotFound {
if requestedCloud {
return nil, err
}
if err := PullHandler(cmd, []string{modelName}); err != nil {
return nil, err
}
return client.Show(cmd.Context(), &api.ShowRequest{Name: modelName})
}
return info, err
}()
if err != nil {
if handleCloudAuthorizationError(err) {
return nil
}
return err
}
ensureCloudStub(cmd.Context(), client, modelName)
opts.Think, err = inferThinkingOption(&info.Capabilities, &opts, false)
if err != nil {
return err
}
audioCapable := slices.Contains(info.Capabilities, model.CapabilityAudio)
opts.MultiModal = slices.Contains(info.Capabilities, model.CapabilityVision) || audioCapable
// TODO: remove the projector info and vision info checks below,
// these are left in for backwards compatibility with older servers
// that don't have the capabilities field in the model info
if len(info.ProjectorInfo) != 0 {
opts.MultiModal = true
}
for k := range info.ModelInfo {
if strings.Contains(k, ".vision.") {
opts.MultiModal = true
break
}
}
applyShowResponseToRunOptions(&opts, info)
if err := loadOrUnloadModel(cmd, &opts); err != nil {
return fmt.Errorf("error loading model: %w", err)
}
@@ -2066,8 +2119,7 @@ func runLauncherAction(cmd *cobra.Command, action tui.TUIAction, deps launcherDe
if err != nil {
return true, fmt.Errorf("launching %s: %w", action.Integration, err)
}
// VS Code is a GUI app — exit the TUI loop after launching
if action.Integration == "vscode" {
if launcherActionExitsLoop(action.Integration) {
return false, nil
}
return true, nil
@@ -2076,6 +2128,15 @@ func runLauncherAction(cmd *cobra.Command, action tui.TUIAction, deps launcherDe
}
}
func launcherActionExitsLoop(integration string) bool {
switch integration {
case "claude-desktop", "vscode":
return true
default:
return false
}
}
func NewCLI() *cobra.Command {
log.SetFlags(log.LstdFlags | log.Lshortfile)
cobra.EnableCommandSorting = false
+16 -15
View File
@@ -209,29 +209,30 @@ func TestRunLauncherAction_RunModelContinuesAfterCancellation(t *testing.T) {
}
}
func TestRunLauncherAction_VSCodeExitsTUILoop(t *testing.T) {
func TestRunLauncherAction_GUIAppsExitTUILoop(t *testing.T) {
setCmdTestHome(t, t.TempDir())
cmd := &cobra.Command{}
cmd.SetContext(context.Background())
// VS Code should exit the TUI loop (return false) after a successful launch.
continueLoop, err := runLauncherAction(cmd, tui.TUIAction{Kind: tui.TUIActionLaunchIntegration, Integration: "vscode"}, launcherDeps{
resolveRunModel: unexpectedRunModelResolution(t),
launchIntegration: func(ctx context.Context, req launch.IntegrationLaunchRequest) error {
return nil
},
runModel: unexpectedModelLaunch(t),
})
if err != nil {
t.Fatalf("expected nil error, got %v", err)
}
if continueLoop {
t.Fatal("expected vscode launch to exit the TUI loop (return false)")
for _, integration := range []string{"claude-desktop", "vscode"} {
continueLoop, err := runLauncherAction(cmd, tui.TUIAction{Kind: tui.TUIActionLaunchIntegration, Integration: integration}, launcherDeps{
resolveRunModel: unexpectedRunModelResolution(t),
launchIntegration: func(ctx context.Context, req launch.IntegrationLaunchRequest) error {
return nil
},
runModel: unexpectedModelLaunch(t),
})
if err != nil {
t.Fatalf("expected nil error for %s, got %v", integration, err)
}
if continueLoop {
t.Fatalf("expected %s launch to exit the TUI loop (return false)", integration)
}
}
// Other integrations should continue the TUI loop (return true).
continueLoop, err = runLauncherAction(cmd, tui.TUIAction{Kind: tui.TUIActionLaunchIntegration, Integration: "claude"}, launcherDeps{
continueLoop, err := runLauncherAction(cmd, tui.TUIAction{Kind: tui.TUIActionLaunchIntegration, Integration: "claude"}, launcherDeps{
resolveRunModel: unexpectedRunModelResolution(t),
launchIntegration: func(ctx context.Context, req launch.IntegrationLaunchRequest) error {
return nil
+889
View File
@@ -0,0 +1,889 @@
package launch
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"os"
"os/exec"
"path/filepath"
"runtime"
"strings"
"time"
"github.com/ollama/ollama/cmd/config"
"github.com/ollama/ollama/cmd/internal/fileutil"
"golang.org/x/term"
)
const (
claudeDesktopIntegrationName = "claude-desktop"
claudeDesktopProfileName = "Ollama"
claudeDesktopProfileID = "00000000-0000-4000-8000-000000000114"
claudeDesktopGatewayBaseURL = "https://ollama.com"
claudeDesktopAPIKeyURL = "https://ollama.com/settings/keys"
claudeDesktopModelLabel = "Ollama Cloud"
claudeDesktopSuccessMessage = "Claude Desktop profile changed to Ollama Cloud."
claudeDesktopRestoreMessage = "To restore the usual Claude profile, run: ollama launch claude-desktop --restore"
claudeDesktopRestoredMessage = "Claude Desktop restored to the usual Claude profile."
)
var (
claudeDesktopGOOS = runtime.GOOS
claudeDesktopUserHome = os.UserHomeDir
claudeDesktopStat = os.Stat
claudeDesktopOpenApp = defaultClaudeDesktopOpenApp
claudeDesktopOpenAppPath = defaultClaudeDesktopOpenAppPath
claudeDesktopQuitApp = defaultClaudeDesktopQuitApp
claudeDesktopIsRunning = defaultClaudeDesktopIsRunning
claudeDesktopRunningAppPath = defaultClaudeDesktopRunningAppPath
claudeDesktopGlob = filepath.Glob
claudeDesktopSleep = time.Sleep
claudeDesktopHTTPClient = http.DefaultClient
claudeDesktopPromptAPIKey = promptClaudeDesktopAPIKey
claudeDesktopValidateAPIKey = validateClaudeDesktopAPIKey
)
// ClaudeDesktop configures and launches Claude Desktop in third-party
// inference mode using Ollama Cloud as the gateway.
type ClaudeDesktop struct{}
func (c *ClaudeDesktop) String() string { return "Claude Desktop" }
func (c *ClaudeDesktop) Supported() error { return claudeDesktopSupported() }
func (c *ClaudeDesktop) Paths() []string {
return nil
}
func (c *ClaudeDesktop) AutodiscoveredModel() string {
return claudeDesktopModelLabel
}
func (c *ClaudeDesktop) ConfigureAutodiscovery() error {
if err := claudeDesktopSupported(); err != nil {
return err
}
targets, err := claudeDesktopTargetPaths()
if err != nil {
return err
}
key, err := claudeDesktopValidatedAPIKey(context.Background(), claudeDesktopTargetProfilePaths(targets))
if err != nil {
return err
}
for _, path := range targets.normalConfigs {
if err := writeClaudeDesktopDeploymentMode(path, "3p"); err != nil {
return err
}
}
for _, target := range targets.thirdPartyProfiles {
if err := writeClaudeDesktopDeploymentMode(target.desktopConfig, "3p"); err != nil {
return err
}
if err := writeClaudeDesktopMeta(target.meta, claudeDesktopProfileID, claudeDesktopProfileName); err != nil {
return err
}
if err := writeClaudeDesktopGatewayProfile(target.profile, key, true); err != nil {
return err
}
}
return nil
}
func (c *ClaudeDesktop) RestoreHint() string {
return claudeDesktopRestoreMessage
}
func (c *ClaudeDesktop) ConfigurationSuccessMessage() string {
return claudeDesktopSuccessMessage + "\n" + claudeDesktopRestoreMessage
}
func (c *ClaudeDesktop) RestoreSuccessMessage() string {
return claudeDesktopRestoredMessage
}
func (c *ClaudeDesktop) AutodiscoveryConfigured() bool {
targets, err := claudeDesktopTargetPaths()
if err != nil {
return false
}
return claudeDesktopTargetsConfigured(targets)
}
func (c *ClaudeDesktop) Onboard() error {
return config.MarkIntegrationOnboarded(claudeDesktopIntegrationName)
}
func (c *ClaudeDesktop) RequiresInteractiveOnboarding() bool {
return false
}
func (c *ClaudeDesktop) SkipModelReadiness() bool {
return true
}
func (c *ClaudeDesktop) Run(_ string, args []string) error {
if err := claudeDesktopSupported(); err != nil {
return err
}
if len(args) > 0 {
return fmt.Errorf("claude-desktop does not accept extra arguments")
}
return claudeDesktopLaunchOrRestart("Restart Claude Desktop to use Ollama?")
}
func (c *ClaudeDesktop) Restore() error {
if err := claudeDesktopSupported(); err != nil {
return err
}
targets, err := claudeDesktopTargetPaths()
if err != nil {
return err
}
for _, path := range targets.normalConfigs {
if err := writeClaudeDesktopDeploymentMode(path, "1p"); err != nil {
return err
}
}
for _, target := range targets.thirdPartyProfiles {
if err := writeClaudeDesktopDeploymentMode(target.desktopConfig, "1p"); err != nil {
return err
}
if err := restoreClaudeDesktopMeta(target.meta); err != nil {
return err
}
if err := restoreClaudeDesktopOllamaProfile(target.profile); err != nil {
return err
}
}
return claudeDesktopLaunchOrRestart("Restart Claude Desktop to use the usual Claude profile?")
}
func claudeDesktopSupported() error {
switch claudeDesktopGOOS {
case "darwin", "windows":
return nil
default:
return fmt.Errorf("Claude Desktop launch is only supported on macOS and Windows")
}
}
func claudeDesktopInstalled() bool {
if claudeDesktopAppPath() != "" {
return true
}
if claudeDesktopGOOS == "windows" && claudeDesktopIsRunning() {
return true
}
for _, dir := range claudeDesktopProfileDirCandidates(false) {
if _, err := claudeDesktopStat(dir); err == nil {
return true
}
}
return false
}
func claudeDesktopAppPath() string {
if claudeDesktopGOOS != "darwin" && claudeDesktopGOOS != "windows" {
return ""
}
for _, path := range claudeDesktopAppCandidates() {
if _, err := claudeDesktopStat(path); err == nil {
return path
}
}
return ""
}
func claudeDesktopAppCandidates() []string {
switch claudeDesktopGOOS {
case "darwin":
return claudeDesktopDarwinAppCandidates()
case "windows":
return claudeDesktopWindowsAppCandidates()
default:
return nil
}
}
func claudeDesktopDarwinAppCandidates() []string {
candidates := []string{"/Applications/Claude.app"}
if home, err := claudeDesktopUserHome(); err == nil {
candidates = append(candidates, filepath.Join(home, "Applications", "Claude.app"))
}
return candidates
}
func claudeDesktopWindowsAppCandidates() []string {
local, err := claudeDesktopLocalAppData()
if err != nil {
return nil
}
candidates := []string{
filepath.Join(local, "Programs", "Claude", "Claude.exe"),
filepath.Join(local, "Programs", "Claude Desktop", "Claude.exe"),
filepath.Join(local, "Claude", "Claude.exe"),
filepath.Join(local, "Claude Nest", "Claude.exe"),
filepath.Join(local, "Claude Desktop", "Claude.exe"),
filepath.Join(local, "AnthropicClaude", "Claude.exe"),
}
for _, pattern := range []string{
filepath.Join(local, "AnthropicClaude", "app-*", "Claude.exe"),
filepath.Join(local, "Programs", "Claude", "app-*", "Claude.exe"),
filepath.Join(local, "Programs", "Claude Desktop", "app-*", "Claude.exe"),
} {
matches, _ := claudeDesktopGlob(pattern)
candidates = append(candidates, matches...)
}
return claudeDesktopDedupePaths(candidates)
}
func claudeDesktopDedupePaths(paths []string) []string {
out := make([]string, 0, len(paths))
seen := make(map[string]bool, len(paths))
for _, path := range paths {
if strings.TrimSpace(path) == "" {
continue
}
key := strings.ToLower(path)
if seen[key] {
continue
}
seen[key] = true
out = append(out, path)
}
return out
}
type claudeDesktopPaths struct {
normalConfig string
desktopConfig string
meta string
profile string
}
type claudeDesktopThirdPartyPaths struct {
desktopConfig string
meta string
profile string
}
type claudeDesktopTargets struct {
normalConfigs []string
thirdPartyProfiles []claudeDesktopThirdPartyPaths
}
func claudeDesktopConfigPaths() (claudeDesktopPaths, error) {
switch claudeDesktopGOOS {
case "darwin":
return claudeDesktopDarwinConfigPaths()
case "windows":
return claudeDesktopWindowsConfigPaths()
default:
return claudeDesktopPaths{}, claudeDesktopSupported()
}
}
func claudeDesktopDarwinConfigPaths() (claudeDesktopPaths, error) {
normalRoots, thirdPartyRoots, err := claudeDesktopDarwinProfileRoots()
if err != nil {
return claudeDesktopPaths{}, err
}
normalBase := normalRoots[0]
thirdPartyBase := thirdPartyRoots[0]
return claudeDesktopPaths{
normalConfig: filepath.Join(normalBase, "claude_desktop_config.json"),
desktopConfig: filepath.Join(thirdPartyBase, "claude_desktop_config.json"),
meta: filepath.Join(thirdPartyBase, "configLibrary", "_meta.json"),
profile: filepath.Join(thirdPartyBase, "configLibrary", claudeDesktopProfileID+".json"),
}, nil
}
func claudeDesktopWindowsConfigPaths() (claudeDesktopPaths, error) {
normalBase, err := claudeDesktopProfileDir(true)
if err != nil {
return claudeDesktopPaths{}, err
}
thirdPartyBase, err := claudeDesktopProfileDir(false)
if err != nil {
return claudeDesktopPaths{}, err
}
return claudeDesktopPaths{
normalConfig: filepath.Join(normalBase, "claude_desktop_config.json"),
desktopConfig: filepath.Join(thirdPartyBase, "claude_desktop_config.json"),
meta: filepath.Join(thirdPartyBase, "configLibrary", "_meta.json"),
profile: filepath.Join(thirdPartyBase, "configLibrary", claudeDesktopProfileID+".json"),
}, nil
}
func claudeDesktopProfileDir(normal bool) (string, error) {
candidates := claudeDesktopProfileDirCandidates(normal)
if len(candidates) == 0 {
return "", fmt.Errorf("Claude Desktop profile directory could not be resolved")
}
for _, candidate := range candidates {
if _, err := claudeDesktopStat(candidate); err == nil {
return candidate, nil
}
}
return candidates[0], nil
}
func claudeDesktopProfileDirCandidates(normal bool) []string {
if claudeDesktopGOOS != "windows" {
return nil
}
normalRoots, thirdPartyRoots, err := claudeDesktopWindowsProfileRoots()
if err != nil {
return nil
}
if normal {
return normalRoots
}
return thirdPartyRoots
}
func claudeDesktopDarwinProfileRoots() ([]string, []string, error) {
home, err := claudeDesktopUserHome()
if err != nil {
return nil, nil, err
}
base := filepath.Join(home, "Library", "Application Support")
return []string{filepath.Join(base, "Claude")}, []string{filepath.Join(base, "Claude-3p")}, nil
}
func claudeDesktopWindowsProfileRoots() ([]string, []string, error) {
local, err := claudeDesktopLocalAppData()
if err != nil {
return nil, nil, err
}
normalRoots := []string{
filepath.Join(local, "Claude"),
filepath.Join(local, "Claude Nest"),
}
thirdPartyRoots := []string{
filepath.Join(local, "Claude-3p"),
filepath.Join(local, "Claude Nest-3p"),
}
return normalRoots, thirdPartyRoots, nil
}
func claudeDesktopTargetPaths() (claudeDesktopTargets, error) {
var (
normalRoots []string
thirdPartyRoots []string
err error
)
switch claudeDesktopGOOS {
case "darwin":
normalRoots, thirdPartyRoots, err = claudeDesktopDarwinProfileRoots()
case "windows":
normalRoots, thirdPartyRoots, err = claudeDesktopWindowsProfileRoots()
default:
err = claudeDesktopSupported()
}
if err != nil {
return claudeDesktopTargets{}, err
}
return newClaudeDesktopTargets(normalRoots, thirdPartyRoots), nil
}
func newClaudeDesktopTargets(normalRoots, thirdPartyRoots []string) claudeDesktopTargets {
targets := claudeDesktopTargets{}
for _, root := range claudeDesktopDedupePaths(normalRoots) {
targets.normalConfigs = append(targets.normalConfigs, filepath.Join(root, "claude_desktop_config.json"))
}
for _, root := range claudeDesktopDedupePaths(thirdPartyRoots) {
targets.thirdPartyProfiles = append(targets.thirdPartyProfiles, claudeDesktopThirdPartyPaths{
desktopConfig: filepath.Join(root, "claude_desktop_config.json"),
meta: filepath.Join(root, "configLibrary", "_meta.json"),
profile: filepath.Join(root, "configLibrary", claudeDesktopProfileID+".json"),
})
}
return targets
}
func claudeDesktopTargetProfilePaths(targets claudeDesktopTargets) []string {
paths := make([]string, 0, len(targets.thirdPartyProfiles))
for _, target := range targets.thirdPartyProfiles {
paths = append(paths, target.profile)
}
return paths
}
func claudeDesktopLocalAppData() (string, error) {
if local := strings.TrimSpace(os.Getenv("LOCALAPPDATA")); local != "" {
return local, nil
}
if home := strings.TrimSpace(os.Getenv("USERPROFILE")); home != "" {
return filepath.Join(home, "AppData", "Local"), nil
}
home, err := claudeDesktopUserHome()
if err != nil {
return "", err
}
return filepath.Join(home, "AppData", "Local"), nil
}
type claudeDesktopAPIKeySource int
const (
claudeDesktopAPIKeySourceNone claudeDesktopAPIKeySource = iota
claudeDesktopAPIKeySourceEnv
claudeDesktopAPIKeySourceProfile
)
func claudeDesktopValidatedAPIKey(ctx context.Context, profilePaths []string) (string, error) {
key, source, err := claudeDesktopAPIKey(profilePaths)
if err != nil {
return "", err
}
if err := claudeDesktopValidateAPIKey(ctx, key); err == nil {
return key, nil
} else if source != claudeDesktopAPIKeySourceProfile || !canPromptClaudeDesktopAPIKey() {
return "", err
}
return promptValidClaudeDesktopAPIKey(ctx)
}
func claudeDesktopAPIKey(profilePaths []string) (string, claudeDesktopAPIKeySource, error) {
if key := strings.TrimSpace(os.Getenv("OLLAMA_API_KEY")); key != "" {
return key, claudeDesktopAPIKeySourceEnv, nil
}
for _, profilePath := range profilePaths {
if key := readClaudeDesktopGatewayAPIKey(profilePath); key != "" {
return key, claudeDesktopAPIKeySourceProfile, nil
}
}
key, err := promptClaudeDesktopAPIKeyValue()
return key, claudeDesktopAPIKeySourceNone, err
}
func canPromptClaudeDesktopAPIKey() bool {
return isInteractiveSession() && !currentLaunchConfirmPolicy.requireYesMessage
}
func promptValidClaudeDesktopAPIKey(ctx context.Context) (string, error) {
key, err := promptClaudeDesktopAPIKeyValue()
if err != nil {
return "", err
}
if err := claudeDesktopValidateAPIKey(ctx, key); err != nil {
return "", err
}
return key, nil
}
func promptClaudeDesktopAPIKeyValue() (string, error) {
if !canPromptClaudeDesktopAPIKey() {
return "", missingClaudeDesktopAPIKeyError()
}
key, err := claudeDesktopPromptAPIKey()
if err != nil {
return "", err
}
key = strings.TrimSpace(key)
if key == "" {
return "", missingClaudeDesktopAPIKeyError()
}
return key, nil
}
func missingClaudeDesktopAPIKeyError() error {
return fmt.Errorf("OLLAMA_API_KEY is required for Claude Desktop. Create an API key at %s, then re-run with OLLAMA_API_KEY set", claudeDesktopAPIKeyURL)
}
func promptClaudeDesktopAPIKey() (string, error) {
fmt.Fprint(os.Stderr, claudeDesktopAPIKeyPrompt())
key, err := term.ReadPassword(int(os.Stdin.Fd()))
fmt.Fprintln(os.Stderr)
if err != nil {
return "", err
}
return string(key), nil
}
func claudeDesktopAPIKeyPrompt() string {
return fmt.Sprintf("Create an Ollama API key at %s\nEnter Ollama API key (input hidden): ", claudeDesktopAPIKeyURL)
}
func readClaudeDesktopGatewayAPIKey(path string) string {
cfg, err := readClaudeDesktopJSON(path)
if err != nil {
return ""
}
key, _ := cfg["inferenceGatewayApiKey"].(string)
return strings.TrimSpace(key)
}
func validateClaudeDesktopAPIKey(ctx context.Context, key string) error {
ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
defer cancel()
if claudeDesktopAPIKeyHasInvalidHeaderChars(key) {
return claudeDesktopAPIKeyVerificationError()
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, claudeDesktopGatewayBaseURL+"/v1/models", nil)
if err != nil {
return claudeDesktopAPIKeyVerificationError()
}
req.Header.Set("Authorization", "Bearer "+key)
req.Header.Set("Accept", "application/json")
resp, err := claudeDesktopHTTPClient.Do(req)
if err != nil {
return claudeDesktopAPIKeyVerificationError()
}
defer resp.Body.Close()
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4<<10))
switch {
case resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden:
return fmt.Errorf("Ollama API key was rejected; create a valid key at %s", claudeDesktopAPIKeyURL)
case resp.StatusCode >= 200 && resp.StatusCode < 300:
return nil
default:
return fmt.Errorf("could not verify Ollama API key; ollama.com returned status %d, try again later", resp.StatusCode)
}
}
func claudeDesktopAPIKeyHasInvalidHeaderChars(key string) bool {
return strings.ContainsFunc(key, func(r rune) bool {
return r < ' ' || r == 0x7f
})
}
func claudeDesktopAPIKeyVerificationError() error {
return fmt.Errorf("could not verify Ollama API key; copy a key from %s and try again", claudeDesktopAPIKeyURL)
}
func writeClaudeDesktopDeploymentMode(path, mode string) error {
cfg, err := readClaudeDesktopJSONAllowMissing(path)
if err != nil {
return fmt.Errorf("parse Claude Desktop config: %w", err)
}
cfg["deploymentMode"] = mode
return writeClaudeDesktopJSON(path, cfg)
}
func writeClaudeDesktopMeta(path, id, name string) error {
meta, err := readClaudeDesktopJSONAllowMissing(path)
if err != nil {
return fmt.Errorf("parse Claude Desktop config metadata: %w", err)
}
meta["appliedId"] = id
entries := make([]any, 0)
for _, entry := range claudeDesktopAnySlice(meta["entries"]) {
entryMap, _ := entry.(map[string]any)
if entryMap == nil {
entries = append(entries, entry)
continue
}
if entryID, _ := entryMap["id"].(string); entryID == id {
continue
}
entries = append(entries, entryMap)
}
entries = append(entries, map[string]any{
"id": id,
"name": name,
})
meta["entries"] = entries
return writeClaudeDesktopJSON(path, meta)
}
func writeClaudeDesktopGatewayProfile(path string, apiKey string, forceChooser bool) error {
cfg, err := readClaudeDesktopJSONAllowMissing(path)
if err != nil {
return fmt.Errorf("parse Claude Desktop Ollama profile: %w", err)
}
cfg["inferenceProvider"] = "gateway"
cfg["inferenceGatewayBaseUrl"] = claudeDesktopGatewayBaseURL
cfg["inferenceGatewayApiKey"] = apiKey
cfg["inferenceGatewayAuthScheme"] = "bearer"
delete(cfg, "inferenceModels")
cfg["disableDeploymentModeChooser"] = forceChooser
return writeClaudeDesktopJSON(path, cfg)
}
func restoreClaudeDesktopMeta(path string) error {
meta, err := readClaudeDesktopJSONAllowMissing(path)
if err != nil {
return fmt.Errorf("parse Claude Desktop config metadata: %w", err)
}
if len(meta) == 0 {
return nil
}
changed := false
if appliedID, _ := meta["appliedId"].(string); appliedID == claudeDesktopProfileID {
delete(meta, "appliedId")
changed = true
}
entries := claudeDesktopAnySlice(meta["entries"])
if entries != nil {
filtered := make([]any, 0, len(entries))
for _, entry := range entries {
entryMap, _ := entry.(map[string]any)
if entryID, _ := entryMap["id"].(string); entryID == claudeDesktopProfileID {
changed = true
continue
}
filtered = append(filtered, entry)
}
meta["entries"] = filtered
}
if !changed {
return nil
}
return writeClaudeDesktopJSON(path, meta)
}
func restoreClaudeDesktopOllamaProfile(path string) error {
cfg, err := readClaudeDesktopJSONAllowMissing(path)
if err != nil {
return fmt.Errorf("parse Claude Desktop Ollama profile: %w", err)
}
if len(cfg) == 0 {
return nil
}
cfg["disableDeploymentModeChooser"] = false
delete(cfg, "inferenceProvider")
delete(cfg, "inferenceGatewayBaseUrl")
delete(cfg, "inferenceGatewayAuthScheme")
delete(cfg, "inferenceModels")
return writeClaudeDesktopJSON(path, cfg)
}
func readClaudeDesktopAppliedID(path string) string {
meta, err := readClaudeDesktopJSON(path)
if err != nil {
return ""
}
applied, _ := meta["appliedId"].(string)
return applied
}
func readClaudeDesktopDeploymentMode(path string) string {
cfg, err := readClaudeDesktopJSON(path)
if err != nil {
return ""
}
mode, _ := cfg["deploymentMode"].(string)
return mode
}
func claudeDesktopTargetsConfigured(targets claudeDesktopTargets) bool {
if len(targets.normalConfigs) == 0 || len(targets.thirdPartyProfiles) == 0 {
return false
}
for _, path := range targets.normalConfigs {
if readClaudeDesktopDeploymentMode(path) != "3p" {
return false
}
}
for _, target := range targets.thirdPartyProfiles {
if readClaudeDesktopDeploymentMode(target.desktopConfig) != "3p" {
return false
}
if !claudeDesktopThirdPartyProfileConfigured(target) {
return false
}
}
return true
}
func claudeDesktopThirdPartyProfileConfigured(target claudeDesktopThirdPartyPaths) bool {
if readClaudeDesktopAppliedID(target.meta) != claudeDesktopProfileID {
return false
}
cfg, err := readClaudeDesktopJSON(target.profile)
if err != nil {
return false
}
if s, _ := cfg["inferenceProvider"].(string); s != "gateway" {
return false
}
if s, _ := cfg["inferenceGatewayBaseUrl"].(string); strings.TrimRight(s, "/") != claudeDesktopGatewayBaseURL {
return false
}
if s, _ := cfg["inferenceGatewayApiKey"].(string); strings.TrimSpace(s) == "" {
return false
}
return true
}
func readClaudeDesktopJSONAllowMissing(path string) (map[string]any, error) {
cfg, err := readClaudeDesktopJSON(path)
if errors.Is(err, os.ErrNotExist) {
return map[string]any{}, nil
}
return cfg, err
}
func readClaudeDesktopJSON(path string) (map[string]any, error) {
data, err := os.ReadFile(path)
if err != nil {
return nil, err
}
var cfg map[string]any
if err := json.Unmarshal(data, &cfg); err != nil {
return nil, err
}
if cfg == nil {
cfg = map[string]any{}
}
return cfg, nil
}
func writeClaudeDesktopJSON(path string, cfg any) error {
data, err := json.MarshalIndent(cfg, "", " ")
if err != nil {
return err
}
data = append(data, '\n')
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
return err
}
return fileutil.WriteWithBackup(path, data)
}
func claudeDesktopAnySlice(value any) []any {
switch v := value.(type) {
case []any:
return v
case nil:
return nil
default:
return nil
}
}
func claudeDesktopLaunchOrRestart(prompt string) error {
if !claudeDesktopIsRunning() {
return claudeDesktopOpenApp()
}
restartAppPath := ""
if claudeDesktopGOOS == "windows" {
restartAppPath = claudeDesktopRunningAppPath()
}
restart, err := ConfirmPrompt(prompt)
if err != nil {
return err
}
if !restart {
fmt.Fprintln(os.Stderr, "\nQuit and reopen Claude Desktop when you're ready for the profile change to take effect.")
return nil
}
if err := claudeDesktopQuitApp(); err != nil {
return fmt.Errorf("quit Claude Desktop: %w", err)
}
if err := waitForClaudeDesktopExit(30 * time.Second); err != nil {
return err
}
if restartAppPath != "" {
return claudeDesktopOpenAppPath(restartAppPath)
}
return claudeDesktopOpenApp()
}
func waitForClaudeDesktopExit(timeout time.Duration) error {
deadline := time.Now().Add(timeout)
for time.Now().Before(deadline) {
if !claudeDesktopIsRunning() {
return nil
}
claudeDesktopSleep(200 * time.Millisecond)
}
return fmt.Errorf("Claude Desktop did not quit; quit it manually and re-run the command")
}
func defaultClaudeDesktopIsRunning() bool {
switch claudeDesktopGOOS {
case "darwin":
out, err := exec.Command("pgrep", "-f", "Claude.app/Contents/MacOS/Claude").Output()
return err == nil && strings.TrimSpace(string(out)) != ""
case "windows":
out, err := exec.Command("powershell.exe", "-NoProfile", "-Command", `(Get-Process claude -ErrorAction SilentlyContinue | Where-Object { $_.MainWindowHandle -ne 0 } | Select-Object -First 1).Id`).Output()
return err == nil && strings.TrimSpace(string(out)) != ""
default:
return false
}
}
func defaultClaudeDesktopOpenApp() error {
switch claudeDesktopGOOS {
case "windows":
if path := claudeDesktopAppPath(); path != "" {
return claudeDesktopOpenAppPath(path)
}
if path := claudeDesktopRunningAppPath(); path != "" {
return claudeDesktopOpenAppPath(path)
}
return fmt.Errorf("Claude Desktop executable was not found; open Claude Desktop manually once and re-run 'ollama launch claude-desktop'")
case "darwin":
return openClaudeDesktopDarwin()
default:
return claudeDesktopSupported()
}
}
func defaultClaudeDesktopOpenAppPath(path string) error {
switch claudeDesktopGOOS {
case "windows":
return exec.Command("powershell.exe", "-NoProfile", "-Command", "Start-Process -FilePath "+quotePowerShellString(path)).Run()
case "darwin":
return openClaudeDesktopDarwin()
default:
return claudeDesktopSupported()
}
}
func openClaudeDesktopDarwin() error {
cmd := exec.Command("open", "-a", "Claude")
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
return cmd.Run()
}
func defaultClaudeDesktopRunningAppPath() string {
if claudeDesktopGOOS != "windows" {
return ""
}
script := `(Get-Process claude -ErrorAction SilentlyContinue | Where-Object { $_.MainWindowHandle -ne 0 -and $_.Path } | Select-Object -First 1 -ExpandProperty Path)`
out, err := exec.Command("powershell.exe", "-NoProfile", "-Command", script).Output()
if err != nil {
return ""
}
return strings.TrimSpace(string(out))
}
func defaultClaudeDesktopQuitApp() error {
if claudeDesktopGOOS == "windows" {
script := `Get-Process claude -ErrorAction SilentlyContinue | Where-Object { $_.MainWindowHandle -ne 0 } | ForEach-Object { [void]$_.CloseMainWindow() }`
return exec.Command("powershell.exe", "-NoProfile", "-Command", script).Run()
}
return exec.Command("osascript", "-e", `tell application "Claude" to quit`).Run()
}
func quotePowerShellString(s string) string {
return "'" + strings.ReplaceAll(s, "'", "''") + "'"
}
+938
View File
@@ -0,0 +1,938 @@
package launch
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"testing"
"time"
)
type roundTripFunc func(*http.Request) (*http.Response, error)
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
return f(req)
}
func withClaudeDesktopPlatform(t *testing.T, goos string) {
t.Helper()
old := claudeDesktopGOOS
claudeDesktopGOOS = goos
t.Cleanup(func() {
claudeDesktopGOOS = old
})
}
func withClaudeDesktopValidation(t *testing.T, fn func(context.Context, string) error) {
t.Helper()
old := claudeDesktopValidateAPIKey
claudeDesktopValidateAPIKey = fn
t.Cleanup(func() {
claudeDesktopValidateAPIKey = old
})
}
func withClaudeDesktopPrompt(t *testing.T, fn func() (string, error)) {
t.Helper()
old := claudeDesktopPromptAPIKey
claudeDesktopPromptAPIKey = fn
t.Cleanup(func() {
claudeDesktopPromptAPIKey = old
})
}
func withClaudeDesktopProcessHooks(t *testing.T, running func() bool, quit func() error, open func() error) {
t.Helper()
oldRunning := claudeDesktopIsRunning
oldQuit := claudeDesktopQuitApp
oldOpen := claudeDesktopOpenApp
oldOpenPath := claudeDesktopOpenAppPath
oldRunningPath := claudeDesktopRunningAppPath
oldSleep := claudeDesktopSleep
claudeDesktopIsRunning = running
claudeDesktopQuitApp = quit
claudeDesktopOpenApp = open
claudeDesktopOpenAppPath = oldOpenPath
claudeDesktopRunningAppPath = oldRunningPath
claudeDesktopSleep = func(time.Duration) {}
t.Cleanup(func() {
claudeDesktopIsRunning = oldRunning
claudeDesktopQuitApp = oldQuit
claudeDesktopOpenApp = oldOpen
claudeDesktopOpenAppPath = oldOpenPath
claudeDesktopRunningAppPath = oldRunningPath
claudeDesktopSleep = oldSleep
})
}
func claudeDesktopReadJSON(t *testing.T, path string) map[string]any {
t.Helper()
data, err := os.ReadFile(path)
if err != nil {
t.Fatalf("read %s: %v", path, err)
}
var cfg map[string]any
if err := json.Unmarshal(data, &cfg); err != nil {
t.Fatalf("parse %s: %v", path, err)
}
return cfg
}
func TestClaudeDesktopIntegration(t *testing.T) {
c := &ClaudeDesktop{}
t.Run("implements Runner", func(t *testing.T) {
var _ Runner = c
})
t.Run("implements managed autodiscovery integration", func(t *testing.T) {
var _ ManagedAutodiscoveryIntegration = c
})
t.Run("does not use local Ollama Cloud auth gate", func(t *testing.T) {
if _, ok := any(c).(ManagedAutodiscoveryCloudIntegration); ok {
t.Fatal("Claude Desktop should validate OLLAMA_API_KEY directly instead of requiring local Ollama Cloud sign-in")
}
})
t.Run("implements restore", func(t *testing.T) {
var _ RestorableIntegration = c
})
t.Run("has restore hint", func(t *testing.T) {
var _ RestoreHintIntegration = c
if !strings.Contains(c.RestoreHint(), "--restore") {
t.Fatalf("expected restore hint to mention --restore, got %q", c.RestoreHint())
}
if strings.Contains(c.RestoreHint(), "Tip:") {
t.Fatalf("restore hint should not use Tip wording, got %q", c.RestoreHint())
}
})
t.Run("has success messages", func(t *testing.T) {
var _ ConfigurationSuccessIntegration = c
var _ RestoreSuccessIntegration = c
if got := c.ConfigurationSuccessMessage(); got != "Claude Desktop profile changed to Ollama Cloud.\nTo restore the usual Claude profile, run: ollama launch claude-desktop --restore" {
t.Fatalf("configuration success message = %q", got)
}
if got := c.RestoreSuccessMessage(); got != "Claude Desktop restored to the usual Claude profile." {
t.Fatalf("restore success message = %q", got)
}
})
t.Run("skips local model readiness", func(t *testing.T) {
var _ ManagedModelReadinessSkipper = c
if !c.SkipModelReadiness() {
t.Fatal("expected Claude Desktop to skip local model readiness")
}
})
}
func TestLaunchIntegration_ClaudeDesktopDoesNotRequireLocalCloudSignIn(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "darwin")
withInteractiveSession(t, true)
withLauncherHooks(t)
t.Setenv("OLLAMA_API_KEY", "test-api-key")
if err := os.MkdirAll(filepath.Join(tmpDir, "Applications", "Claude.app"), 0o755); err != nil {
t.Fatal(err)
}
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/status":
w.WriteHeader(http.StatusNotFound)
fmt.Fprint(w, `{"error":"not found"}`)
case "/api/me":
w.WriteHeader(http.StatusUnauthorized)
fmt.Fprint(w, `{"error":"unauthorized","signin_url":"https://example.com/signin"}`)
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
t.Setenv("OLLAMA_HOST", srv.URL)
var validatedKey string
withClaudeDesktopValidation(t, func(_ context.Context, key string) error {
validatedKey = key
return nil
})
DefaultSignIn = func(modelName, signInURL string) (string, error) {
t.Fatalf("Claude Desktop launch should not require local Ollama Cloud sign-in, got %s at %s", modelName, signInURL)
return "", nil
}
var openCalls int
withClaudeDesktopProcessHooks(t,
func() bool { return false },
func() error { return nil },
func() error {
openCalls++
return nil
},
)
if err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{Name: "claude-desktop"}); err != nil {
t.Fatalf("LaunchIntegration returned error: %v", err)
}
if validatedKey != "test-api-key" {
t.Fatalf("validated key = %q, want test API key", validatedKey)
}
if openCalls != 1 {
t.Fatalf("open calls = %d, want 1", openCalls)
}
}
func TestClaudeDesktopConfigureWritesOllamaCloudProfile(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "darwin")
t.Setenv("OLLAMA_API_KEY", "test-api-key")
var validatedKey string
withClaudeDesktopValidation(t, func(_ context.Context, key string) error {
validatedKey = key
return nil
})
paths, err := claudeDesktopConfigPaths()
if err != nil {
t.Fatal(err)
}
if err := os.MkdirAll(filepath.Dir(paths.desktopConfig), 0o755); err != nil {
t.Fatal(err)
}
if err := os.MkdirAll(filepath.Dir(paths.meta), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(paths.desktopConfig, []byte(`{"existing":true}`), 0o644); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(paths.meta, []byte(`{"entries":[{"id":"custom","name":"Custom"}]}`), 0o644); err != nil {
t.Fatal(err)
}
if err := (&ClaudeDesktop{}).ConfigureAutodiscovery(); err != nil {
t.Fatalf("Configure returned error: %v", err)
}
if validatedKey != "test-api-key" {
t.Fatalf("validated key = %q, want test API key", validatedKey)
}
desktopConfig := claudeDesktopReadJSON(t, paths.desktopConfig)
if desktopConfig["existing"] != true {
t.Fatalf("existing desktop config key was not preserved: %v", desktopConfig)
}
if desktopConfig["deploymentMode"] != "3p" {
t.Fatalf("deploymentMode = %v, want 3p", desktopConfig["deploymentMode"])
}
normalConfig := claudeDesktopReadJSON(t, paths.normalConfig)
if normalConfig["deploymentMode"] != "3p" {
t.Fatalf("normal deploymentMode = %v, want 3p", normalConfig["deploymentMode"])
}
meta := claudeDesktopReadJSON(t, paths.meta)
if meta["appliedId"] != claudeDesktopProfileID {
t.Fatalf("appliedId = %v, want %s", meta["appliedId"], claudeDesktopProfileID)
}
entries, _ := meta["entries"].([]any)
if len(entries) != 2 {
t.Fatalf("entries len = %d, want 2: %v", len(entries), entries)
}
profile := claudeDesktopReadJSON(t, paths.profile)
if profile["inferenceProvider"] != "gateway" {
t.Fatalf("inferenceProvider = %v, want gateway", profile["inferenceProvider"])
}
if profile["inferenceGatewayBaseUrl"] != claudeDesktopGatewayBaseURL {
t.Fatalf("base URL = %v, want %s", profile["inferenceGatewayBaseUrl"], claudeDesktopGatewayBaseURL)
}
if profile["inferenceGatewayApiKey"] != "test-api-key" {
t.Fatal("expected configured API key to be written")
}
if profile["inferenceGatewayAuthScheme"] != "bearer" {
t.Fatalf("auth scheme = %v, want bearer", profile["inferenceGatewayAuthScheme"])
}
if profile["disableDeploymentModeChooser"] != true {
t.Fatalf("disableDeploymentModeChooser = %v, want true", profile["disableDeploymentModeChooser"])
}
if _, ok := profile["inferenceModels"]; ok {
t.Fatalf("inferenceModels should be omitted so Claude can discover models, got %v", profile["inferenceModels"])
}
}
func TestClaudeDesktopConfigureAutodiscoveryRemovesExistingModelCatalog(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "darwin")
t.Setenv("OLLAMA_API_KEY", "test-api-key")
withClaudeDesktopValidation(t, func(context.Context, string) error { return nil })
paths, err := claudeDesktopConfigPaths()
if err != nil {
t.Fatal(err)
}
if err := os.MkdirAll(filepath.Dir(paths.profile), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(paths.profile, []byte(`{"inferenceModels":["qwen3.5"],"inferenceGatewayApiKey":"old"}`), 0o644); err != nil {
t.Fatal(err)
}
if err := (&ClaudeDesktop{}).ConfigureAutodiscovery(); err != nil {
t.Fatalf("ConfigureAutodiscovery returned error: %v", err)
}
profile := claudeDesktopReadJSON(t, paths.profile)
if _, ok := profile["inferenceModels"]; ok {
t.Fatalf("inferenceModels should be removed, got %v", profile["inferenceModels"])
}
if profile["inferenceGatewayApiKey"] != "test-api-key" {
t.Fatal("expected env API key to replace the old key")
}
}
func TestClaudeDesktopWindowsConfigPathsUseLocalAppData(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "windows")
t.Setenv("LOCALAPPDATA", filepath.Join(tmpDir, "LocalAppData"))
paths, err := claudeDesktopConfigPaths()
if err != nil {
t.Fatal(err)
}
if want := filepath.Join(tmpDir, "LocalAppData", "Claude-3p", "claude_desktop_config.json"); paths.desktopConfig != want {
t.Fatalf("desktop config = %q, want %q", paths.desktopConfig, want)
}
if want := filepath.Join(tmpDir, "LocalAppData", "Claude", "claude_desktop_config.json"); paths.normalConfig != want {
t.Fatalf("normal config = %q, want %q", paths.normalConfig, want)
}
}
func TestClaudeDesktopWindowsConfigPathsFallbackToNestProfile(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "windows")
local := filepath.Join(tmpDir, "LocalAppData")
t.Setenv("LOCALAPPDATA", local)
if err := os.MkdirAll(filepath.Join(local, "Claude Nest-3p"), 0o755); err != nil {
t.Fatal(err)
}
paths, err := claudeDesktopConfigPaths()
if err != nil {
t.Fatal(err)
}
if want := filepath.Join(local, "Claude Nest-3p", "claude_desktop_config.json"); paths.desktopConfig != want {
t.Fatalf("desktop config = %q, want %q", paths.desktopConfig, want)
}
}
func TestClaudeDesktopAutodiscoveryConfiguredOnWindows(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "windows")
t.Setenv("LOCALAPPDATA", filepath.Join(tmpDir, "LocalAppData"))
t.Setenv("OLLAMA_API_KEY", "test-api-key")
withClaudeDesktopValidation(t, func(context.Context, string) error { return nil })
c := &ClaudeDesktop{}
if err := c.ConfigureAutodiscovery(); err != nil {
t.Fatalf("Configure returned error: %v", err)
}
if !c.AutodiscoveryConfigured() {
t.Fatal("expected Claude Desktop autodiscovery config to be detected on Windows")
}
}
func TestClaudeDesktopConfigureAutodiscoveryTouchesAllWindowsProfileCandidates(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "windows")
local := filepath.Join(tmpDir, "LocalAppData")
t.Setenv("LOCALAPPDATA", local)
t.Setenv("OLLAMA_API_KEY", "test-api-key")
withClaudeDesktopValidation(t, func(context.Context, string) error { return nil })
targets, err := claudeDesktopTargetPaths()
if err != nil {
t.Fatal(err)
}
if len(targets.normalConfigs) != 2 {
t.Fatalf("normal config target count = %d, want 2", len(targets.normalConfigs))
}
if len(targets.thirdPartyProfiles) != 2 {
t.Fatalf("third-party target count = %d, want 2", len(targets.thirdPartyProfiles))
}
c := &ClaudeDesktop{}
if err := c.ConfigureAutodiscovery(); err != nil {
t.Fatalf("ConfigureAutodiscovery returned error: %v", err)
}
for _, path := range targets.normalConfigs {
cfg := claudeDesktopReadJSON(t, path)
if cfg["deploymentMode"] != "3p" {
t.Fatalf("%s deploymentMode = %v, want 3p", path, cfg["deploymentMode"])
}
}
for _, target := range targets.thirdPartyProfiles {
cfg := claudeDesktopReadJSON(t, target.desktopConfig)
if cfg["deploymentMode"] != "3p" {
t.Fatalf("%s deploymentMode = %v, want 3p", target.desktopConfig, cfg["deploymentMode"])
}
meta := claudeDesktopReadJSON(t, target.meta)
if meta["appliedId"] != claudeDesktopProfileID {
t.Fatalf("%s appliedId = %v, want %s", target.meta, meta["appliedId"], claudeDesktopProfileID)
}
profile := claudeDesktopReadJSON(t, target.profile)
if profile["inferenceProvider"] != "gateway" {
t.Fatalf("%s inferenceProvider = %v, want gateway", target.profile, profile["inferenceProvider"])
}
if profile["inferenceGatewayBaseUrl"] != claudeDesktopGatewayBaseURL {
t.Fatalf("%s base URL = %v, want %s", target.profile, profile["inferenceGatewayBaseUrl"], claudeDesktopGatewayBaseURL)
}
if profile["inferenceGatewayApiKey"] != "test-api-key" {
t.Fatalf("%s should contain the configured API key", target.profile)
}
if _, ok := profile["inferenceModels"]; ok {
t.Fatalf("%s inferenceModels should be omitted, got %v", target.profile, profile["inferenceModels"])
}
}
if !c.AutodiscoveryConfigured() {
t.Fatal("expected all Windows profile candidates to be considered configured")
}
if err := writeClaudeDesktopDeploymentMode(targets.thirdPartyProfiles[1].desktopConfig, "1p"); err != nil {
t.Fatal(err)
}
if c.AutodiscoveryConfigured() {
t.Fatal("expected a stale Windows candidate to force reconfiguration")
}
}
func TestClaudeDesktopInstalledOnWindowsRecognizesLocalProfileDir(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "windows")
local := filepath.Join(tmpDir, "LocalAppData")
t.Setenv("LOCALAPPDATA", local)
withClaudeDesktopProcessHooks(t, func() bool { return false }, func() error { return nil }, func() error { return nil })
if err := os.MkdirAll(filepath.Join(local, "Claude-3p"), 0o755); err != nil {
t.Fatal(err)
}
if !claudeDesktopInstalled() {
t.Fatal("expected Claude Desktop to be installed when the Windows profile directory exists")
}
}
func TestClaudeDesktopWindowsAppPathFindsAnthropicClaudeInstall(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "windows")
local := filepath.Join(tmpDir, "LocalAppData")
t.Setenv("LOCALAPPDATA", local)
want := filepath.Join(local, "AnthropicClaude", "app-1.2.3", "Claude.exe")
if err := os.MkdirAll(filepath.Dir(want), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(want, []byte(""), 0o755); err != nil {
t.Fatal(err)
}
if got := claudeDesktopAppPath(); got != want {
t.Fatalf("claudeDesktopAppPath() = %q, want %q", got, want)
}
}
func TestWaitForClaudeDesktopExitUsesRunningHook(t *testing.T) {
withClaudeDesktopPlatform(t, "windows")
runningChecks := 0
withClaudeDesktopProcessHooks(t,
func() bool {
runningChecks++
return runningChecks == 1
},
func() error { return nil },
func() error { return nil },
)
if err := waitForClaudeDesktopExit(time.Second); err != nil {
t.Fatalf("waitForClaudeDesktopExit returned error: %v", err)
}
if runningChecks < 2 {
t.Fatalf("expected running hook to be checked until the visible window exits, got %d checks", runningChecks)
}
}
func TestClaudeDesktopWindowsRestartUsesCapturedDesktopPath(t *testing.T) {
withClaudeDesktopPlatform(t, "windows")
restoreConfirm := withLaunchConfirmPolicy(launchConfirmPolicy{yes: true})
defer restoreConfirm()
desktopPath := `C:\Users\parth\AppData\Local\AnthropicClaude\app-1.2.3\Claude.exe`
running := true
var openedPath string
withClaudeDesktopProcessHooks(t,
func() bool { return running },
func() error {
running = false
return nil
},
func() error {
t.Fatal("expected restart to open the captured Desktop executable path, not the generic launcher")
return nil
},
)
claudeDesktopRunningAppPath = func() string { return desktopPath }
claudeDesktopOpenAppPath = func(path string) error {
openedPath = path
return nil
}
if err := (&ClaudeDesktop{}).Run("qwen3.5", nil); err != nil {
t.Fatalf("Run returned error: %v", err)
}
if openedPath != desktopPath {
t.Fatalf("opened path = %q, want %q", openedPath, desktopPath)
}
}
func TestClaudeDesktopWindowsOpenDoesNotFallBackToClaudeCommand(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "windows")
t.Setenv("LOCALAPPDATA", filepath.Join(tmpDir, "LocalAppData"))
oldRunningPath := claudeDesktopRunningAppPath
claudeDesktopRunningAppPath = func() string { return "" }
t.Cleanup(func() { claudeDesktopRunningAppPath = oldRunningPath })
err := defaultClaudeDesktopOpenApp()
if err == nil || !strings.Contains(err.Error(), "Claude Desktop executable was not found") {
t.Fatalf("defaultClaudeDesktopOpenApp error = %v, want executable-not-found error", err)
}
}
func TestClaudeDesktopConfigureStopsBeforeWriteWhenKeyValidationFails(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "darwin")
t.Setenv("OLLAMA_API_KEY", "bad-key")
withClaudeDesktopValidation(t, func(context.Context, string) error {
return errors.New("invalid key")
})
err := (&ClaudeDesktop{}).ConfigureAutodiscovery()
if err == nil || !strings.Contains(err.Error(), "invalid key") {
t.Fatalf("Configure error = %v, want invalid key", err)
}
paths, err := claudeDesktopConfigPaths()
if err != nil {
t.Fatal(err)
}
if _, err := os.Stat(paths.desktopConfig); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("desktop config should not be written after validation failure, stat err = %v", err)
}
}
func TestValidateClaudeDesktopAPIKeyUsesClaudeModelsRoute(t *testing.T) {
oldClient := claudeDesktopHTTPClient
var gotPath, gotAuth string
claudeDesktopHTTPClient = &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
gotPath = req.URL.Path
gotAuth = req.Header.Get("Authorization")
return &http.Response{
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader(`{"data":[]}`)),
Header: make(http.Header),
}, nil
})}
t.Cleanup(func() {
claudeDesktopHTTPClient = oldClient
})
if err := validateClaudeDesktopAPIKey(context.Background(), "test-key"); err != nil {
t.Fatalf("validateClaudeDesktopAPIKey returned error: %v", err)
}
if gotPath != "/v1/models" {
t.Fatalf("validation path = %q, want /v1/models", gotPath)
}
if gotAuth != "Bearer test-key" {
t.Fatalf("Authorization header = %q, want bearer key", gotAuth)
}
}
func TestValidateClaudeDesktopAPIKeyHidesInvalidHeaderDetails(t *testing.T) {
err := validateClaudeDesktopAPIKey(context.Background(), "bad\nkey")
if err == nil {
t.Fatal("expected validation error for key with newline")
}
if !strings.Contains(err.Error(), "could not verify Ollama API key") {
t.Fatalf("validation error = %v, want friendly verification message", err)
}
if strings.Contains(err.Error(), "invalid header") || strings.Contains(err.Error(), "net/http") {
t.Fatalf("validation error should not expose transport internals: %v", err)
}
if !strings.Contains(err.Error(), "https://ollama.com/settings/keys") {
t.Fatalf("validation error should include settings link: %v", err)
}
}
func TestClaudeDesktopConfigureRequiresAPIKey(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "darwin")
t.Setenv("OLLAMA_API_KEY", "")
withClaudeDesktopValidation(t, func(context.Context, string) error {
t.Fatal("validation should not run without an API key")
return nil
})
err := (&ClaudeDesktop{}).ConfigureAutodiscovery()
if err == nil || !strings.Contains(err.Error(), "OLLAMA_API_KEY is required") {
t.Fatalf("Configure error = %v, want missing key guidance", err)
}
}
func TestClaudeDesktopAPIKeyPromptIncludesSettingsLink(t *testing.T) {
prompt := claudeDesktopAPIKeyPrompt()
if !strings.Contains(prompt, "Enter Ollama API key") {
t.Fatalf("prompt should ask for the API key, got %q", prompt)
}
if !strings.Contains(prompt, "https://ollama.com/settings/keys") {
t.Fatalf("prompt should include API key settings link, got %q", prompt)
}
}
func TestClaudeDesktopConfigureReusesExistingAPIKey(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "darwin")
t.Setenv("OLLAMA_API_KEY", "")
paths, err := claudeDesktopConfigPaths()
if err != nil {
t.Fatal(err)
}
if err := os.MkdirAll(filepath.Dir(paths.profile), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(paths.profile, []byte(`{"inferenceGatewayApiKey":"existing-key"}`), 0o644); err != nil {
t.Fatal(err)
}
var validatedKey string
withClaudeDesktopValidation(t, func(_ context.Context, key string) error {
validatedKey = key
return nil
})
if err := (&ClaudeDesktop{}).ConfigureAutodiscovery(); err != nil {
t.Fatalf("ConfigureAutodiscovery returned error: %v", err)
}
if validatedKey != "existing-key" {
t.Fatalf("validated key = %q, want existing-key", validatedKey)
}
}
func TestClaudeDesktopConfigureReplacesInvalidExistingAPIKey(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "darwin")
withInteractiveSession(t, true)
t.Setenv("OLLAMA_API_KEY", "")
paths, err := claudeDesktopConfigPaths()
if err != nil {
t.Fatal(err)
}
if err := os.MkdirAll(filepath.Dir(paths.profile), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(paths.profile, []byte(`{"inferenceGatewayApiKey":"stale-key"}`), 0o644); err != nil {
t.Fatal(err)
}
var validated []string
withClaudeDesktopValidation(t, func(_ context.Context, key string) error {
validated = append(validated, key)
if key == "stale-key" {
return errors.New("invalid key")
}
return nil
})
withClaudeDesktopPrompt(t, func() (string, error) {
return "replacement-key", nil
})
if err := (&ClaudeDesktop{}).ConfigureAutodiscovery(); err != nil {
t.Fatalf("ConfigureAutodiscovery returned error: %v", err)
}
if diff := compareStrings(validated, []string{"stale-key", "replacement-key"}); diff != "" {
t.Fatalf("validated keys mismatch: %s", diff)
}
profile := claudeDesktopReadJSON(t, paths.profile)
if profile["inferenceGatewayApiKey"] != "replacement-key" {
t.Fatalf("configured key = %v, want replacement-key", profile["inferenceGatewayApiKey"])
}
}
func TestClaudeDesktopConfigureReusesExistingAPIKeyFromAnyWindowsProfile(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "windows")
local := filepath.Join(tmpDir, "LocalAppData")
t.Setenv("LOCALAPPDATA", local)
t.Setenv("OLLAMA_API_KEY", "")
targets, err := claudeDesktopTargetPaths()
if err != nil {
t.Fatal(err)
}
fallbackProfile := targets.thirdPartyProfiles[1].profile
if err := os.MkdirAll(filepath.Dir(fallbackProfile), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(fallbackProfile, []byte(`{"inferenceGatewayApiKey":"fallback-key"}`), 0o644); err != nil {
t.Fatal(err)
}
var validatedKey string
withClaudeDesktopValidation(t, func(_ context.Context, key string) error {
validatedKey = key
return nil
})
if err := (&ClaudeDesktop{}).ConfigureAutodiscovery(); err != nil {
t.Fatalf("ConfigureAutodiscovery returned error: %v", err)
}
if validatedKey != "fallback-key" {
t.Fatalf("validated key = %q, want fallback-key", validatedKey)
}
for _, target := range targets.thirdPartyProfiles {
profile := claudeDesktopReadJSON(t, target.profile)
if profile["inferenceGatewayApiKey"] != "fallback-key" {
t.Fatalf("%s should reuse fallback key, got %v", target.profile, profile["inferenceGatewayApiKey"])
}
}
}
func TestClaudeDesktopAutodiscoveryConfiguredRequiresAppliedOllamaProfile(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "darwin")
t.Setenv("OLLAMA_API_KEY", "test-api-key")
withClaudeDesktopValidation(t, func(context.Context, string) error { return nil })
c := &ClaudeDesktop{}
if err := c.ConfigureAutodiscovery(); err != nil {
t.Fatalf("Configure returned error: %v", err)
}
if !c.AutodiscoveryConfigured() {
t.Fatal("expected Claude Desktop autodiscovery config to be detected")
}
paths, err := claudeDesktopConfigPaths()
if err != nil {
t.Fatal(err)
}
if err := os.WriteFile(paths.meta, []byte(`{"appliedId":"custom"}`), 0o644); err != nil {
t.Fatal(err)
}
if c.AutodiscoveryConfigured() {
t.Fatal("expected another applied profile to hide Claude Desktop autodiscovery config")
}
}
func TestClaudeDesktopAutodiscoveryConfiguredRequiresAPIKey(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "darwin")
t.Setenv("OLLAMA_API_KEY", "test-api-key")
withClaudeDesktopValidation(t, func(context.Context, string) error { return nil })
c := &ClaudeDesktop{}
if err := c.ConfigureAutodiscovery(); err != nil {
t.Fatalf("Configure returned error: %v", err)
}
paths, err := claudeDesktopConfigPaths()
if err != nil {
t.Fatal(err)
}
profile := claudeDesktopReadJSON(t, paths.profile)
delete(profile, "inferenceGatewayApiKey")
data, err := json.Marshal(profile)
if err != nil {
t.Fatal(err)
}
if err := os.WriteFile(paths.profile, data, 0o644); err != nil {
t.Fatal(err)
}
if c.AutodiscoveryConfigured() {
t.Fatal("expected missing gateway API key to force Claude Desktop reconfiguration")
}
}
func TestClaudeDesktopRestoreSwitchesBackToFirstPartyMode(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "darwin")
withClaudeDesktopProcessHooks(t, func() bool { return false }, func() error { return nil }, func() error { return nil })
paths, err := claudeDesktopConfigPaths()
if err != nil {
t.Fatal(err)
}
if err := os.MkdirAll(filepath.Dir(paths.profile), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(paths.meta, []byte(`{"appliedId":"`+claudeDesktopProfileID+`","entries":[{"id":"`+claudeDesktopProfileID+`","name":"Ollama"}]}`), 0o644); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(paths.profile, []byte(`{"disableDeploymentModeChooser":true,"inferenceGatewayApiKey":"keep","inferenceProvider":"gateway","inferenceGatewayBaseUrl":"https://ollama.com","inferenceGatewayAuthScheme":"bearer","inferenceModels":["legacy"]}`), 0o644); err != nil {
t.Fatal(err)
}
if err := (&ClaudeDesktop{}).Restore(); err != nil {
t.Fatalf("Restore returned error: %v", err)
}
desktopConfig := claudeDesktopReadJSON(t, paths.desktopConfig)
if desktopConfig["deploymentMode"] != "1p" {
t.Fatalf("deploymentMode = %v, want 1p", desktopConfig["deploymentMode"])
}
normalConfig := claudeDesktopReadJSON(t, paths.normalConfig)
if normalConfig["deploymentMode"] != "1p" {
t.Fatalf("normal deploymentMode = %v, want 1p", normalConfig["deploymentMode"])
}
profile := claudeDesktopReadJSON(t, paths.profile)
if profile["disableDeploymentModeChooser"] != false {
t.Fatalf("disableDeploymentModeChooser = %v, want false", profile["disableDeploymentModeChooser"])
}
if profile["inferenceGatewayApiKey"] != "keep" {
t.Fatal("restore should leave existing Ollama profile credentials in place")
}
for _, key := range []string{"inferenceProvider", "inferenceGatewayBaseUrl", "inferenceGatewayAuthScheme", "inferenceModels"} {
if _, ok := profile[key]; ok {
t.Fatalf("restore should clear stale %s from the Ollama profile: %v", key, profile)
}
}
meta := claudeDesktopReadJSON(t, paths.meta)
if _, ok := meta["appliedId"]; ok {
t.Fatalf("restore should clear the applied Ollama third-party profile: %v", meta)
}
if (&ClaudeDesktop{}).AutodiscoveryConfigured() {
t.Fatal("restore should leave Claude Desktop autodiscovery unconfigured")
}
}
func TestClaudeDesktopRestoreTouchesAllWindowsProfileCandidates(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withClaudeDesktopPlatform(t, "windows")
local := filepath.Join(tmpDir, "LocalAppData")
t.Setenv("LOCALAPPDATA", local)
withClaudeDesktopProcessHooks(t, func() bool { return false }, func() error { return nil }, func() error { return nil })
targets, err := claudeDesktopTargetPaths()
if err != nil {
t.Fatal(err)
}
if len(targets.normalConfigs) != 2 {
t.Fatalf("normal config target count = %d, want 2", len(targets.normalConfigs))
}
if len(targets.thirdPartyProfiles) != 2 {
t.Fatalf("third-party target count = %d, want 2", len(targets.thirdPartyProfiles))
}
for _, target := range targets.thirdPartyProfiles {
if err := os.MkdirAll(filepath.Dir(target.profile), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(target.meta, []byte(`{"appliedId":"`+claudeDesktopProfileID+`","entries":[{"id":"`+claudeDesktopProfileID+`","name":"Ollama"}]}`), 0o644); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(target.profile, []byte(`{"disableDeploymentModeChooser":true,"inferenceGatewayApiKey":"keep","inferenceProvider":"gateway","inferenceGatewayBaseUrl":"https://ollama.com","inferenceGatewayAuthScheme":"bearer","inferenceModels":["legacy"]}`), 0o644); err != nil {
t.Fatal(err)
}
}
if err := (&ClaudeDesktop{}).Restore(); err != nil {
t.Fatalf("Restore returned error: %v", err)
}
for _, path := range targets.normalConfigs {
cfg := claudeDesktopReadJSON(t, path)
if cfg["deploymentMode"] != "1p" {
t.Fatalf("%s deploymentMode = %v, want 1p", path, cfg["deploymentMode"])
}
}
for _, target := range targets.thirdPartyProfiles {
cfg := claudeDesktopReadJSON(t, target.desktopConfig)
if cfg["deploymentMode"] != "1p" {
t.Fatalf("%s deploymentMode = %v, want 1p", target.desktopConfig, cfg["deploymentMode"])
}
meta := claudeDesktopReadJSON(t, target.meta)
if _, ok := meta["appliedId"]; ok {
t.Fatalf("%s should not keep the Ollama applied profile: %v", target.meta, meta)
}
profile := claudeDesktopReadJSON(t, target.profile)
if profile["disableDeploymentModeChooser"] != false {
t.Fatalf("%s disableDeploymentModeChooser = %v, want false", target.profile, profile["disableDeploymentModeChooser"])
}
if profile["inferenceGatewayApiKey"] != "keep" {
t.Fatalf("%s should preserve gateway API key", target.profile)
}
for _, key := range []string{"inferenceProvider", "inferenceGatewayBaseUrl", "inferenceGatewayAuthScheme", "inferenceModels"} {
if _, ok := profile[key]; ok {
t.Fatalf("%s should clear stale %s: %v", target.profile, key, profile)
}
}
}
}
func TestClaudeDesktopRunRestartsRunningAppWhenConfirmed(t *testing.T) {
withClaudeDesktopPlatform(t, "darwin")
restoreConfirm := withLaunchConfirmPolicy(launchConfirmPolicy{yes: true})
defer restoreConfirm()
running := true
var quitCalls, openCalls int
withClaudeDesktopProcessHooks(t,
func() bool { return running },
func() error {
quitCalls++
running = false
return nil
},
func() error {
openCalls++
return nil
},
)
if err := (&ClaudeDesktop{}).Run("qwen3.5", nil); err != nil {
t.Fatalf("Run returned error: %v", err)
}
if quitCalls != 1 || openCalls != 1 {
t.Fatalf("quit/open calls = %d/%d, want 1/1", quitCalls, openCalls)
}
}
func TestClaudeDesktopRunRejectsExtraArgs(t *testing.T) {
withClaudeDesktopPlatform(t, "darwin")
err := (&ClaudeDesktop{}).Run("qwen3.5", []string{"--foo"})
if err == nil || !strings.Contains(err.Error(), "does not accept extra arguments") {
t.Fatalf("Run error = %v, want extra args rejection", err)
}
}
+66
View File
@@ -61,6 +61,9 @@ func TestLaunchCmd(t *testing.T) {
if !strings.Contains(cmd.Long, "hermes") {
t.Error("Long description should mention hermes")
}
if !strings.Contains(cmd.Long, "kimi") {
t.Error("Long description should mention kimi")
}
})
t.Run("flags exist", func(t *testing.T) {
@@ -70,6 +73,9 @@ func TestLaunchCmd(t *testing.T) {
if cmd.Flags().Lookup("config") == nil {
t.Error("--config flag should exist")
}
if cmd.Flags().Lookup("restore") == nil {
t.Error("--restore flag should exist")
}
if cmd.Flags().Lookup("yes") == nil {
t.Error("--yes flag should exist")
}
@@ -204,6 +210,27 @@ func TestLaunchCmdTUICallback(t *testing.T) {
t.Error("TUI callback should NOT be called when flags or extra args are provided without an integration")
}
})
t.Run("--restore flag without integration returns error", func(t *testing.T) {
tuiCalled := false
mockTUI := func(cmd *cobra.Command) {
tuiCalled = true
}
cmd := LaunchCmd(mockCheck, mockTUI)
cmd.SetArgs([]string{"--restore"})
err := cmd.Execute()
if err == nil {
t.Fatal("expected --restore without an integration to fail")
}
if !strings.Contains(err.Error(), "require an integration name") {
t.Fatalf("expected integration-name guidance, got %v", err)
}
if tuiCalled {
t.Error("TUI callback should NOT be called when --restore is provided without an integration")
}
})
}
func TestLaunchCmdNilHeartbeat(t *testing.T) {
@@ -270,6 +297,8 @@ func TestLaunchCmdModelFlagClearsDisabledCloudOverride(t *testing.T) {
switch r.URL.Path {
case "/api/status":
fmt.Fprintf(w, `{"cloud":{"disabled":true,"source":"config"}}`)
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`)
case "/api/show":
@@ -326,6 +355,41 @@ func TestLaunchCmdModelFlagClearsDisabledCloudOverride(t *testing.T) {
}
}
func TestLaunchCmdAutodiscoveryDefaultLaunchDoesNotForceConfigure(t *testing.T) {
tmpDir := t.TempDir()
setLaunchTestHome(t, tmpDir)
withInteractiveSession(t, true)
withLauncherHooks(t)
runner := &launcherManagedAutodiscoveryRunner{
autodiscoveryConfigured: true,
}
restore := OverrideIntegration("stubauto", runner)
defer restore()
if err := config.SaveIntegration("stubauto", []string{"Ollama Cloud"}); err != nil {
t.Fatalf("failed to save managed integration config: %v", err)
}
if err := config.MarkIntegrationOnboarded("stubauto"); err != nil {
t.Fatalf("failed to mark integration onboarded: %v", err)
}
cmd := LaunchCmd(func(cmd *cobra.Command, args []string) error { return nil }, func(cmd *cobra.Command) {
t.Fatal("TUI callback should not run for direct integration launch")
})
cmd.SetArgs([]string{"stubauto"})
if err := cmd.Execute(); err != nil {
t.Fatalf("launch command failed: %v", err)
}
if runner.autodiscoveryConfigures != 0 {
t.Fatalf("expected default autodiscovery launch to reuse existing config, got %d configures", runner.autodiscoveryConfigures)
}
if runner.ranModel != "Ollama Cloud" {
t.Fatalf("expected launch to run autodiscovery label, got %q", runner.ranModel)
}
}
func TestLaunchCmdYes_AutoConfirmsLaunchPromptPath(t *testing.T) {
tmpDir := t.TempDir()
setLaunchTestHome(t, tmpDir)
@@ -471,6 +535,8 @@ func TestLaunchCmdIntegrationArgPromptsForModelWithSavedSelection(t *testing.T)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"llama3.2"},{"name":"qwen3:8b"}]}`)
case "/api/show":
+76
View File
@@ -0,0 +1,76 @@
package launch
import (
"fmt"
"os"
"os/exec"
"path/filepath"
"runtime"
"github.com/ollama/ollama/envconfig"
)
// Copilot implements Runner for GitHub Copilot CLI integration.
type Copilot struct{}
func (c *Copilot) String() string { return "Copilot CLI" }
func (c *Copilot) args(model string, extra []string) []string {
var args []string
if model != "" {
args = append(args, "--model", model)
}
args = append(args, extra...)
return args
}
func (c *Copilot) findPath() (string, error) {
if p, err := exec.LookPath("copilot"); err == nil {
return p, nil
}
home, err := os.UserHomeDir()
if err != nil {
return "", err
}
name := "copilot"
if runtime.GOOS == "windows" {
name = "copilot.exe"
}
fallback := filepath.Join(home, ".local", "bin", name)
if _, err := os.Stat(fallback); err != nil {
return "", err
}
return fallback, nil
}
func (c *Copilot) Run(model string, args []string) error {
copilotPath, err := c.findPath()
if err != nil {
return fmt.Errorf("copilot is not installed, install from https://docs.github.com/en/copilot/how-tos/set-up/install-copilot-cli")
}
cmd := exec.Command(copilotPath, c.args(model, args)...)
cmd.Stdin = os.Stdin
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
cmd.Env = append(os.Environ(), c.envVars(model)...)
return cmd.Run()
}
// envVars returns the environment variables that configure Copilot CLI
// to use Ollama as its model provider.
func (c *Copilot) envVars(model string) []string {
env := []string{
"COPILOT_PROVIDER_BASE_URL=" + envconfig.Host().String() + "/v1",
"COPILOT_PROVIDER_API_KEY=",
"COPILOT_PROVIDER_WIRE_API=responses",
}
if model != "" {
env = append(env, "COPILOT_MODEL="+model)
}
return env
}
+161
View File
@@ -0,0 +1,161 @@
package launch
import (
"os"
"path/filepath"
"runtime"
"slices"
"strings"
"testing"
)
func TestCopilotIntegration(t *testing.T) {
c := &Copilot{}
t.Run("String", func(t *testing.T) {
if got := c.String(); got != "Copilot CLI" {
t.Errorf("String() = %q, want %q", got, "Copilot CLI")
}
})
t.Run("implements Runner", func(t *testing.T) {
var _ Runner = c
})
}
func TestCopilotFindPath(t *testing.T) {
c := &Copilot{}
t.Run("finds copilot in PATH", func(t *testing.T) {
tmpDir := t.TempDir()
name := "copilot"
if runtime.GOOS == "windows" {
name = "copilot.exe"
}
fakeBin := filepath.Join(tmpDir, name)
os.WriteFile(fakeBin, []byte("#!/bin/sh\n"), 0o755)
t.Setenv("PATH", tmpDir)
got, err := c.findPath()
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if got != fakeBin {
t.Errorf("findPath() = %q, want %q", got, fakeBin)
}
})
t.Run("returns error when not in PATH", func(t *testing.T) {
t.Setenv("PATH", t.TempDir()) // empty dir, no copilot binary
_, err := c.findPath()
if err == nil {
t.Fatal("expected error, got nil")
}
})
t.Run("falls back to ~/.local/bin/copilot", func(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
t.Setenv("PATH", t.TempDir()) // empty dir, no copilot binary
name := "copilot"
if runtime.GOOS == "windows" {
name = "copilot.exe"
}
fallback := filepath.Join(tmpDir, ".local", "bin", name)
os.MkdirAll(filepath.Dir(fallback), 0o755)
os.WriteFile(fallback, []byte("#!/bin/sh\n"), 0o755)
got, err := c.findPath()
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if got != fallback {
t.Errorf("findPath() = %q, want %q", got, fallback)
}
})
t.Run("returns error when neither PATH nor fallback exists", func(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
t.Setenv("PATH", t.TempDir()) // empty dir, no copilot binary
_, err := c.findPath()
if err == nil {
t.Fatal("expected error, got nil")
}
})
}
func TestCopilotArgs(t *testing.T) {
c := &Copilot{}
tests := []struct {
name string
model string
args []string
want []string
}{
{"with model", "llama3.2", nil, []string{"--model", "llama3.2"}},
{"empty model", "", nil, nil},
{"with model and extra", "llama3.2", []string{"--verbose"}, []string{"--model", "llama3.2", "--verbose"}},
{"empty model with help", "", []string{"--help"}, []string{"--help"}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := c.args(tt.model, tt.args)
if !slices.Equal(got, tt.want) {
t.Errorf("args(%q, %v) = %v, want %v", tt.model, tt.args, got, tt.want)
}
})
}
}
func TestCopilotEnvVars(t *testing.T) {
c := &Copilot{}
envMap := func(envs []string) map[string]string {
m := make(map[string]string)
for _, e := range envs {
k, v, _ := strings.Cut(e, "=")
m[k] = v
}
return m
}
t.Run("sets required provider env vars with model", func(t *testing.T) {
got := envMap(c.envVars("llama3.2"))
if got["COPILOT_PROVIDER_BASE_URL"] == "" {
t.Error("COPILOT_PROVIDER_BASE_URL should be set")
}
if !strings.HasSuffix(got["COPILOT_PROVIDER_BASE_URL"], "/v1") {
t.Errorf("COPILOT_PROVIDER_BASE_URL = %q, want /v1 suffix", got["COPILOT_PROVIDER_BASE_URL"])
}
if _, ok := got["COPILOT_PROVIDER_API_KEY"]; !ok {
t.Error("COPILOT_PROVIDER_API_KEY should be set (empty)")
}
if got["COPILOT_PROVIDER_WIRE_API"] != "responses" {
t.Errorf("COPILOT_PROVIDER_WIRE_API = %q, want %q", got["COPILOT_PROVIDER_WIRE_API"], "responses")
}
if got["COPILOT_MODEL"] != "llama3.2" {
t.Errorf("COPILOT_MODEL = %q, want %q", got["COPILOT_MODEL"], "llama3.2")
}
})
t.Run("omits COPILOT_MODEL when model is empty", func(t *testing.T) {
got := envMap(c.envVars(""))
if _, ok := got["COPILOT_MODEL"]; ok {
t.Errorf("COPILOT_MODEL should not be set for empty model, got %q", got["COPILOT_MODEL"])
}
})
t.Run("uses custom OLLAMA_HOST", func(t *testing.T) {
t.Setenv("OLLAMA_HOST", "http://myhost:9999")
got := envMap(c.envVars("test"))
if !strings.Contains(got["COPILOT_PROVIDER_BASE_URL"], "myhost:9999") {
t.Errorf("COPILOT_PROVIDER_BASE_URL = %q, want custom host", got["COPILOT_PROVIDER_BASE_URL"])
}
})
}
+33 -316
View File
@@ -4,18 +4,15 @@ import (
"bufio"
"bytes"
"context"
"errors"
"fmt"
"net/http"
"os"
"os/exec"
pathpkg "path"
"path/filepath"
"runtime"
"slices"
"strconv"
"strings"
"time"
"gopkg.in/yaml.v3"
@@ -66,23 +63,13 @@ var hermesMessagingEnvGroups = [][]string{
// switching UX after startup.
type Hermes struct{}
type hermesConfigBackend struct {
displayPath string
read func() ([]byte, error)
write func([]byte) error
}
func (h *Hermes) String() string { return "Hermes Agent" }
func (h *Hermes) Run(_ string, args []string) error {
// Hermes reads its primary model from config.yaml. launch configures that
// default model ahead of time so we can keep runtime invocation simple and
// still let Hermes discover additional models later via its own UX.
if hermesGOOS == "windows" {
return h.runWindows(args)
}
bin, err := h.findUnixBinary()
bin, err := h.binary()
if err != nil {
return err
}
@@ -95,21 +82,21 @@ func (h *Hermes) Run(_ string, args []string) error {
}
func (h *Hermes) Paths() []string {
backend, err := h.configBackend()
configPath, err := hermesConfigPath()
if err != nil {
return nil
}
return []string{backend.displayPath}
return []string{configPath}
}
func (h *Hermes) Configure(model string) error {
backend, err := h.configBackend()
configPath, err := hermesConfigPath()
if err != nil {
return err
}
cfg := map[string]any{}
if data, err := backend.read(); err == nil {
if data, err := os.ReadFile(configPath); err == nil {
if err := yaml.Unmarshal(data, &cfg); err != nil {
return fmt.Errorf("parse hermes config: %w", err)
}
@@ -142,15 +129,18 @@ func (h *Hermes) Configure(model string) error {
if err != nil {
return err
}
return backend.write(data)
if err := os.MkdirAll(filepath.Dir(configPath), 0o755); err != nil {
return err
}
return fileutil.WriteWithBackup(configPath, data)
}
func (h *Hermes) CurrentModel() string {
backend, err := h.configBackend()
configPath, err := hermesConfigPath()
if err != nil {
return ""
}
data, err := backend.read()
data, err := os.ReadFile(configPath)
if err != nil {
return ""
}
@@ -188,14 +178,7 @@ func (h *Hermes) RefreshRuntimeAfterConfigure() error {
}
func (h *Hermes) installed() bool {
if hermesGOOS == "windows" {
if _, err := hermesLookPath("hermes"); err == nil {
return true
}
return h.wslHasHermes()
}
_, err := h.findUnixBinary()
_, err := h.binary()
return err == nil
}
@@ -205,7 +188,7 @@ func (h *Hermes) ensureInstalled() error {
}
if hermesGOOS == "windows" {
return h.ensureInstalledWindows()
return hermesWindowsHint()
}
var missing []string
@@ -239,42 +222,6 @@ func (h *Hermes) ensureInstalled() error {
return nil
}
func (h *Hermes) ensureInstalledWindows() error {
// Hermes upstream support is WSL-oriented, so Windows launch uses a hybrid
// WSL handoff that stays on the same install path as upstream Hermes.
if _, err := hermesLookPath("hermes"); err == nil {
return nil
}
if !h.wslAvailable() {
return hermesWindowsHint(fmt.Errorf("hermes is not installed"))
}
if h.wslHasHermes() {
return nil
}
ok, err := ConfirmPromptWithOptions("Hermes runs through WSL2 on Windows. Install it in WSL now?", ConfirmOptions{
YesLabel: "Use WSL",
NoLabel: "Show manual steps",
})
if err != nil {
return err
}
if !ok {
return hermesWindowsHint(fmt.Errorf("hermes is not installed"))
}
fmt.Fprintf(os.Stderr, "\nInstalling Hermes in WSL...\n")
if err := h.runWSL("bash", "-lc", hermesInstallScript); err != nil {
return hermesWindowsHint(fmt.Errorf("failed to install hermes in WSL: %w", err))
}
if !h.wslHasHermes() {
return hermesWindowsHint(fmt.Errorf("hermes install finished but the WSL binary was not found"))
}
fmt.Fprintf(os.Stderr, "%sHermes installed successfully in WSL%s\n\n", ansiGreen, ansiReset)
return nil
}
func (h *Hermes) listModels(defaultModel string) []string {
client := hermesOllamaClient()
resp, err := client.List(context.Background())
@@ -306,11 +253,15 @@ func (h *Hermes) listModels(defaultModel string) []string {
return models
}
func (h *Hermes) findUnixBinary() (string, error) {
func (h *Hermes) binary() (string, error) {
if path, err := hermesLookPath("hermes"); err == nil {
return path, nil
}
if hermesGOOS == "windows" {
return "", hermesWindowsHint()
}
home, err := hermesUserHome()
if err != nil {
return "", err
@@ -323,70 +274,6 @@ func (h *Hermes) findUnixBinary() (string, error) {
return "", fmt.Errorf("hermes is not installed")
}
func (h *Hermes) runWindows(args []string) error {
if path, err := hermesLookPath("hermes"); err == nil {
if err := h.runGatewaySetupPreflight(args, func() error {
return hermesAttachedCommand(path, "gateway", "setup").Run()
}); err != nil {
return err
}
return hermesAttachedCommand(path, args...).Run()
}
if !h.wslAvailable() {
return hermesWindowsHint(fmt.Errorf("hermes is not installed"))
}
if err := h.runGatewaySetupPreflight(args, func() error {
return h.runWSL("hermes", "gateway", "setup")
}); err != nil {
return err
}
if err := h.runWSL(append([]string{"hermes"}, args...)...); err != nil {
return hermesWindowsHint(err)
}
return nil
}
func (h *Hermes) runWSL(args ...string) error {
if !h.wslAvailable() {
return fmt.Errorf("wsl.exe is not available")
}
return hermesAttachedCommand("wsl.exe", "bash", "-lc", shellQuoteArgs(args)).Run()
}
func (h *Hermes) runWSLCombinedOutput(args ...string) ([]byte, error) {
if !h.wslAvailable() {
return nil, fmt.Errorf("wsl.exe is not available")
}
return hermesCommand("wsl.exe", "bash", "-lc", shellQuoteArgs(args)).CombinedOutput()
}
func (h *Hermes) wslAvailable() bool {
_, err := hermesLookPath("wsl.exe")
return err == nil
}
func (h *Hermes) wslHasHermes() bool {
if !h.wslAvailable() {
return false
}
cmd := hermesCommand("wsl.exe", "bash", "-lc", "command -v hermes >/dev/null 2>&1")
return cmd.Run() == nil
}
func (h *Hermes) configBackend() (*hermesConfigBackend, error) {
if hermesGOOS == "windows" {
if _, err := hermesLookPath("hermes"); err == nil {
return hermesLocalConfigBackend()
}
if h.wslAvailable() {
return h.wslConfigBackend()
}
}
return hermesLocalConfigBackend()
}
func hermesConfigPath() (string, error) {
home, err := hermesUserHome()
if err != nil {
@@ -395,110 +282,6 @@ func hermesConfigPath() (string, error) {
return filepath.Join(home, ".hermes", "config.yaml"), nil
}
func hermesLocalConfigBackend() (*hermesConfigBackend, error) {
configPath, err := hermesConfigPath()
if err != nil {
return nil, err
}
return &hermesConfigBackend{
displayPath: configPath,
read: func() ([]byte, error) {
return os.ReadFile(configPath)
},
write: func(data []byte) error {
if err := os.MkdirAll(filepath.Dir(configPath), 0o755); err != nil {
return err
}
return fileutil.WriteWithBackup(configPath, data)
},
}, nil
}
func (h *Hermes) wslConfigBackend() (*hermesConfigBackend, error) {
home, err := h.wslHome()
if err != nil {
return nil, err
}
configPath := pathpkg.Join(home, ".hermes", "config.yaml")
return &hermesConfigBackend{
displayPath: configPath,
read: func() ([]byte, error) {
return h.readWSLFile(configPath)
},
write: func(data []byte) error {
return h.writeWSLConfig(configPath, data)
},
}, nil
}
func (h *Hermes) wslHome() (string, error) {
if !h.wslAvailable() {
return "", fmt.Errorf("wsl.exe is not available")
}
cmd := hermesCommand("wsl.exe", "bash", "-lc", `printf %s "$HOME"`)
out, err := cmd.Output()
if err != nil {
return "", err
}
home := strings.TrimSpace(string(out))
if home == "" {
return "", fmt.Errorf("could not resolve WSL home directory")
}
return home, nil
}
func (h *Hermes) readWSLFile(path string) ([]byte, error) {
pathArg := shellQuoteArgs([]string{path})
cmd := hermesCommand("wsl.exe", "bash", "-lc", fmt.Sprintf("if [ -f %s ]; then cat %s; else exit 42; fi", pathArg, pathArg))
out, err := cmd.Output()
if err == nil {
return out, nil
}
var exitErr *exec.ExitError
if errors.As(err, &exitErr) && exitErr.ExitCode() == 42 {
return nil, os.ErrNotExist
}
return nil, err
}
func (h *Hermes) writeWSLConfig(path string, data []byte) error {
if existing, err := h.readWSLFile(path); err == nil {
if !bytes.Equal(existing, data) {
if err := hermesBackupData(path, existing); err != nil {
return fmt.Errorf("backup failed: %w", err)
}
}
} else if !os.IsNotExist(err) {
return fmt.Errorf("read existing file: %w", err)
}
dir := pathpkg.Dir(path)
dirArg := shellQuoteArgs([]string{dir})
pathArg := shellQuoteArgs([]string{path})
script := fmt.Sprintf(
"dir=%s; path=%s; mkdir -p \"$dir\" && tmp=$(mktemp \"$dir/.tmp-XXXXXX\") && cat > \"$tmp\" && mv \"$tmp\" \"$path\"",
dirArg,
pathArg,
)
cmd := hermesCommand("wsl.exe", "bash", "-lc", script)
cmd.Stdin = bytes.NewReader(data)
if out, err := cmd.CombinedOutput(); err != nil {
if msg := strings.TrimSpace(string(out)); msg != "" {
return fmt.Errorf("%w: %s", err, msg)
}
return err
}
return nil
}
func hermesBackupData(path string, data []byte) error {
if err := os.MkdirAll(fileutil.BackupDir(), 0o755); err != nil {
return err
}
backupPath := filepath.Join(fileutil.BackupDir(), fmt.Sprintf("%s.%d", filepath.Base(path), time.Now().Unix()))
return os.WriteFile(backupPath, data, 0o644)
}
func hermesBaseURL() string {
return strings.TrimRight(hermesOllamaURL().String(), "/") + "/v1"
}
@@ -554,8 +337,11 @@ func (h *Hermes) messagingConfigured() bool {
func (h *Hermes) gatewayEnvVars() (map[string]string, error) {
envVars := make(map[string]string)
data, err := h.readGatewayEnvFile()
switch {
envFilePath, err := hermesEnvPath()
if err != nil {
return nil, err
}
switch data, err := os.ReadFile(envFilePath); {
case err == nil:
for key, value := range hermesParseEnvFile(data) {
envVars[key] = value
@@ -566,12 +352,10 @@ func (h *Hermes) gatewayEnvVars() (map[string]string, error) {
return nil, err
}
if h.usesLocalRuntimeEnv() {
for _, group := range hermesMessagingEnvGroups {
for _, key := range group {
if value, ok := os.LookupEnv(key); ok {
envVars[key] = value
}
for _, group := range hermesMessagingEnvGroups {
for _, key := range group {
if value, ok := os.LookupEnv(key); ok {
envVars[key] = value
}
}
}
@@ -579,39 +363,6 @@ func (h *Hermes) gatewayEnvVars() (map[string]string, error) {
return envVars, nil
}
func (h *Hermes) readGatewayEnvFile() ([]byte, error) {
if hermesGOOS == "windows" {
if _, err := hermesLookPath("hermes"); err == nil {
path, err := hermesEnvPath()
if err != nil {
return nil, err
}
return os.ReadFile(path)
}
if h.wslAvailable() {
home, err := h.wslHome()
if err != nil {
return nil, err
}
return h.readWSLFile(pathpkg.Join(home, ".hermes", ".env"))
}
}
path, err := hermesEnvPath()
if err != nil {
return nil, err
}
return os.ReadFile(path)
}
func (h *Hermes) usesLocalRuntimeEnv() bool {
if hermesGOOS != "windows" {
return true
}
_, err := hermesLookPath("hermes")
return err == nil
}
func (h *Hermes) gatewayRunning() (bool, error) {
status, err := h.gatewayStatusOutput()
if err != nil {
@@ -621,19 +372,7 @@ func (h *Hermes) gatewayRunning() (bool, error) {
}
func (h *Hermes) gatewayStatusOutput() (string, error) {
if hermesGOOS == "windows" {
if path, err := hermesLookPath("hermes"); err == nil {
out, err := hermesCommand(path, "gateway", "status").CombinedOutput()
return string(out), err
}
if !h.wslAvailable() {
return "", hermesWindowsHint(fmt.Errorf("hermes is not installed"))
}
out, err := h.runWSLCombinedOutput("hermes", "gateway", "status")
return string(out), err
}
bin, err := h.findUnixBinary()
bin, err := h.binary()
if err != nil {
return "", err
}
@@ -642,20 +381,7 @@ func (h *Hermes) gatewayStatusOutput() (string, error) {
}
func (h *Hermes) restartGateway() error {
if hermesGOOS == "windows" {
if path, err := hermesLookPath("hermes"); err == nil {
return hermesAttachedCommand(path, "gateway", "restart").Run()
}
if !h.wslAvailable() {
return hermesWindowsHint(fmt.Errorf("hermes is not installed"))
}
if err := h.runWSL("hermes", "gateway", "restart"); err != nil {
return hermesWindowsHint(err)
}
return nil
}
bin, err := h.findUnixBinary()
bin, err := h.binary()
if err != nil {
return err
}
@@ -938,14 +664,6 @@ func mergeHermesToolsets(current any) any {
}
}
func shellQuoteArgs(args []string) string {
quoted := make([]string, 0, len(args))
for _, arg := range args {
quoted = append(quoted, "'"+strings.ReplaceAll(arg, "'", `'\''`)+"'")
}
return strings.Join(quoted, " ")
}
func hermesAttachedCommand(name string, args ...string) *exec.Cmd {
cmd := hermesCommand(name, args...)
cmd.Stdin = os.Stdin
@@ -954,9 +672,8 @@ func hermesAttachedCommand(name string, args ...string) *exec.Cmd {
return cmd
}
func hermesWindowsHint(err error) error {
if hermesGOOS != "windows" {
return err
}
return fmt.Errorf("%w\n\nHermes runs on Windows through WSL2.\nQuick setup: wsl --install\nInstaller docs: https://hermes-agent.nousresearch.com/docs/getting-started/installation/", err)
func hermesWindowsHint() error {
return fmt.Errorf("Hermes on Windows requires WSL2. Install WSL with: wsl --install\n" +
"Then run 'ollama launch hermes' from inside your WSL shell.\n" +
"Docs: https://hermes-agent.nousresearch.com/docs/getting-started/installation/")
}
+19 -146
View File
@@ -49,15 +49,6 @@ func withHermesUserHome(t *testing.T, dir string) {
})
}
func withHermesLookPath(t *testing.T, fn func(string) (string, error)) {
t.Helper()
old := hermesLookPath
hermesLookPath = fn
t.Cleanup(func() {
hermesLookPath = old
})
}
func clearHermesMessagingEnvVars(t *testing.T) {
t.Helper()
for _, group := range hermesMessagingEnvGroups {
@@ -112,6 +103,8 @@ func TestHermesConfigurePreservesExistingConfigAndEnablesWeb(t *testing.T) {
switch r.URL.Path {
case "/api/show":
fmt.Fprint(w, `{"model_info":{"general.context_length":131072}}`)
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"gemma4"},{"name":"qwen3.5"},{"name":"llama3.3"}]}`)
default:
@@ -224,6 +217,8 @@ func TestHermesConfigureUpdatesMatchingCustomProviderWithoutDroppingFields(t *te
switch r.URL.Path {
case "/api/show":
fmt.Fprint(w, `{"model_info":{"general.context_length":131072}}`)
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"gemma4"},{"name":"qwen3.5"},{"name":"llama3.3"}]}`)
default:
@@ -300,6 +295,8 @@ func TestHermesConfigureUsesLaunchResolvedHostForModelDiscovery(t *testing.T) {
switch r.URL.Path {
case "/api/show":
fmt.Fprint(w, `{"model_info":{"general.context_length":131072}}`)
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"gemma4"},{"name":"qwen3.5"},{"name":"llama3.3"}]}`)
default:
@@ -365,6 +362,8 @@ func TestHermesConfigureMigratesLegacyManagedAliases(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"gemma4"},{"name":"qwen3.5"}]}`)
default:
@@ -896,64 +895,6 @@ fi
}
}
func TestHermesRefreshRuntimeAfterConfigure_WindowsWSLRestartsRunningGateway(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses POSIX shell test binaries to simulate WSL")
}
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withHermesPlatform(t, "windows")
t.Setenv("PATH", tmpDir+string(os.PathListSeparator)+os.Getenv("PATH"))
wslPath := filepath.Join(tmpDir, "wsl.exe")
wslScript := `#!/bin/sh
printf '[%s]\n' "$*" >> "$HOME/wsl-invocations.log"
exec /bin/sh -lc "$3"
`
if err := os.WriteFile(wslPath, []byte(wslScript), 0o755); err != nil {
t.Fatal(err)
}
hermesBin := filepath.Join(tmpDir, "hermes")
hermesScript := `#!/bin/sh
printf '[%s]\n' "$*" >> "$HOME/hermes-invocations.log"
if [ "$1" = "gateway" ] && [ "$2" = "status" ]; then
printf '✓ Gateway is running (PID: 321)\n'
fi
`
if err := os.WriteFile(hermesBin, []byte(hermesScript), 0o755); err != nil {
t.Fatal(err)
}
withHermesLookPath(t, func(file string) (string, error) {
if file == "wsl.exe" {
return wslPath, nil
}
return "", os.ErrNotExist
})
h := &Hermes{}
if err := h.RefreshRuntimeAfterConfigure(); err != nil {
t.Fatalf("RefreshRuntimeAfterConfigure returned error: %v", err)
}
data, err := os.ReadFile(filepath.Join(tmpDir, "hermes-invocations.log"))
if err != nil {
t.Fatal(err)
}
lines := strings.Split(strings.TrimSpace(string(data)), "\n")
if len(lines) != 2 {
t.Fatalf("expected WSL status then restart invocations, got %v", lines)
}
if lines[0] != "[gateway status]" {
t.Fatalf("expected WSL gateway status first, got %q", lines[0])
}
if lines[1] != "[gateway restart]" {
t.Fatalf("expected WSL gateway restart second, got %q", lines[1])
}
}
func TestHermesMessagingConfiguredRecognizesSupportedGatewayVars(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
@@ -1002,82 +943,7 @@ func TestHermesMessagingConfiguredRecognizesSupportedGatewayVars(t *testing.T) {
}
}
func TestHermesRunWindowsWSL_UsesGatewaySetupPreflight(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses POSIX shell test binaries to simulate WSL")
}
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withLauncherHooks(t)
withInteractiveSession(t, true)
withHermesPlatform(t, "windows")
clearHermesMessagingEnvVars(t)
t.Setenv("PATH", tmpDir+string(os.PathListSeparator)+os.Getenv("PATH"))
wslPath := filepath.Join(tmpDir, "wsl.exe")
wslScript := `#!/bin/sh
printf '[%s]\n' "$*" >> "$HOME/wsl-invocations.log"
exec /bin/sh -lc "$3"
`
if err := os.WriteFile(wslPath, []byte(wslScript), 0o755); err != nil {
t.Fatal(err)
}
hermesBin := filepath.Join(tmpDir, "hermes")
hermesScript := `#!/bin/sh
printf '[%s]\n' "$*" >> "$HOME/hermes-invocations.log"
if [ "$1" = "gateway" ] && [ "$2" = "setup" ]; then
/bin/mkdir -p "$HOME/.hermes"
printf 'TELEGRAM_BOT_TOKEN=configured\n' > "$HOME/.hermes/.env"
fi
`
if err := os.WriteFile(hermesBin, []byte(hermesScript), 0o755); err != nil {
t.Fatal(err)
}
withHermesLookPath(t, func(file string) (string, error) {
if file == "wsl.exe" {
return wslPath, nil
}
return "", os.ErrNotExist
})
promptCount := 0
DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) {
promptCount++
if prompt != hermesGatewaySetupTitle {
t.Fatalf("unexpected prompt %q", prompt)
}
return true, nil
}
h := &Hermes{}
if err := h.Run("", nil); err != nil {
t.Fatalf("Run returned error: %v", err)
}
if promptCount != 1 {
t.Fatalf("expected one messaging prompt, got %d", promptCount)
}
data, err := os.ReadFile(filepath.Join(tmpDir, "hermes-invocations.log"))
if err != nil {
t.Fatal(err)
}
lines := strings.Split(strings.TrimSpace(string(data)), "\n")
if len(lines) != 2 {
t.Fatalf("expected WSL hermes to run setup then launch, got %v", lines)
}
if lines[0] != "[gateway setup]" {
t.Fatalf("expected WSL gateway setup first, got %q", lines[0])
}
if lines[1] != "[]" {
t.Fatalf("expected WSL default hermes launch second, got %q", lines[1])
}
}
func TestHermesEnsureInstalledWindowsWithoutWSLGivesGuidance(t *testing.T) {
func TestHermesEnsureInstalledWindowsShowsWSLGuidance(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
withHermesPlatform(t, "windows")
@@ -1086,10 +952,17 @@ func TestHermesEnsureInstalledWindowsWithoutWSLGivesGuidance(t *testing.T) {
h := &Hermes{}
err := h.ensureInstalled()
if err == nil {
t.Fatal("expected missing WSL guidance error")
t.Fatal("expected WSL guidance error")
}
if !strings.Contains(err.Error(), "wsl --install") {
t.Fatalf("expected WSL guidance, got %v", err)
msg := err.Error()
if !strings.Contains(msg, "wsl --install") {
t.Fatalf("expected install command in guidance, got %v", err)
}
if !strings.Contains(msg, "hermes-agent.nousresearch.com") {
t.Fatalf("expected docs link in guidance, got %v", err)
}
if strings.Contains(msg, "hermes is not installed") {
t.Fatalf("guidance should not lead with 'hermes is not installed', got %v", err)
}
}
+164 -27
View File
@@ -53,9 +53,13 @@ func TestIntegrationLookup(t *testing.T) {
{"claude lowercase", "claude", true, "Claude Code"},
{"claude uppercase", "CLAUDE", true, "Claude Code"},
{"claude mixed case", "Claude", true, "Claude Code"},
{"claude desktop", "claude-desktop", true, "Claude Desktop"},
{"claude desktop alias", "claude-app", true, "Claude Desktop"},
{"codex", "codex", true, "Codex"},
{"kimi", "kimi", true, "Kimi Code CLI"},
{"droid", "droid", true, "Droid"},
{"opencode", "opencode", true, "OpenCode"},
{"pool", "pool", true, "Pool"},
{"unknown integration", "unknown", false, ""},
{"empty string", "", false, ""},
}
@@ -74,8 +78,7 @@ func TestIntegrationLookup(t *testing.T) {
}
func TestIntegrationRegistry(t *testing.T) {
expectedIntegrations := []string{"claude", "codex", "droid", "opencode", "hermes"}
expectedIntegrations := []string{"claude", "claude-desktop", "codex", "kimi", "droid", "opencode", "hermes", "pool"}
for _, name := range expectedIntegrations {
t.Run(name, func(t *testing.T) {
r, ok := integrations[name]
@@ -89,6 +92,15 @@ func TestIntegrationRegistry(t *testing.T) {
}
}
func TestHiddenIntegrationsExcludedFromVisibleLists(t *testing.T) {
for _, info := range ListIntegrationInfos() {
switch info.Name {
case "cline", "vscode", "kimi":
t.Fatalf("hidden integration %q should not appear in ListIntegrationInfos", info.Name)
}
}
}
func TestHasLocalModel(t *testing.T) {
tests := []struct {
name string
@@ -291,7 +303,7 @@ func TestParseArgs(t *testing.T) {
func TestIsCloudModel(t *testing.T) {
// isCloudModel now only uses Show API, so nil client always returns false
t.Run("nil client returns false", func(t *testing.T) {
models := []string{"glm-5.1:cloud", "kimi-k2.5:cloud", "local-model"}
models := []string{"glm-5.1:cloud", "kimi-k2.6:cloud", "local-model"}
for _, model := range models {
if isCloudModel(context.Background(), nil, model) {
t.Errorf("isCloudModel(%q) with nil client should return false", model)
@@ -308,10 +320,18 @@ func names(items []ModelItem) []string {
return out
}
func recommendedNames(extra ...string) []string {
out := make([]string, 0, len(recommendedModels)+len(extra))
for _, item := range recommendedModels {
out = append(out, item.Name)
}
return append(out, extra...)
}
func TestBuildModelList_NoExistingModels(t *testing.T) {
items, _, _, _ := buildModelList(nil, nil, "")
want := []string{"kimi-k2.5:cloud", "qwen3.5:cloud", "glm-5.1:cloud", "minimax-m2.7:cloud", "gemma4", "qwen3.5"}
want := recommendedNames()
if diff := cmp.Diff(want, names(items)); diff != "" {
t.Errorf("with no existing models, items should be recommended in order (-want +got):\n%s", diff)
}
@@ -340,7 +360,7 @@ func TestBuildModelList_OnlyLocalModels_CloudRecsStillFirst(t *testing.T) {
// Cloud recs always come first among recommended, regardless of installed inventory.
// Cloud disablement is handled upstream in loadSelectableModels via filterCloudItems.
want := []string{"kimi-k2.5:cloud", "qwen3.5:cloud", "glm-5.1:cloud", "minimax-m2.7:cloud", "gemma4", "qwen3.5", "llama3.2", "qwen2.5"}
want := recommendedNames("llama3.2", "qwen2.5")
if diff := cmp.Diff(want, got); diff != "" {
t.Errorf("cloud recs pinned first even when no cloud models installed (-want +got):\n%s", diff)
}
@@ -356,13 +376,13 @@ func TestBuildModelList_BothCloudAndLocal_RegularSort(t *testing.T) {
got := names(items)
// All recs pinned at top (cloud before local in mixed case), then non-recs
want := []string{"kimi-k2.5:cloud", "qwen3.5:cloud", "glm-5.1:cloud", "minimax-m2.7:cloud", "gemma4", "qwen3.5", "llama3.2"}
want := recommendedNames("llama3.2")
if diff := cmp.Diff(want, got); diff != "" {
t.Errorf("recs pinned at top, cloud recs first in mixed case (-want +got):\n%s", diff)
}
}
func TestBuildModelList_PreCheckedFirst(t *testing.T) {
func TestBuildModelList_PreCheckedNonRecommendedFirstInMore(t *testing.T) {
existing := []modelInfo{
{Name: "llama3.2:latest", Remote: false},
{Name: "glm-5.1:cloud", Remote: true},
@@ -371,8 +391,9 @@ func TestBuildModelList_PreCheckedFirst(t *testing.T) {
items, _, _, _ := buildModelList(existing, []string{"llama3.2"}, "")
got := names(items)
if got[0] != "llama3.2" {
t.Errorf("pre-checked model should be first, got %v", got)
want := recommendedNames("llama3.2")
if diff := cmp.Diff(want, got); diff != "" {
t.Errorf("recommended block should stay fixed while checked non-recommended models lead More (-want +got):\n%s", diff)
}
}
@@ -427,7 +448,7 @@ func TestBuildModelList_ExistingRecommendedMarked(t *testing.T) {
if !strings.HasSuffix(item.Description, "(not downloaded)") {
t.Errorf("non-installed recommended %q should have '(not downloaded)' suffix, got %q", item.Name, item.Description)
}
case "minimax-m2.7:cloud", "kimi-k2.5:cloud", "qwen3.5:cloud":
case "minimax-m2.7:cloud", "kimi-k2.6:cloud", "qwen3.5:cloud":
if strings.HasSuffix(item.Description, "(not downloaded)") {
t.Errorf("cloud model %q should not have '(not downloaded)' suffix, got %q", item.Name, item.Description)
}
@@ -445,9 +466,9 @@ func TestBuildModelList_ExistingCloudModelsNotPushedToBottom(t *testing.T) {
got := names(items)
// gemma4 and glm-5.1:cloud are installed so they sort normally;
// kimi-k2.5:cloud, qwen3.5:cloud, and qwen3.5 are not installed so they go to the bottom
// qwen3.5:cloud and qwen3.5 are not installed so they go to the bottom
// All recs: cloud first in mixed case, then local, in rec order within each
want := []string{"kimi-k2.5:cloud", "qwen3.5:cloud", "glm-5.1:cloud", "minimax-m2.7:cloud", "gemma4", "qwen3.5"}
want := recommendedNames()
if diff := cmp.Diff(want, got); diff != "" {
t.Errorf("all recs, cloud first in mixed case (-want +got):\n%s", diff)
}
@@ -456,23 +477,23 @@ func TestBuildModelList_ExistingCloudModelsNotPushedToBottom(t *testing.T) {
func TestBuildModelList_HasRecommendedCloudModel_OnlyNonInstalledAtBottom(t *testing.T) {
existing := []modelInfo{
{Name: "llama3.2:latest", Remote: false},
{Name: "kimi-k2.5:cloud", Remote: true},
{Name: "kimi-k2.6:cloud", Remote: true},
}
items, _, _, _ := buildModelList(existing, nil, "")
got := names(items)
// kimi-k2.5:cloud is installed so it sorts normally;
// kimi-k2.6:cloud is installed so it sorts normally;
// the rest of the recommendations are not installed so they go to the bottom
// All recs pinned at top (cloud first in mixed case), then non-recs
want := []string{"kimi-k2.5:cloud", "qwen3.5:cloud", "glm-5.1:cloud", "minimax-m2.7:cloud", "gemma4", "qwen3.5", "llama3.2"}
want := recommendedNames("llama3.2")
if diff := cmp.Diff(want, got); diff != "" {
t.Errorf("recs pinned at top, cloud first in mixed case (-want +got):\n%s", diff)
}
for _, item := range items {
isCloud := strings.HasSuffix(item.Name, ":cloud")
isInstalled := slices.Contains([]string{"kimi-k2.5:cloud", "llama3.2"}, item.Name)
isInstalled := slices.Contains([]string{"kimi-k2.6:cloud", "llama3.2"}, item.Name)
if isInstalled || isCloud {
if strings.HasSuffix(item.Description, "(not downloaded)") {
t.Errorf("installed or cloud model %q should not have '(not downloaded)' suffix, got %q", item.Name, item.Description)
@@ -539,8 +560,8 @@ func TestBuildModelList_ReturnsExistingAndCloudMaps(t *testing.T) {
if !cloudModels["glm-5.1:cloud"] {
t.Error("glm-5.1:cloud should be in cloudModels")
}
if !cloudModels["kimi-k2.5:cloud"] {
t.Error("kimi-k2.5:cloud should be in cloudModels (recommended cloud)")
if !cloudModels["kimi-k2.6:cloud"] {
t.Error("kimi-k2.6:cloud should be in cloudModels (recommended cloud)")
}
if !cloudModels["qwen3.5:cloud"] {
t.Error("qwen3.5:cloud should be in cloudModels (recommended cloud)")
@@ -560,7 +581,7 @@ func TestBuildModelList_RecommendedFieldSet(t *testing.T) {
for _, item := range items {
switch item.Name {
case "gemma4", "qwen3.5", "glm-5.1:cloud", "kimi-k2.5:cloud", "qwen3.5:cloud":
case "gemma4", "qwen3.5", "glm-5.1:cloud", "kimi-k2.6:cloud", "qwen3.5:cloud":
if !item.Recommended {
t.Errorf("%q should have Recommended=true", item.Name)
}
@@ -618,7 +639,7 @@ func TestBuildModelList_RecsAboveNonRecs(t *testing.T) {
lastRecIdx := -1
firstNonRecIdx := len(got)
for i, name := range got {
isRec := name == "gemma4" || name == "qwen3.5" || name == "minimax-m2.7:cloud" || name == "glm-5.1:cloud" || name == "kimi-k2.5:cloud" || name == "qwen3.5:cloud"
isRec := name == "gemma4" || name == "qwen3.5" || name == "minimax-m2.7:cloud" || name == "glm-5.1:cloud" || name == "kimi-k2.6:cloud" || name == "qwen3.5:cloud"
if isRec && i > lastRecIdx {
lastRecIdx = i
}
@@ -631,17 +652,32 @@ func TestBuildModelList_RecsAboveNonRecs(t *testing.T) {
}
}
func TestBuildModelList_CheckedBeforeRecs(t *testing.T) {
func TestBuildModelList_CheckedRecommendedDoesNotReshuffleRecommendedOrder(t *testing.T) {
existing := []modelInfo{
{Name: "llama3.2:latest", Remote: false},
{Name: "glm-5.1:cloud", Remote: true},
}
items, _, _, _ := buildModelList(existing, []string{"llama3.2"}, "")
items, _, _, _ := buildModelList(existing, []string{"qwen3.5:cloud", "glm-5.1:cloud"}, "")
got := names(items)
if got[0] != "llama3.2" {
t.Errorf("checked model should be first even before recs, got %v", got)
want := recommendedNames("llama3.2")
if diff := cmp.Diff(want, got); diff != "" {
t.Errorf("checked recommended models should not reshuffle the fixed recommended order (-want +got):\n%s", diff)
}
}
func TestBuildModelList_StaleSavedKimiK25DoesNotReshuffleRecommendedOrder(t *testing.T) {
existing := []modelInfo{
{Name: "kimi-k2.5:cloud", Remote: true},
}
items, _, _, _ := buildModelList(existing, []string{"kimi-k2.5:cloud", "qwen3.5:cloud", "glm-5.1:cloud", "minimax-m2.7:cloud"}, "kimi-k2.5:cloud")
got := names(items)
want := recommendedNames("kimi-k2.5:cloud")
if diff := cmp.Diff(want, got); diff != "" {
t.Errorf("stale saved kimi-k2.5 should stay in More without reshuffling the fixed recommended order (-want +got):\n%s", diff)
}
}
@@ -1330,6 +1366,30 @@ func TestEnsureAuth_SkipsWhenNoCloudSelected(t *testing.T) {
}
}
func TestEnsureAuth_EmptyWhoamiRequiresSignIn(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/status":
w.WriteHeader(http.StatusNotFound)
fmt.Fprintf(w, `{"error":"not found"}`)
case "/api/me":
w.WriteHeader(http.StatusOK)
fmt.Fprintf(w, `{}`)
default:
w.WriteHeader(http.StatusNotFound)
}
}))
defer srv.Close()
u, _ := url.Parse(srv.URL)
client := api.NewClient(u, srv.Client())
err := ensureAuth(context.Background(), client, map[string]bool{"cloud-model:cloud": true}, []string{"cloud-model:cloud"})
if err == nil || !strings.Contains(err.Error(), "cloud-model:cloud requires sign in") {
t.Fatalf("ensureAuth error = %v, want sign-in required", err)
}
}
func TestEnsureAuth_PreservesCancelledSignInHook(t *testing.T) {
oldSignIn := DefaultSignIn
DefaultSignIn = func(modelName, signInURL string) (string, error) {
@@ -1450,6 +1510,11 @@ func TestIntegration_InstallHint(t *testing.T) {
input: "claude",
wantURL: "https://code.claude.com/docs/en/quickstart",
},
{
name: "claude desktop has hint",
input: "claude-desktop",
wantURL: "https://claude.com/download",
},
{
name: "codex has hint",
input: "codex",
@@ -1460,6 +1525,11 @@ func TestIntegration_InstallHint(t *testing.T) {
input: "openclaw",
wantURL: "https://docs.openclaw.ai",
},
{
name: "pool has hint",
input: "pool",
wantURL: "https://github.com/poolsideai/pool",
},
{
name: "unknown has no hint",
input: "unknown",
@@ -1515,7 +1585,28 @@ func TestListIntegrationInfos(t *testing.T) {
for _, info := range infos {
got = append(got, info.Name)
}
if diff := compareStrings(got, integrationOrder); diff != "" {
want := append([]string(nil), integrationOrder...)
if poolsideGOOS == "windows" {
filtered := make([]string, 0, len(want))
for _, name := range want {
if name != "pool" {
filtered = append(filtered, name)
}
}
want = filtered
}
if claudeDesktopSupported() != nil {
filtered := make([]string, 0, len(want))
for _, name := range want {
if name != "claude-desktop" {
filtered = append(filtered, name)
}
}
want = filtered
}
if diff := compareStrings(got, want); diff != "" {
t.Fatalf("launcher integration order mismatch: %s", diff)
}
})
@@ -1533,6 +1624,12 @@ func TestListIntegrationInfos(t *testing.T) {
t.Run("includes known integrations", func(t *testing.T) {
known := map[string]bool{"claude": false, "codex": false, "opencode": false}
if claudeDesktopSupported() == nil {
known["claude-desktop"] = false
}
if poolsideGOOS != "windows" {
known["pool"] = false
}
for _, info := range infos {
if _, ok := known[info.Name]; ok {
known[info.Name] = true
@@ -1568,6 +1665,30 @@ func TestListIntegrationInfos(t *testing.T) {
})
}
func TestListIntegrationInfos_HidesPoolsideOnWindows(t *testing.T) {
prev := poolsideGOOS
poolsideGOOS = "windows"
t.Cleanup(func() { poolsideGOOS = prev })
for _, info := range ListIntegrationInfos() {
if info.Name == "pool" {
t.Fatal("expected pool to be hidden on Windows")
}
}
}
func TestListIntegrationInfos_HidesClaudeDesktopOnUnsupportedPlatform(t *testing.T) {
prev := claudeDesktopGOOS
claudeDesktopGOOS = "linux"
t.Cleanup(func() { claudeDesktopGOOS = prev })
for _, info := range ListIntegrationInfos() {
if info.Name == "claude-desktop" {
t.Fatal("expected claude-desktop to be hidden on unsupported platforms")
}
}
}
func TestBuildModelList_Descriptions(t *testing.T) {
t.Run("installed recommended has base description", func(t *testing.T) {
existing := []modelInfo{
@@ -1594,7 +1715,7 @@ func TestBuildModelList_Descriptions(t *testing.T) {
for _, item := range items {
if item.Name == "qwen3.5" {
if !strings.Contains(item.Description, "~11GB") {
if !strings.Contains(item.Description, "~14GB") {
t.Errorf("not-installed qwen3.5 should show VRAM hint, got %q", item.Description)
}
return
@@ -1611,7 +1732,7 @@ func TestBuildModelList_Descriptions(t *testing.T) {
for _, item := range items {
if item.Name == "qwen3.5" {
if strings.Contains(item.Description, "~11GB") {
if strings.Contains(item.Description, "~14GB") {
t.Errorf("installed qwen3.5 should not show VRAM hint, got %q", item.Description)
}
return
@@ -1630,6 +1751,7 @@ func TestIntegration_Editor(t *testing.T) {
{"opencode", true},
{"openclaw", true},
{"claude", false},
{"claude-desktop", false},
{"codex", false},
{"nonexistent", false},
}
@@ -1656,6 +1778,7 @@ func TestIntegration_AutoInstallable(t *testing.T) {
{"pi", true},
{"hermes", true},
{"claude", false},
{"claude-desktop", false},
{"codex", false},
{"opencode", false},
}
@@ -1673,6 +1796,20 @@ func TestIntegration_AutoInstallable(t *testing.T) {
}
}
func TestEnsureIntegrationInstalled_PoolsideUnsupportedOnWindows(t *testing.T) {
prev := poolsideGOOS
poolsideGOOS = "windows"
t.Cleanup(func() { poolsideGOOS = prev })
err := EnsureIntegrationInstalled("pool", &Poolside{})
if err == nil {
t.Fatal("expected Windows unsupported error")
}
if !strings.Contains(err.Error(), "not currently supported on Windows") {
t.Fatalf("expected Windows warning, got %v", err)
}
}
func TestIntegrationModels(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
+315
View File
@@ -0,0 +1,315 @@
package launch
import (
"context"
"encoding/json"
"fmt"
"os"
"os/exec"
"path/filepath"
"runtime"
"strings"
"time"
"github.com/ollama/ollama/api"
"github.com/ollama/ollama/envconfig"
)
// Kimi implements Runner for Kimi Code CLI integration.
type Kimi struct{}
const (
kimiDefaultModelAlias = "ollama"
kimiDefaultMaxContextSize = 32768
)
var (
kimiGOOS = runtime.GOOS
kimiModelShowTimeout = 5 * time.Second
)
func (k *Kimi) String() string { return "Kimi Code CLI" }
func (k *Kimi) args(config string, extra []string) []string {
args := []string{"--config", config}
args = append(args, extra...)
return args
}
func (k *Kimi) Run(model string, args []string) error {
if strings.TrimSpace(model) == "" {
return fmt.Errorf("model is required")
}
if err := validateKimiPassthroughArgs(args); err != nil {
return err
}
config, err := buildKimiInlineConfig(model, resolveKimiMaxContextSize(model))
if err != nil {
return fmt.Errorf("failed to build kimi config: %w", err)
}
bin, err := ensureKimiInstalled()
if err != nil {
return err
}
cmd := exec.Command(bin, k.args(config, args)...)
cmd.Stdin = os.Stdin
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
return cmd.Run()
}
func findKimiBinary() (string, error) {
if path, err := exec.LookPath("kimi"); err == nil {
return path, nil
}
home, _ := os.UserHomeDir()
var candidates []string
switch kimiGOOS {
case "windows":
candidates = appendWindowsKimiCandidates(candidates, filepath.Join(home, ".local", "bin"))
candidates = appendWindowsKimiCandidates(candidates, filepath.Join(home, "bin"))
if appData := strings.TrimSpace(os.Getenv("APPDATA")); appData != "" {
candidates = appendWindowsKimiCandidates(candidates, filepath.Join(appData, "uv", "bin"))
}
if localAppData := strings.TrimSpace(os.Getenv("LOCALAPPDATA")); localAppData != "" {
candidates = appendWindowsKimiCandidates(candidates, filepath.Join(localAppData, "uv", "bin"))
}
default:
candidates = append(candidates,
filepath.Join(home, ".local", "bin", "kimi"),
filepath.Join(home, "bin", "kimi"),
filepath.Join(home, ".local", "share", "uv", "tools", "kimi-cli", "bin", "kimi"),
filepath.Join(home, ".local", "share", "uv", "tools", "kimi", "bin", "kimi"),
)
if xdgDataHome := strings.TrimSpace(os.Getenv("XDG_DATA_HOME")); xdgDataHome != "" {
candidates = append(candidates,
filepath.Join(xdgDataHome, "uv", "tools", "kimi-cli", "bin", "kimi"),
filepath.Join(xdgDataHome, "uv", "tools", "kimi", "bin", "kimi"),
)
}
// WSL users can inherit Windows env vars while launching from Linux shells.
if profile := windowsPathToWSL(os.Getenv("USERPROFILE")); profile != "" {
candidates = appendWindowsKimiCandidates(candidates, filepath.Join(profile, ".local", "bin"))
}
if appData := windowsPathToWSL(os.Getenv("APPDATA")); appData != "" {
candidates = appendWindowsKimiCandidates(candidates, filepath.Join(appData, "uv", "bin"))
}
if localAppData := windowsPathToWSL(os.Getenv("LOCALAPPDATA")); localAppData != "" {
candidates = appendWindowsKimiCandidates(candidates, filepath.Join(localAppData, "uv", "bin"))
}
}
for _, candidate := range candidates {
if info, err := os.Stat(candidate); err == nil && !info.IsDir() {
return candidate, nil
}
}
return "", fmt.Errorf("kimi binary not found")
}
func appendWindowsKimiCandidates(candidates []string, dir string) []string {
if strings.TrimSpace(dir) == "" {
return candidates
}
return append(candidates,
filepath.Join(dir, "kimi.exe"),
filepath.Join(dir, "kimi.cmd"),
filepath.Join(dir, "kimi.bat"),
)
}
func windowsPathToWSL(path string) string {
trimmed := strings.TrimSpace(path)
if len(trimmed) < 3 || trimmed[1] != ':' {
return ""
}
drive := strings.ToLower(string(trimmed[0]))
rest := strings.ReplaceAll(trimmed[2:], "\\", "/")
rest = strings.TrimPrefix(rest, "/")
if rest == "" {
return filepath.Join("/mnt", drive)
}
return filepath.Join("/mnt", drive, rest)
}
func validateKimiPassthroughArgs(args []string) error {
for _, arg := range args {
switch {
case arg == "--config", strings.HasPrefix(arg, "--config="):
return fmt.Errorf("conflicting extra argument %q: ollama launch kimi manages --config", arg)
case arg == "--config-file", strings.HasPrefix(arg, "--config-file="):
return fmt.Errorf("conflicting extra argument %q: ollama launch kimi manages --config-file", arg)
case arg == "--model", strings.HasPrefix(arg, "--model="):
return fmt.Errorf("conflicting extra argument %q: ollama launch kimi manages --model", arg)
case arg == "-m", strings.HasPrefix(arg, "-m="):
return fmt.Errorf("conflicting extra argument %q: ollama launch kimi manages -m/--model", arg)
}
}
return nil
}
func buildKimiInlineConfig(model string, maxContextSize int) (string, error) {
cfg := map[string]any{
"default_model": kimiDefaultModelAlias,
"providers": map[string]any{
kimiDefaultModelAlias: map[string]any{
"type": "openai_legacy",
"base_url": envconfig.ConnectableHost().String() + "/v1",
"api_key": "ollama",
},
},
"models": map[string]any{
kimiDefaultModelAlias: map[string]any{
"provider": kimiDefaultModelAlias,
"model": model,
"max_context_size": maxContextSize,
},
},
}
data, err := json.Marshal(cfg)
if err != nil {
return "", err
}
return string(data), nil
}
func resolveKimiMaxContextSize(model string) int {
if l, ok := lookupCloudModelLimit(model); ok {
return l.Context
}
client, err := api.ClientFromEnvironment()
if err != nil {
return kimiDefaultMaxContextSize
}
ctx, cancel := context.WithTimeout(context.Background(), kimiModelShowTimeout)
defer cancel()
resp, err := client.Show(ctx, &api.ShowRequest{Model: model})
if err != nil {
return kimiDefaultMaxContextSize
}
if n, ok := modelInfoContextLength(resp.ModelInfo); ok {
return n
}
return kimiDefaultMaxContextSize
}
func modelInfoContextLength(modelInfo map[string]any) (int, bool) {
for key, val := range modelInfo {
if !strings.HasSuffix(key, ".context_length") {
continue
}
switch v := val.(type) {
case float64:
if v > 0 {
return int(v), true
}
case int:
if v > 0 {
return v, true
}
case int64:
if v > 0 {
return int(v), true
}
}
}
return 0, false
}
func ensureKimiInstalled() (string, error) {
if path, err := findKimiBinary(); err == nil {
return path, nil
}
if err := checkKimiInstallerDependencies(); err != nil {
return "", err
}
ok, err := ConfirmPrompt("Kimi is not installed. Install now?")
if err != nil {
return "", err
}
if !ok {
return "", fmt.Errorf("kimi installation cancelled")
}
bin, args, err := kimiInstallerCommand(kimiGOOS)
if err != nil {
return "", err
}
fmt.Fprintf(os.Stderr, "\nInstalling Kimi...\n")
cmd := exec.Command(bin, args...)
cmd.Stdin = os.Stdin
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
if err := cmd.Run(); err != nil {
return "", fmt.Errorf("failed to install kimi: %w", err)
}
path, err := findKimiBinary()
if err != nil {
return "", fmt.Errorf("kimi was installed but the binary was not found on PATH\n\nYou may need to restart your shell")
}
fmt.Fprintf(os.Stderr, "%sKimi installed successfully%s\n\n", ansiGreen, ansiReset)
return path, nil
}
func checkKimiInstallerDependencies() error {
switch kimiGOOS {
case "windows":
if _, err := exec.LookPath("powershell"); err != nil {
return fmt.Errorf("kimi is not installed and required dependencies are missing\n\nInstall the following first:\n PowerShell: https://learn.microsoft.com/powershell/\n\nThen re-run:\n ollama launch kimi")
}
default:
var missing []string
if _, err := exec.LookPath("curl"); err != nil {
missing = append(missing, "curl: https://curl.se/")
}
if _, err := exec.LookPath("bash"); err != nil {
missing = append(missing, "bash: https://www.gnu.org/software/bash/")
}
if len(missing) > 0 {
return fmt.Errorf("kimi is not installed and required dependencies are missing\n\nInstall the following first:\n %s\n\nThen re-run:\n ollama launch kimi", strings.Join(missing, "\n "))
}
}
return nil
}
func kimiInstallerCommand(goos string) (string, []string, error) {
switch goos {
case "windows":
return "powershell", []string{
"-NoProfile",
"-ExecutionPolicy",
"Bypass",
"-Command",
"Invoke-RestMethod https://code.kimi.com/install.ps1 | Invoke-Expression",
}, nil
case "darwin", "linux":
return "bash", []string{
"-c",
"curl -LsSf https://code.kimi.com/install.sh | bash",
}, nil
default:
return "", nil, fmt.Errorf("unsupported platform for kimi install: %s", goos)
}
}
+636
View File
@@ -0,0 +1,636 @@
package launch
import (
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"runtime"
"slices"
"strings"
"testing"
)
func assertKimiBinPath(t *testing.T, bin string) {
t.Helper()
base := strings.ToLower(filepath.Base(bin))
if !strings.HasPrefix(base, "kimi") {
t.Fatalf("bin = %q, want path to kimi executable", bin)
}
}
func TestKimiIntegration(t *testing.T) {
k := &Kimi{}
t.Run("String", func(t *testing.T) {
if got := k.String(); got != "Kimi Code CLI" {
t.Errorf("String() = %q, want %q", got, "Kimi Code CLI")
}
})
t.Run("implements Runner", func(t *testing.T) {
var _ Runner = k
})
}
func TestKimiArgs(t *testing.T) {
k := &Kimi{}
got := k.args(`{"foo":"bar"}`, []string{"--quiet", "--print"})
want := []string{"--config", `{"foo":"bar"}`, "--quiet", "--print"}
if !slices.Equal(got, want) {
t.Fatalf("args() = %v, want %v", got, want)
}
}
func TestWindowsPathToWSL(t *testing.T) {
tests := []struct {
name string
in string
want string
valid bool
}{
{
name: "user profile path",
in: `C:\Users\parth`,
want: filepath.Join("/mnt", "c", "Users", "parth"),
valid: true,
},
{
name: "path with trailing slash",
in: `D:\tools\bin\`,
want: filepath.Join("/mnt", "d", "tools", "bin"),
valid: true,
},
{
name: "non windows path",
in: "/home/parth",
valid: false,
},
{
name: "empty",
in: "",
valid: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := windowsPathToWSL(tt.in)
if !tt.valid {
if got != "" {
t.Fatalf("windowsPathToWSL(%q) = %q, want empty", tt.in, got)
}
return
}
if got != tt.want {
t.Fatalf("windowsPathToWSL(%q) = %q, want %q", tt.in, got, tt.want)
}
})
}
}
func TestFindKimiBinaryFallbacks(t *testing.T) {
oldGOOS := kimiGOOS
t.Cleanup(func() { kimiGOOS = oldGOOS })
t.Run("linux/ubuntu uv tool path", func(t *testing.T) {
homeDir := t.TempDir()
setTestHome(t, homeDir)
t.Setenv("PATH", t.TempDir())
kimiGOOS = "linux"
target := filepath.Join(homeDir, ".local", "share", "uv", "tools", "kimi-cli", "bin", "kimi")
if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
t.Fatalf("failed to create candidate dir: %v", err)
}
if err := os.WriteFile(target, []byte("#!/bin/sh\nexit 0\n"), 0o755); err != nil {
t.Fatalf("failed to write kimi candidate: %v", err)
}
got, err := findKimiBinary()
if err != nil {
t.Fatalf("findKimiBinary() error = %v", err)
}
if got != target {
t.Fatalf("findKimiBinary() = %q, want %q", got, target)
}
})
t.Run("windows appdata uv bin", func(t *testing.T) {
setTestHome(t, t.TempDir())
t.Setenv("PATH", t.TempDir())
kimiGOOS = "windows"
appDataDir := t.TempDir()
t.Setenv("APPDATA", appDataDir)
t.Setenv("LOCALAPPDATA", "")
target := filepath.Join(appDataDir, "uv", "bin", "kimi.cmd")
if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
t.Fatalf("failed to create candidate dir: %v", err)
}
if err := os.WriteFile(target, []byte("@echo off\r\nexit /b 0\r\n"), 0o755); err != nil {
t.Fatalf("failed to write kimi candidate: %v", err)
}
got, err := findKimiBinary()
if err != nil {
t.Fatalf("findKimiBinary() error = %v", err)
}
if got != target {
t.Fatalf("findKimiBinary() = %q, want %q", got, target)
}
})
}
func TestValidateKimiPassthroughArgs_RejectsConflicts(t *testing.T) {
tests := []struct {
name string
args []string
want string
}{
{name: "--config", args: []string{"--config", "{}"}, want: "--config"},
{name: "--config=", args: []string{"--config={}"}, want: "--config={"},
{name: "--config-file", args: []string{"--config-file", "x.toml"}, want: "--config-file"},
{name: "--config-file=", args: []string{"--config-file=x.toml"}, want: "--config-file=x.toml"},
{name: "--model", args: []string{"--model", "foo"}, want: "--model"},
{name: "--model=", args: []string{"--model=foo"}, want: "--model=foo"},
{name: "-m", args: []string{"-m", "foo"}, want: "-m"},
{name: "-m=", args: []string{"-m=foo"}, want: "-m=foo"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := validateKimiPassthroughArgs(tt.args)
if err == nil {
t.Fatalf("expected error for args %v", tt.args)
}
if !strings.Contains(err.Error(), tt.want) {
t.Fatalf("error %q does not contain %q", err.Error(), tt.want)
}
})
}
}
func TestBuildKimiInlineConfig(t *testing.T) {
t.Setenv("OLLAMA_HOST", "http://127.0.0.1:11434")
cfg, err := buildKimiInlineConfig("llama3.2", 65536)
if err != nil {
t.Fatalf("buildKimiInlineConfig() error = %v", err)
}
var parsed map[string]any
if err := json.Unmarshal([]byte(cfg), &parsed); err != nil {
t.Fatalf("config is not valid JSON: %v", err)
}
if parsed["default_model"] != "ollama" {
t.Fatalf("default_model = %v, want ollama", parsed["default_model"])
}
providers, ok := parsed["providers"].(map[string]any)
if !ok {
t.Fatalf("providers missing or wrong type: %T", parsed["providers"])
}
ollamaProvider, ok := providers["ollama"].(map[string]any)
if !ok {
t.Fatalf("providers.ollama missing or wrong type: %T", providers["ollama"])
}
if ollamaProvider["type"] != "openai_legacy" {
t.Fatalf("provider type = %v, want openai_legacy", ollamaProvider["type"])
}
if ollamaProvider["base_url"] != "http://127.0.0.1:11434/v1" {
t.Fatalf("provider base_url = %v, want http://127.0.0.1:11434/v1", ollamaProvider["base_url"])
}
if ollamaProvider["api_key"] != "ollama" {
t.Fatalf("provider api_key = %v, want ollama", ollamaProvider["api_key"])
}
models, ok := parsed["models"].(map[string]any)
if !ok {
t.Fatalf("models missing or wrong type: %T", parsed["models"])
}
ollamaModel, ok := models["ollama"].(map[string]any)
if !ok {
t.Fatalf("models.ollama missing or wrong type: %T", models["ollama"])
}
if ollamaModel["provider"] != "ollama" {
t.Fatalf("model provider = %v, want ollama", ollamaModel["provider"])
}
if ollamaModel["model"] != "llama3.2" {
t.Fatalf("model model = %v, want llama3.2", ollamaModel["model"])
}
if ollamaModel["max_context_size"] != float64(65536) {
t.Fatalf("model max_context_size = %v, want 65536", ollamaModel["max_context_size"])
}
}
func TestBuildKimiInlineConfig_UsesConnectableHostForUnspecifiedBind(t *testing.T) {
t.Setenv("OLLAMA_HOST", "http://0.0.0.0:11434")
cfg, err := buildKimiInlineConfig("llama3.2", 65536)
if err != nil {
t.Fatalf("buildKimiInlineConfig() error = %v", err)
}
var parsed map[string]any
if err := json.Unmarshal([]byte(cfg), &parsed); err != nil {
t.Fatalf("config is not valid JSON: %v", err)
}
providers, ok := parsed["providers"].(map[string]any)
if !ok {
t.Fatalf("providers missing or wrong type: %T", parsed["providers"])
}
ollamaProvider, ok := providers["ollama"].(map[string]any)
if !ok {
t.Fatalf("providers.ollama missing or wrong type: %T", providers["ollama"])
}
if got, _ := ollamaProvider["base_url"].(string); got != "http://127.0.0.1:11434/v1" {
t.Fatalf("provider base_url = %q, want %q", got, "http://127.0.0.1:11434/v1")
}
}
func TestResolveKimiMaxContextSize(t *testing.T) {
t.Run("uses cloud limit when known", func(t *testing.T) {
got := resolveKimiMaxContextSize("kimi-k2.5:cloud")
if got != 262_144 {
t.Fatalf("resolveKimiMaxContextSize() = %d, want 262144", got)
}
})
t.Run("uses model show context length for local models", func(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/show" {
http.NotFound(w, r)
return
}
fmt.Fprint(w, `{"model_info":{"llama.context_length":131072}}`)
}))
defer srv.Close()
t.Setenv("OLLAMA_HOST", srv.URL)
got := resolveKimiMaxContextSize("llama3.2")
if got != 131_072 {
t.Fatalf("resolveKimiMaxContextSize() = %d, want 131072", got)
}
})
t.Run("falls back to default when show fails", func(t *testing.T) {
srv := httptest.NewServer(http.NotFoundHandler())
defer srv.Close()
t.Setenv("OLLAMA_HOST", srv.URL)
oldTimeout := kimiModelShowTimeout
kimiModelShowTimeout = 100 * 1000 * 1000 // 100ms
t.Cleanup(func() { kimiModelShowTimeout = oldTimeout })
got := resolveKimiMaxContextSize("llama3.2")
if got != kimiDefaultMaxContextSize {
t.Fatalf("resolveKimiMaxContextSize() = %d, want %d", got, kimiDefaultMaxContextSize)
}
})
}
func TestKimiRun_RejectsConflictingArgsBeforeInstall(t *testing.T) {
k := &Kimi{}
oldConfirm := DefaultConfirmPrompt
DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) {
t.Fatalf("did not expect install prompt, got %q", prompt)
return false, nil
}
t.Cleanup(func() { DefaultConfirmPrompt = oldConfirm })
err := k.Run("llama3.2", []string{"--model", "other"})
if err == nil || !strings.Contains(err.Error(), "--model") {
t.Fatalf("expected conflict error mentioning --model, got %v", err)
}
}
func TestKimiRun_PassesInlineConfigAndExtraArgs(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses POSIX shell fake binary")
}
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
logPath := filepath.Join(tmpDir, "kimi-args.log")
script := fmt.Sprintf(`#!/bin/sh
for arg in "$@"; do
printf "%%s\n" "$arg" >> %q
done
exit 0
`, logPath)
if err := os.WriteFile(filepath.Join(tmpDir, "kimi"), []byte(script), 0o755); err != nil {
t.Fatalf("failed to write fake kimi: %v", err)
}
t.Setenv("PATH", tmpDir)
srv := httptest.NewServer(http.NotFoundHandler())
defer srv.Close()
t.Setenv("OLLAMA_HOST", srv.URL)
k := &Kimi{}
if err := k.Run("llama3.2", []string{"--quiet", "--print"}); err != nil {
t.Fatalf("Run() error = %v", err)
}
data, err := os.ReadFile(logPath)
if err != nil {
t.Fatalf("failed to read args log: %v", err)
}
lines := strings.Split(strings.TrimSpace(string(data)), "\n")
if len(lines) < 4 {
t.Fatalf("expected at least 4 args, got %v", lines)
}
if lines[0] != "--config" {
t.Fatalf("first arg = %q, want --config", lines[0])
}
var cfg map[string]any
if err := json.Unmarshal([]byte(lines[1]), &cfg); err != nil {
t.Fatalf("config arg is not valid JSON: %v", err)
}
providers := cfg["providers"].(map[string]any)
ollamaProvider := providers["ollama"].(map[string]any)
if ollamaProvider["type"] != "openai_legacy" {
t.Fatalf("provider type = %v, want openai_legacy", ollamaProvider["type"])
}
if lines[2] != "--quiet" || lines[3] != "--print" {
t.Fatalf("extra args = %v, want [--quiet --print]", lines[2:])
}
}
func TestEnsureKimiInstalled(t *testing.T) {
oldGOOS := kimiGOOS
t.Cleanup(func() { kimiGOOS = oldGOOS })
withConfirm := func(t *testing.T, fn func(prompt string) (bool, error)) {
t.Helper()
oldConfirm := DefaultConfirmPrompt
DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) {
return fn(prompt)
}
t.Cleanup(func() { DefaultConfirmPrompt = oldConfirm })
}
t.Run("already installed", func(t *testing.T) {
setTestHome(t, t.TempDir())
tmpDir := t.TempDir()
t.Setenv("PATH", tmpDir)
writeFakeBinary(t, tmpDir, "kimi")
kimiGOOS = runtime.GOOS
withConfirm(t, func(prompt string) (bool, error) {
t.Fatalf("did not expect prompt, got %q", prompt)
return false, nil
})
bin, err := ensureKimiInstalled()
if err != nil {
t.Fatalf("ensureKimiInstalled() error = %v", err)
}
assertKimiBinPath(t, bin)
})
t.Run("missing dependencies", func(t *testing.T) {
setTestHome(t, t.TempDir())
tmpDir := t.TempDir()
t.Setenv("PATH", tmpDir)
kimiGOOS = "linux"
withConfirm(t, func(prompt string) (bool, error) {
t.Fatalf("did not expect prompt, got %q", prompt)
return false, nil
})
_, err := ensureKimiInstalled()
if err == nil || !strings.Contains(err.Error(), "required dependencies are missing") {
t.Fatalf("expected missing dependency error, got %v", err)
}
})
t.Run("missing and user declines install", func(t *testing.T) {
setTestHome(t, t.TempDir())
tmpDir := t.TempDir()
t.Setenv("PATH", tmpDir)
writeFakeBinary(t, tmpDir, "curl")
writeFakeBinary(t, tmpDir, "bash")
kimiGOOS = "linux"
withConfirm(t, func(prompt string) (bool, error) {
if !strings.Contains(prompt, "Kimi is not installed.") {
t.Fatalf("unexpected prompt: %q", prompt)
}
return false, nil
})
_, err := ensureKimiInstalled()
if err == nil || !strings.Contains(err.Error(), "installation cancelled") {
t.Fatalf("expected cancellation error, got %v", err)
}
})
t.Run("missing and user confirms install succeeds", func(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses POSIX shell fake binaries")
}
setTestHome(t, t.TempDir())
tmpDir := t.TempDir()
t.Setenv("PATH", tmpDir)
kimiGOOS = "linux"
writeFakeBinary(t, tmpDir, "curl")
installLog := filepath.Join(tmpDir, "bash.log")
kimiPath := filepath.Join(tmpDir, "kimi")
bashScript := fmt.Sprintf(`#!/bin/sh
echo "$@" >> %q
if [ "$1" = "-c" ]; then
/bin/cat > %q <<'EOS'
#!/bin/sh
exit 0
EOS
/bin/chmod +x %q
fi
exit 0
`, installLog, kimiPath, kimiPath)
if err := os.WriteFile(filepath.Join(tmpDir, "bash"), []byte(bashScript), 0o755); err != nil {
t.Fatalf("failed to write fake bash: %v", err)
}
withConfirm(t, func(prompt string) (bool, error) {
return true, nil
})
bin, err := ensureKimiInstalled()
if err != nil {
t.Fatalf("ensureKimiInstalled() error = %v", err)
}
assertKimiBinPath(t, bin)
logData, err := os.ReadFile(installLog)
if err != nil {
t.Fatalf("failed to read install log: %v", err)
}
if !strings.Contains(string(logData), "https://code.kimi.com/install.sh") {
t.Fatalf("expected install.sh command in log, got:\n%s", string(logData))
}
})
t.Run("install succeeds and kimi is in home local bin without PATH update", func(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses POSIX shell fake binaries")
}
homeDir := t.TempDir()
setTestHome(t, homeDir)
tmpBin := t.TempDir()
t.Setenv("PATH", tmpBin)
kimiGOOS = "linux"
writeFakeBinary(t, tmpBin, "curl")
installedKimi := filepath.Join(homeDir, ".local", "bin", "kimi")
bashScript := fmt.Sprintf(`#!/bin/sh
if [ "$1" = "-c" ]; then
/bin/mkdir -p %q
/bin/cat > %q <<'EOS'
#!/bin/sh
exit 0
EOS
/bin/chmod +x %q
fi
exit 0
`, filepath.Dir(installedKimi), installedKimi, installedKimi)
if err := os.WriteFile(filepath.Join(tmpBin, "bash"), []byte(bashScript), 0o755); err != nil {
t.Fatalf("failed to write fake bash: %v", err)
}
withConfirm(t, func(prompt string) (bool, error) {
return true, nil
})
bin, err := ensureKimiInstalled()
if err != nil {
t.Fatalf("ensureKimiInstalled() error = %v", err)
}
if bin != installedKimi {
t.Fatalf("bin = %q, want %q", bin, installedKimi)
}
})
t.Run("install command fails", func(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses POSIX shell fake binaries")
}
setTestHome(t, t.TempDir())
tmpDir := t.TempDir()
t.Setenv("PATH", tmpDir)
kimiGOOS = "linux"
writeFakeBinary(t, tmpDir, "curl")
if err := os.WriteFile(filepath.Join(tmpDir, "bash"), []byte("#!/bin/sh\nexit 1\n"), 0o755); err != nil {
t.Fatalf("failed to write fake bash: %v", err)
}
withConfirm(t, func(prompt string) (bool, error) {
return true, nil
})
_, err := ensureKimiInstalled()
if err == nil || !strings.Contains(err.Error(), "failed to install kimi") {
t.Fatalf("expected install failure error, got %v", err)
}
})
t.Run("install succeeds but binary missing on PATH", func(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses POSIX shell fake binaries")
}
setTestHome(t, t.TempDir())
tmpDir := t.TempDir()
t.Setenv("PATH", tmpDir)
kimiGOOS = "linux"
writeFakeBinary(t, tmpDir, "curl")
if err := os.WriteFile(filepath.Join(tmpDir, "bash"), []byte("#!/bin/sh\nexit 0\n"), 0o755); err != nil {
t.Fatalf("failed to write fake bash: %v", err)
}
withConfirm(t, func(prompt string) (bool, error) {
return true, nil
})
_, err := ensureKimiInstalled()
if err == nil || !strings.Contains(err.Error(), "binary was not found on PATH") {
t.Fatalf("expected PATH guidance error, got %v", err)
}
})
}
func TestKimiInstallerCommand(t *testing.T) {
tests := []struct {
name string
goos string
wantBin string
wantParts []string
wantErr bool
}{
{
name: "linux",
goos: "linux",
wantBin: "bash",
wantParts: []string{"-c", "install.sh"},
},
{
name: "darwin",
goos: "darwin",
wantBin: "bash",
wantParts: []string{"-c", "install.sh"},
},
{
name: "windows",
goos: "windows",
wantBin: "powershell",
wantParts: []string{"-Command", "install.ps1"},
},
{
name: "unsupported",
goos: "freebsd",
wantErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
bin, args, err := kimiInstallerCommand(tt.goos)
if tt.wantErr {
if err == nil {
t.Fatal("expected error")
}
return
}
if err != nil {
t.Fatalf("kimiInstallerCommand() error = %v", err)
}
if bin != tt.wantBin {
t.Fatalf("bin = %q, want %q", bin, tt.wantBin)
}
joined := strings.Join(args, " ")
for _, part := range tt.wantParts {
if !strings.Contains(joined, part) {
t.Fatalf("args %q missing %q", joined, part)
}
}
})
}
}
+383 -31
View File
@@ -4,6 +4,7 @@ import (
"context"
"errors"
"fmt"
"log/slog"
"net/http"
"os"
"slices"
@@ -120,6 +121,7 @@ type IntegrationLaunchRequest struct {
ModelOverride string
ForceConfigure bool
ConfigureOnly bool
Restore bool
ExtraArgs []string
Policy *LaunchPolicy
}
@@ -153,6 +155,46 @@ type ManagedSingleModel interface {
Onboard() error
}
// ManagedModelListConfigurer lets managed single-model integrations receive
// the launcher's model list while still preserving one primary selected model.
type ManagedModelListConfigurer interface {
ConfigureWithModels(primary string, models []string) error
}
// ManagedAutodiscoveryIntegration is for managed integrations that do not need
// a launcher-selected model because the app discovers available models itself.
type ManagedAutodiscoveryIntegration interface {
Paths() []string
AutodiscoveredModel() string
AutodiscoveryConfigured() bool
ConfigureAutodiscovery() error
Onboard() error
}
// ManagedAutodiscoveryCloudIntegration marks an autodiscovery integration whose
// discovered model catalog depends on the user's local Ollama Cloud auth state.
type ManagedAutodiscoveryCloudIntegration interface {
UsesOllamaCloud() bool
}
// RestoreHintIntegration can provide a short restore command after launch
// switches an app into a launch-managed mode.
type RestoreHintIntegration interface {
RestoreHint() string
}
// ConfigurationSuccessIntegration can print a short message after launcher
// successfully switches an app into a launch-managed mode.
type ConfigurationSuccessIntegration interface {
ConfigurationSuccessMessage() string
}
// RestoreSuccessIntegration can print a short message after launcher restores
// an app back to its default mode.
type RestoreSuccessIntegration interface {
RestoreSuccessMessage() string
}
// ManagedRuntimeRefresher lets managed integrations refresh any long-lived
// background runtime after launch rewrites their config.
type ManagedRuntimeRefresher interface {
@@ -171,6 +213,25 @@ type ManagedInteractiveOnboarding interface {
RequiresInteractiveOnboarding() bool
}
// ManagedModelReadinessSkipper lets managed integrations opt out of local
// Ollama model readiness checks when the configured runtime is not the local
// daemon.
type ManagedModelReadinessSkipper interface {
SkipModelReadiness() bool
}
// RestorableIntegration lets integrations switch back from a launch-managed
// mode to the application's normal/default mode.
type RestorableIntegration interface {
Restore() error
}
// SupportedIntegration lets an integration report platform support separately
// from whether the underlying app binary is installed.
type SupportedIntegration interface {
Supported() error
}
type modelInfo struct {
Name string
Remote bool
@@ -182,9 +243,12 @@ type ModelInfo = modelInfo
// ModelItem represents a model for selection UIs.
type ModelItem struct {
Name string
Description string
Recommended bool
Name string
Description string
Recommended bool
VRAMBytes int64
ContextLength int
MaxOutputTokens int
}
// LaunchCmd returns the cobra command for launching integrations.
@@ -193,6 +257,7 @@ func LaunchCmd(checkServerHeartbeat func(cmd *cobra.Command, args []string) erro
var modelFlag string
var configFlag bool
var yesFlag bool
var restoreFlag bool
cmd := &cobra.Command{
Use: "launch [INTEGRATION] [-- [EXTRA_ARGS...]]",
@@ -203,26 +268,37 @@ Without arguments, this is equivalent to running 'ollama' directly.
Flags and extra arguments require an integration name.
Supported integrations:
claude Claude Code
cline Cline
codex Codex
droid Droid
hermes Hermes Agent
opencode OpenCode
openclaw OpenClaw (aliases: clawdbot, moltbot)
pi Pi
vscode    VS Code (aliases: code)
claude Claude Code
claude-desktop Claude Desktop (aliases: claude-app)
cline Cline
codex Codex
copilot Copilot CLI (aliases: copilot-cli)
droid Droid
hermes Hermes Agent
kimi Kimi Code CLI
opencode OpenCode
openclaw OpenClaw (aliases: clawdbot, moltbot)
pi Pi
pool Pool
vscode VS Code (aliases: code)
Examples:
ollama launch
ollama launch claude
ollama launch claude --model <model>
ollama launch claude-desktop
ollama launch claude-desktop --restore
ollama launch hermes
ollama launch droid --config (does not auto-launch)
ollama launch codex -- -p myprofile (pass extra args to integration)
ollama launch codex -- --sandbox workspace-write`,
Args: cobra.ArbitraryArgs,
PreRunE: checkServerHeartbeat,
Args: cobra.ArbitraryArgs,
PreRunE: func(cmd *cobra.Command, args []string) error {
if restoreFlag || launchCommandCanSkipHeartbeat(args) {
return nil
}
return checkServerHeartbeat(cmd, args)
},
RunE: func(cmd *cobra.Command, args []string) error {
policy := defaultLaunchPolicy(isInteractiveSession(), yesFlag)
// reset when done to make sure state doens't leak between launches
@@ -251,7 +327,7 @@ Examples:
}
if name == "" {
if cmd.Flags().Changed("model") || cmd.Flags().Changed("config") || cmd.Flags().Changed("yes") || len(passArgs) > 0 {
if cmd.Flags().Changed("model") || cmd.Flags().Changed("config") || cmd.Flags().Changed("yes") || cmd.Flags().Changed("restore") || len(passArgs) > 0 {
return fmt.Errorf("flags and extra args require an integration name, for example: 'ollama launch claude --model qwen3.5'")
}
runTUI(cmd)
@@ -268,11 +344,20 @@ Examples:
}
headlessYes := yesFlag && !isInteractiveSession()
forceConfigure := configFlag || (modelFlag == "" && !headlessYes)
if forceConfigure && !configFlag && modelFlag == "" {
if _, runner, err := LookupIntegration(name); err == nil {
if _, ok := runner.(ManagedAutodiscoveryIntegration); ok {
forceConfigure = false
}
}
}
err := LaunchIntegration(cmd.Context(), IntegrationLaunchRequest{
Name: name,
ModelOverride: modelFlag,
ForceConfigure: configFlag || (modelFlag == "" && !headlessYes),
ForceConfigure: forceConfigure,
ConfigureOnly: configFlag,
Restore: restoreFlag,
ExtraArgs: passArgs,
Policy: &policy,
})
@@ -285,10 +370,19 @@ Examples:
cmd.Flags().StringVar(&modelFlag, "model", "", "Model to use")
cmd.Flags().BoolVar(&configFlag, "config", false, "Configure without launching")
cmd.Flags().BoolVar(&restoreFlag, "restore", false, "Restore an integration to its default profile")
cmd.Flags().BoolVarP(&yesFlag, "yes", "y", false, "Automatically answer yes to confirmation prompts")
return cmd
}
func launchCommandCanSkipHeartbeat(args []string) bool {
if len(args) == 0 {
return false
}
name, _, err := LookupIntegration(args[0])
return err == nil && name == "claude-desktop"
}
type launcherClient struct {
apiClient *api.Client
modelInventory []ModelInfo
@@ -342,8 +436,13 @@ func LaunchIntegration(ctx context.Context, req IntegrationLaunchRequest) error
}
policy := launchIntegrationPolicy(req)
if req.Restore {
return restoreIntegration(name, runner, req)
}
if policy.Confirm == LaunchConfirmAutoApprove && !isInteractiveSession() && req.ModelOverride == "" {
return fmt.Errorf("headless --yes launch for %s requires --model <model>", name)
if _, ok := runner.(ManagedAutodiscoveryIntegration); !ok {
return fmt.Errorf("headless --yes launch for %s requires --model <model>", name)
}
}
launchClient, saved, err := prepareIntegrationLaunch(name, policy)
@@ -351,6 +450,13 @@ func LaunchIntegration(ctx context.Context, req IntegrationLaunchRequest) error
return err
}
if autodiscovery, ok := runner.(ManagedAutodiscoveryIntegration); ok {
if err := EnsureIntegrationInstalled(name, runner); err != nil {
return err
}
return launchClient.launchManagedAutodiscoveryIntegration(ctx, name, runner, autodiscovery, saved, req)
}
if managed, ok := runner.(ManagedSingleModel); ok {
if err := EnsureIntegrationInstalled(name, runner); err != nil {
return err
@@ -370,6 +476,24 @@ func LaunchIntegration(ctx context.Context, req IntegrationLaunchRequest) error
return launchClient.launchSingleIntegration(ctx, name, runner, saved, req)
}
func restoreIntegration(name string, runner Runner, req IntegrationLaunchRequest) error {
if req.ModelOverride != "" || req.ConfigureOnly || len(req.ExtraArgs) > 0 {
return fmt.Errorf("--restore cannot be combined with --model, --config, or extra args")
}
restorable, ok := runner.(RestorableIntegration)
if !ok {
return fmt.Errorf("%s does not support --restore", name)
}
if err := EnsureIntegrationInstalled(name, runner); err != nil {
return err
}
if err := restorable.Restore(); err != nil {
return err
}
printRestoreSuccess(restorable)
return nil
}
func launchIntegrationPolicy(req IntegrationLaunchRequest) LaunchPolicy {
// TUI does not set a policy, whereas ollama launch <app> does as it can
// have flags which change the behavior.
@@ -420,7 +544,12 @@ func (c *launcherClient) buildLauncherIntegrationState(ctx context.Context, info
}
var currentModel string
var usable bool
if managed, ok := integration.spec.Runner.(ManagedSingleModel); ok {
if autodiscovery, ok := integration.spec.Runner.(ManagedAutodiscoveryIntegration); ok {
currentModel, usable, err = c.launcherManagedAutodiscoveryState(ctx, info.Name, autodiscovery)
if err != nil {
return LauncherIntegrationState{}, err
}
} else if managed, ok := integration.spec.Runner.(ManagedSingleModel); ok {
currentModel, usable, err = c.launcherManagedModelState(ctx, info.Name, managed)
if err != nil {
return LauncherIntegrationState{}, err
@@ -482,6 +611,10 @@ func (c *launcherClient) launcherManagedModelState(ctx context.Context, name str
return "", false, nil
}
if skips, ok := managed.(ManagedModelReadinessSkipper); ok && skips.SkipModelReadiness() {
return current, true, nil
}
usable, err := c.savedModelUsable(ctx, current)
if err != nil {
return current, false, err
@@ -489,6 +622,20 @@ func (c *launcherClient) launcherManagedModelState(ctx context.Context, name str
return current, usable, nil
}
func (c *launcherClient) launcherManagedAutodiscoveryState(ctx context.Context, name string, autodiscovery ManagedAutodiscoveryIntegration) (string, bool, error) {
if autodiscovery.AutodiscoveryConfigured() {
return autodiscovery.AutodiscoveredModel(), c.managedAutodiscoveryUsable(ctx, autodiscovery), nil
}
cfg, loadErr := loadStoredIntegrationConfig(name)
if loadErr == nil {
if current := primaryModelFromConfig(cfg); current != "" {
return current, false, nil
}
}
return "", false, nil
}
func (c *launcherClient) resolveRunModel(ctx context.Context, req RunModelRequest) (string, error) {
current := config.LastModel()
if !req.ForcePicker && current != "" && c.policy.Confirm == LaunchConfirmAutoApprove && !isInteractiveSession() {
@@ -586,8 +733,12 @@ func (c *launcherClient) launchManagedSingleIntegration(ctx context.Context, nam
return nil
}
if current == "" || needsConfigure || req.ModelOverride != "" || target != current {
if err := prepareManagedSingleIntegration(name, runner, managed, target); err != nil {
if needsConfigure || req.ModelOverride != "" || (current != "" && target != current) || !savedMatchesModels(saved, []string{target}) {
configureModels, err := c.managedSingleConfigureModels(ctx, managed, target)
if err != nil {
return err
}
if err := prepareManagedSingleIntegration(name, runner, managed, target, configureModels); err != nil {
return err
}
if refresher, ok := managed.(ManagedRuntimeRefresher); ok {
@@ -613,15 +764,146 @@ func (c *launcherClient) launchManagedSingleIntegration(ctx context.Context, nam
return runIntegration(runner, target, req.ExtraArgs)
}
func (c *launcherClient) launchManagedAutodiscoveryIntegration(ctx context.Context, name string, runner Runner, autodiscovery ManagedAutodiscoveryIntegration, saved *config.IntegrationConfig, req IntegrationLaunchRequest) error {
if req.ModelOverride != "" {
return fmt.Errorf("%s discovers models automatically; omit --model", runner)
}
target := autodiscovery.AutodiscoveredModel()
if err := c.ensureManagedAutodiscoveryUsable(ctx, autodiscovery, target); err != nil {
return err
}
needsConfigure := req.ForceConfigure || req.ConfigureOnly || !autodiscovery.AutodiscoveryConfigured() || !savedMatchesModels(saved, []string{target})
if needsConfigure {
if err := prepareManagedAutodiscoveryIntegration(name, runner, autodiscovery, target); err != nil {
return err
}
if refresher, ok := autodiscovery.(ManagedRuntimeRefresher); ok {
if err := refresher.RefreshRuntimeAfterConfigure(); err != nil {
return err
}
}
}
if !managedIntegrationOnboarded(saved, autodiscovery) {
if !isInteractiveSession() && managedRequiresInteractiveOnboarding(autodiscovery) {
return fmt.Errorf("%s still needs interactive gateway setup; run 'ollama launch %s' in a terminal to finish onboarding", runner, name)
}
if err := autodiscovery.Onboard(); err != nil {
return err
}
}
if !printConfigurationSuccess(autodiscovery) {
printRestoreHint(autodiscovery)
}
if req.ConfigureOnly {
return nil
}
return runIntegration(runner, target, req.ExtraArgs)
}
func (c *launcherClient) managedAutodiscoveryUsable(ctx context.Context, autodiscovery ManagedAutodiscoveryIntegration) bool {
if !managedAutodiscoveryUsesOllamaCloud(autodiscovery) {
return true
}
return c.ollamaCloudSignedIn(ctx)
}
func (c *launcherClient) ensureManagedAutodiscoveryUsable(ctx context.Context, autodiscovery ManagedAutodiscoveryIntegration, label string) error {
if !managedAutodiscoveryUsesOllamaCloud(autodiscovery) {
return nil
}
return ensureCloudAuth(ctx, c.apiClient, label)
}
func managedAutodiscoveryUsesOllamaCloud(autodiscovery ManagedAutodiscoveryIntegration) bool {
cloud, ok := autodiscovery.(ManagedAutodiscoveryCloudIntegration)
return ok && cloud.UsesOllamaCloud()
}
func printRestoreHint(integration any) {
hint, ok := integration.(RestoreHintIntegration)
if !ok {
return
}
if msg := strings.TrimSpace(hint.RestoreHint()); msg != "" {
fmt.Fprintln(os.Stderr, msg)
}
}
func printConfigurationSuccess(integration any) bool {
success, ok := integration.(ConfigurationSuccessIntegration)
if !ok {
return false
}
if msg := strings.TrimSpace(success.ConfigurationSuccessMessage()); msg != "" {
fmt.Fprintln(os.Stderr, msg)
return true
}
return false
}
func printRestoreSuccess(integration any) {
success, ok := integration.(RestoreSuccessIntegration)
if !ok {
return
}
if msg := strings.TrimSpace(success.RestoreSuccessMessage()); msg != "" {
fmt.Fprintln(os.Stderr, msg)
}
}
func (c *launcherClient) ollamaCloudSignedIn(ctx context.Context) bool {
if disabled, known := cloudStatusDisabled(ctx, c.apiClient); known && disabled {
return false
}
user, err := c.apiClient.Whoami(ctx)
return err == nil && user != nil && user.Name != ""
}
func (c *launcherClient) managedSingleConfigureModels(ctx context.Context, managed ManagedSingleModel, target string) ([]string, error) {
models := []string{target}
if _, ok := managed.(ManagedModelListConfigurer); !ok {
return models, nil
}
items, _, err := c.loadSelectableModels(ctx, []string{target}, target, "no models available")
if err != nil {
// Managed integrations that can use a model catalog should still be
// configurable with an explicit target even if the broader inventory
// cannot be loaded in the moment.
//nolint:nilerr
return models, nil
}
for _, item := range items {
models = append(models, item.Name)
}
return dedupeModelList(models), nil
}
func (c *launcherClient) resolveSingleIntegrationTarget(ctx context.Context, runner Runner, current string, req IntegrationLaunchRequest) (string, bool, error) {
target := req.ModelOverride
needsConfigure := req.ForceConfigure
skipReadiness := false
if skipper, ok := runner.(ManagedModelReadinessSkipper); ok {
skipReadiness = skipper.SkipModelReadiness()
}
if target == "" {
target = current
usable, err := c.savedModelUsable(ctx, target)
if err != nil {
return "", false, err
usable := skipReadiness && target != ""
if !skipReadiness {
var err error
usable, err = c.savedModelUsable(ctx, target)
if err != nil {
return "", false, err
}
}
if !usable {
needsConfigure = true
@@ -629,13 +911,19 @@ func (c *launcherClient) resolveSingleIntegrationTarget(ctx context.Context, run
}
if needsConfigure {
selected, err := c.selectSingleModelWithSelector(ctx, fmt.Sprintf("Select model for %s:", runner), target, DefaultSingleSelector)
selected, err := c.selectSingleModelWithSelectorReady(ctx, fmt.Sprintf("Select model for %s:", runner), target, DefaultSingleSelector, !skipReadiness)
if err != nil {
return "", false, err
}
target = selected
} else if err := c.ensureModelsReady(ctx, []string{target}); err != nil {
return "", false, err
} else if !skipReadiness {
if err := c.ensureModelsReady(ctx, []string{target}); err != nil {
return "", false, err
}
}
if target == "" {
return "", false, nil
}
return target, needsConfigure, nil
@@ -645,7 +933,7 @@ func savedIntegrationOnboarded(saved *config.IntegrationConfig) bool {
return saved != nil && saved.Onboarded
}
func managedIntegrationOnboarded(saved *config.IntegrationConfig, managed ManagedSingleModel) bool {
func managedIntegrationOnboarded(saved *config.IntegrationConfig, managed any) bool {
if !savedIntegrationOnboarded(saved) {
return false
}
@@ -659,7 +947,7 @@ func managedIntegrationOnboarded(saved *config.IntegrationConfig, managed Manage
// Most managed integrations treat onboarding as an interactive terminal step.
// Hermes opts out because its launch-owned onboarding is just bookkeeping, so
// headless launches should not be blocked once config is already prepared.
func managedRequiresInteractiveOnboarding(managed ManagedSingleModel) bool {
func managedRequiresInteractiveOnboarding(managed any) bool {
onboarding, ok := managed.(ManagedInteractiveOnboarding)
if !ok {
return true
@@ -668,6 +956,10 @@ func managedRequiresInteractiveOnboarding(managed ManagedSingleModel) bool {
}
func (c *launcherClient) selectSingleModelWithSelector(ctx context.Context, title, current string, selector SingleSelector) (string, error) {
return c.selectSingleModelWithSelectorReady(ctx, title, current, selector, true)
}
func (c *launcherClient) selectSingleModelWithSelectorReady(ctx context.Context, title, current string, selector SingleSelector, ensureReady bool) (string, error) {
if selector == nil {
return "", fmt.Errorf("no selector configured")
}
@@ -681,8 +973,13 @@ func (c *launcherClient) selectSingleModelWithSelector(ctx context.Context, titl
if err != nil {
return "", err
}
if err := c.ensureModelsReady(ctx, []string{selected}); err != nil {
return "", err
if selected == "" {
return "", ErrCancelled
}
if ensureReady {
if err := c.ensureModelsReady(ctx, []string{selected}); err != nil {
return "", err
}
}
return selected, nil
}
@@ -717,9 +1014,10 @@ func (c *launcherClient) loadSelectableModels(ctx context.Context, preChecked []
if err := c.loadModelInventoryOnce(ctx); err != nil {
return nil, nil, err
}
recommendations := c.recommendations(ctx)
cloudDisabled, _ := cloudStatusDisabled(ctx, c.apiClient)
items, orderedChecked, _, _ := buildModelList(c.modelInventory, preChecked, current)
items, orderedChecked, _, _ := buildModelListWithRecommendations(c.modelInventory, recommendations, preChecked, current)
if cloudDisabled {
items = filterCloudItems(items)
orderedChecked = c.filterDisabledCloudModels(ctx, orderedChecked)
@@ -730,6 +1028,60 @@ func (c *launcherClient) loadSelectableModels(ctx context.Context, preChecked []
return items, orderedChecked, nil
}
func (c *launcherClient) recommendations(ctx context.Context) []ModelItem {
recommendations, err := c.requestRecommendations(ctx)
if err != nil || len(recommendations) == 0 {
// Fail open: recommendation issues should not block launch flows.
// Fall back to built-in recommendations until server data is available.
fallback := append([]ModelItem(nil), recommendedModels...)
setDynamicCloudModelLimits(cloudModelLimitsFromRecommendations(fallback))
return fallback
}
setDynamicCloudModelLimits(cloudModelLimitsFromRecommendations(recommendations))
return recommendations
}
func (c *launcherClient) requestRecommendations(ctx context.Context) ([]ModelItem, error) {
resp, err := c.apiClient.ModelRecommendationsExperimental(ctx)
if err != nil {
return nil, err
}
items := make([]ModelItem, 0, len(resp.Recommendations))
seen := make(map[string]struct{}, len(resp.Recommendations))
for _, rec := range resp.Recommendations {
name := strings.TrimSpace(rec.Model)
if name == "" {
continue
}
if _, ok := seen[name]; ok {
continue
}
seen[name] = struct{}{}
if isCloudModelName(name) && (rec.ContextLength <= 0 || rec.MaxOutputTokens <= 0) {
slog.Warn("skipping cloud recommendation with missing limits", "model", name)
continue
}
description := strings.TrimSpace(rec.Description)
if description == "" {
description = "Recommended model"
}
items = append(items, ModelItem{
Name: name,
Description: description,
Recommended: true,
VRAMBytes: rec.VRAMBytes,
ContextLength: rec.ContextLength,
MaxOutputTokens: rec.MaxOutputTokens,
})
}
return items, nil
}
func (c *launcherClient) ensureModelsReady(ctx context.Context, models []string) error {
models = dedupeModelList(models)
if len(models) == 0 {
+569 -22
View File
@@ -13,6 +13,7 @@ import (
"strings"
"testing"
"github.com/google/go-cmp/cmp"
"github.com/ollama/ollama/cmd/config"
)
@@ -49,6 +50,22 @@ func (r *launcherSingleRunner) Run(model string, args []string) error {
func (r *launcherSingleRunner) String() string { return "StubSingle" }
type launcherRestorableRunner struct {
launcherSingleRunner
restored bool
restoreErr error
restoreSuccessMessage string
}
func (r *launcherRestorableRunner) Restore() error {
r.restored = true
return r.restoreErr
}
func (r *launcherRestorableRunner) RestoreSuccessMessage() string {
return r.restoreSuccessMessage
}
type launcherManagedRunner struct {
paths []string
currentModel string
@@ -98,6 +115,45 @@ type launcherHeadlessManagedRunner struct {
func (r *launcherHeadlessManagedRunner) RequiresInteractiveOnboarding() bool { return false }
type launcherManagedListRunner struct {
launcherManagedRunner
configuredModelLists [][]string
}
func (r *launcherManagedListRunner) ConfigureWithModels(primary string, models []string) error {
r.configuredModelLists = append(r.configuredModelLists, append([]string(nil), models...))
return r.Configure(primary)
}
type launcherManagedAutodiscoveryRunner struct {
launcherManagedRunner
autodiscoveryConfigures int
autodiscoveryConfigured bool
usesCloud bool
restoreHint string
configSuccessMessage string
}
func (r *launcherManagedAutodiscoveryRunner) AutodiscoveredModel() string { return "Ollama Cloud" }
func (r *launcherManagedAutodiscoveryRunner) UsesOllamaCloud() bool { return r.usesCloud }
func (r *launcherManagedAutodiscoveryRunner) RestoreHint() string { return r.restoreHint }
func (r *launcherManagedAutodiscoveryRunner) ConfigurationSuccessMessage() string {
return r.configSuccessMessage
}
func (r *launcherManagedAutodiscoveryRunner) AutodiscoveryConfigured() bool {
return r.autodiscoveryConfigured
}
func (r *launcherManagedAutodiscoveryRunner) ConfigureAutodiscovery() error {
r.autodiscoveryConfigures++
r.autodiscoveryConfigured = true
return nil
}
func setLaunchTestHome(t *testing.T, dir string) {
t.Helper()
t.Setenv("HOME", dir)
@@ -196,6 +252,8 @@ func TestBuildLauncherState_ManagedSingleIntegrationUsesCurrentModel(t *testing.
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"gemma4"}]}`)
case "/api/show":
@@ -229,6 +287,8 @@ func TestBuildLauncherState_ManagedSingleIntegrationShowsSavedModelWhenLiveConfi
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"gemma4"}]}`)
case "/api/show":
@@ -268,6 +328,8 @@ func TestLaunchIntegration_ManagedSingleIntegrationConfiguresOnboardsAndRuns(t *
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"gemma4"}]}`)
case "/api/show":
@@ -325,6 +387,8 @@ func TestLaunchIntegration_ManagedSingleIntegrationReOnboardsWhenSavedFlagIsStal
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"gemma4"}]}`)
case "/api/show":
@@ -409,7 +473,7 @@ func TestLaunchIntegration_ManagedSingleIntegrationConfigOnlySkipsFinalRun(t *te
}
}
func TestLaunchIntegration_ManagedSingleIntegrationRepairsMissingLiveConfigUsingSavedModel(t *testing.T) {
func TestLaunchIntegration_ManagedSingleIntegrationSkipsRewriteWhenSavedMatches(t *testing.T) {
tmpDir := t.TempDir()
setLaunchTestHome(t, tmpDir)
withInteractiveSession(t, true)
@@ -417,6 +481,8 @@ func TestLaunchIntegration_ManagedSingleIntegrationRepairsMissingLiveConfigUsing
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"gemma4"}]}`)
case "/api/show":
@@ -436,29 +502,30 @@ func TestLaunchIntegration_ManagedSingleIntegrationRepairsMissingLiveConfigUsing
withIntegrationOverride(t, "stubmanaged", runner)
DefaultSingleSelector = func(title string, items []ModelItem, current string) (string, error) {
t.Fatal("selector should not be called when saved model is reused for repair")
t.Fatal("selector should not be called when saved model matches target")
return "", nil
}
DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) {
return true, nil
t.Fatal("confirm prompt should not run when saved model matches target")
return false, nil
}
if err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{Name: "stubmanaged"}); err != nil {
t.Fatalf("LaunchIntegration returned error: %v", err)
}
if diff := compareStrings(runner.configured, []string{"gemma4"}); diff != "" {
t.Fatalf("expected missing live config to be rewritten from saved model: %s", diff)
if len(runner.configured) != 0 {
t.Fatalf("expected Configure to be skipped when saved matches, got %v", runner.configured)
}
if runner.refreshCalls != 1 {
t.Fatalf("expected repaired config to refresh runtime once, got %d", runner.refreshCalls)
if runner.refreshCalls != 0 {
t.Fatalf("expected no runtime refresh when config is unchanged, got %d", runner.refreshCalls)
}
if runner.ranModel != "gemma4" {
t.Fatalf("expected launch to use repaired saved model, got %q", runner.ranModel)
t.Fatalf("expected launch to run saved model, got %q", runner.ranModel)
}
}
func TestLaunchIntegration_ManagedSingleIntegrationConfigureOnlyRepairsMissingLiveConfig(t *testing.T) {
func TestLaunchIntegration_ManagedSingleIntegrationRewritesWhenSavedDiffers(t *testing.T) {
tmpDir := t.TempDir()
setLaunchTestHome(t, tmpDir)
withInteractiveSession(t, true)
@@ -466,6 +533,62 @@ func TestLaunchIntegration_ManagedSingleIntegrationConfigureOnlyRepairsMissingLi
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"gemma4"}]}`)
case "/api/show":
fmt.Fprint(w, `{"model_info":{"general.context_length":131072}}`)
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
t.Setenv("OLLAMA_HOST", srv.URL)
if err := config.SaveIntegration("stubmanaged", []string{"old-model"}); err != nil {
t.Fatalf("failed to save managed integration config: %v", err)
}
runner := &launcherManagedRunner{}
withIntegrationOverride(t, "stubmanaged", runner)
DefaultSingleSelector = func(title string, items []ModelItem, current string) (string, error) {
t.Fatal("selector should not be called when model override is provided")
return "", nil
}
DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) {
return true, nil
}
if err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{
Name: "stubmanaged",
ModelOverride: "gemma4",
}); err != nil {
t.Fatalf("LaunchIntegration returned error: %v", err)
}
if diff := compareStrings(runner.configured, []string{"gemma4"}); diff != "" {
t.Fatalf("expected Configure to run when saved differs from target: %s", diff)
}
if runner.refreshCalls != 1 {
t.Fatalf("expected runtime refresh once after configure, got %d", runner.refreshCalls)
}
if runner.ranModel != "gemma4" {
t.Fatalf("expected launch to run configured model, got %q", runner.ranModel)
}
}
func TestLaunchIntegration_ManagedSingleIntegrationRewritesWhenLiveConfigDrifts(t *testing.T) {
tmpDir := t.TempDir()
setLaunchTestHome(t, tmpDir)
withInteractiveSession(t, true)
withLauncherHooks(t)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"gemma4"},{"name":"qwen3:8b"}]}`)
case "/api/show":
fmt.Fprint(w, `{"model_info":{"general.context_length":131072}}`)
default:
@@ -479,32 +602,39 @@ func TestLaunchIntegration_ManagedSingleIntegrationConfigureOnlyRepairsMissingLi
t.Fatalf("failed to save managed integration config: %v", err)
}
runner := &launcherManagedRunner{}
runner := &launcherManagedRunner{
currentModel: "qwen3:8b",
}
withIntegrationOverride(t, "stubmanaged", runner)
DefaultSingleSelector = func(title string, items []ModelItem, current string) (string, error) {
t.Fatal("selector should not be called when saved model is reused for repair")
t.Fatal("selector should not be called when live config already provides the target")
return "", nil
}
DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) {
return true, nil
}
if err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{
Name: "stubmanaged",
ConfigureOnly: true,
}); err != nil {
if err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{Name: "stubmanaged"}); err != nil {
t.Fatalf("LaunchIntegration returned error: %v", err)
}
if diff := compareStrings(runner.configured, []string{"gemma4"}); diff != "" {
t.Fatalf("expected configure-only flow to rewrite missing live config: %s", diff)
if diff := compareStrings(runner.configured, []string{"qwen3:8b"}); diff != "" {
t.Fatalf("expected Configure to reconcile stale saved config to live target: %s", diff)
}
if runner.refreshCalls != 1 {
t.Fatalf("expected configure-only repair to refresh runtime once, got %d", runner.refreshCalls)
t.Fatalf("expected runtime refresh once after drift reconciliation, got %d", runner.refreshCalls)
}
if runner.ranModel != "" {
t.Fatalf("expected configure-only flow to skip final launch, got %q", runner.ranModel)
if runner.ranModel != "qwen3:8b" {
t.Fatalf("expected launch to run live configured model, got %q", runner.ranModel)
}
saved, err := config.LoadIntegration("stubmanaged")
if err != nil {
t.Fatalf("failed to reload managed integration config: %v", err)
}
if diff := compareStrings(saved.Models, []string{"qwen3:8b"}); diff != "" {
t.Fatalf("saved models mismatch after drift reconciliation: %s", diff)
}
}
@@ -552,6 +682,378 @@ func TestLaunchIntegration_ManagedSingleIntegrationStopsWhenRuntimeRefreshFails(
}
}
func TestLaunchIntegration_ManagedSingleIntegrationCanConfigureWithModelList(t *testing.T) {
tmpDir := t.TempDir()
setLaunchTestHome(t, tmpDir)
withInteractiveSession(t, true)
withLauncherHooks(t)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"gemma4"},{"name":"qwen3:8b"}]}`)
case "/api/show":
fmt.Fprint(w, `{"model_info":{"general.context_length":131072}}`)
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
t.Setenv("OLLAMA_HOST", srv.URL)
runner := &launcherManagedListRunner{}
withIntegrationOverride(t, "stubmanaged", runner)
DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) {
return true, nil
}
if err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{
Name: "stubmanaged",
ModelOverride: "gemma4",
}); err != nil {
t.Fatalf("LaunchIntegration returned error: %v", err)
}
if diff := compareStringSlices(runner.configuredModelLists, [][]string{{"gemma4", "kimi-k2.6:cloud", "qwen3.5:cloud", "glm-5.1:cloud", "minimax-m2.7:cloud", "qwen3.5", "qwen3:8b"}}); diff != "" {
t.Fatalf("configured model list mismatch (-want +got):\n%s", diff)
}
if diff := compareStrings(runner.configured, []string{"gemma4"}); diff != "" {
t.Fatalf("configured primary mismatch: %s", diff)
}
}
func TestLaunchIntegration_ManagedAutodiscoverySkipsModelPicker(t *testing.T) {
tmpDir := t.TempDir()
setLaunchTestHome(t, tmpDir)
withInteractiveSession(t, true)
withLauncherHooks(t)
runner := &launcherManagedAutodiscoveryRunner{}
withIntegrationOverride(t, "stubmanaged", runner)
DefaultSingleSelector = func(title string, items []ModelItem, current string) (string, error) {
t.Fatal("model selector should not run for autodiscovery integrations")
return "", nil
}
DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) {
return true, nil
}
if err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{Name: "stubmanaged"}); err != nil {
t.Fatalf("LaunchIntegration returned error: %v", err)
}
if runner.autodiscoveryConfigures != 1 {
t.Fatalf("expected one autodiscovery configure, got %d", runner.autodiscoveryConfigures)
}
if runner.ranModel != "Ollama Cloud" {
t.Fatalf("expected launch to run autodiscovery label, got %q", runner.ranModel)
}
saved, err := config.LoadIntegration("stubmanaged")
if err != nil {
t.Fatalf("failed to reload managed integration config: %v", err)
}
if diff := compareStrings(saved.Models, []string{"Ollama Cloud"}); diff != "" {
t.Fatalf("saved models mismatch: %s", diff)
}
}
func TestLaunchIntegration_ManagedAutodiscoveryPrintsConfigurationSuccessAfterConfigure(t *testing.T) {
tmpDir := t.TempDir()
setLaunchTestHome(t, tmpDir)
withInteractiveSession(t, true)
withLauncherHooks(t)
runner := &launcherManagedAutodiscoveryRunner{
restoreHint: "run restore command",
configSuccessMessage: "configured successfully\nrestore via success message",
}
withIntegrationOverride(t, "stubmanaged", runner)
DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) {
return true, nil
}
stderr := captureStderr(t, func() {
if err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{Name: "stubmanaged"}); err != nil {
t.Fatalf("LaunchIntegration returned error: %v", err)
}
})
if runner.autodiscoveryConfigures != 1 {
t.Fatalf("expected one autodiscovery configure, got %d", runner.autodiscoveryConfigures)
}
if !strings.Contains(stderr, "configured successfully") {
t.Fatalf("expected configuration success in stderr, got %q", stderr)
}
if !strings.Contains(stderr, "restore via success message") {
t.Fatalf("expected restore guidance in configuration success, got %q", stderr)
}
if strings.Contains(stderr, "run restore command") {
t.Fatalf("restore hint should not print separately after configure, got %q", stderr)
}
}
func TestLaunchIntegration_ManagedAutodiscoveryPrintsRestoreHintWhenAlreadyConfigured(t *testing.T) {
tmpDir := t.TempDir()
setLaunchTestHome(t, tmpDir)
withInteractiveSession(t, true)
withLauncherHooks(t)
runner := &launcherManagedAutodiscoveryRunner{
autodiscoveryConfigured: true,
restoreHint: "run restore command",
}
withIntegrationOverride(t, "stubmanaged", runner)
if err := config.SaveIntegration("stubmanaged", []string{"Ollama Cloud"}); err != nil {
t.Fatalf("failed to save managed integration config: %v", err)
}
if err := config.MarkIntegrationOnboarded("stubmanaged"); err != nil {
t.Fatalf("failed to mark integration onboarded: %v", err)
}
stderr := captureStderr(t, func() {
if err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{Name: "stubmanaged"}); err != nil {
t.Fatalf("LaunchIntegration returned error: %v", err)
}
})
if runner.autodiscoveryConfigures != 0 {
t.Fatalf("expected configured autodiscovery integration not to reconfigure, got %d configures", runner.autodiscoveryConfigures)
}
if runner.ranModel != "Ollama Cloud" {
t.Fatalf("expected launch to run autodiscovery label, got %q", runner.ranModel)
}
if !strings.Contains(stderr, "run restore command") {
t.Fatalf("expected restore hint in stderr, got %q", stderr)
}
}
func TestLaunchIntegration_ManagedAutodiscoveryPrintsConfigurationSuccessWhenAlreadyConfigured(t *testing.T) {
tmpDir := t.TempDir()
setLaunchTestHome(t, tmpDir)
withInteractiveSession(t, true)
withLauncherHooks(t)
runner := &launcherManagedAutodiscoveryRunner{
autodiscoveryConfigured: true,
restoreHint: "run restore command",
configSuccessMessage: "configured successfully\nrestore via success message",
}
withIntegrationOverride(t, "stubmanaged", runner)
if err := config.SaveIntegration("stubmanaged", []string{"Ollama Cloud"}); err != nil {
t.Fatalf("failed to save managed integration config: %v", err)
}
if err := config.MarkIntegrationOnboarded("stubmanaged"); err != nil {
t.Fatalf("failed to mark integration onboarded: %v", err)
}
stderr := captureStderr(t, func() {
if err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{Name: "stubmanaged"}); err != nil {
t.Fatalf("LaunchIntegration returned error: %v", err)
}
})
if runner.autodiscoveryConfigures != 0 {
t.Fatalf("expected configured autodiscovery integration not to reconfigure, got %d configures", runner.autodiscoveryConfigures)
}
if !strings.Contains(stderr, "configured successfully") {
t.Fatalf("expected configuration success in stderr, got %q", stderr)
}
if !strings.Contains(stderr, "restore via success message") {
t.Fatalf("expected restore guidance in configuration success, got %q", stderr)
}
if strings.Contains(stderr, "run restore command") {
t.Fatalf("restore hint should not print separately when success message exists, got %q", stderr)
}
}
func TestLaunchIntegration_RestorePrintsSuccessMessage(t *testing.T) {
tmpDir := t.TempDir()
setLaunchTestHome(t, tmpDir)
withLauncherHooks(t)
runner := &launcherRestorableRunner{
restoreSuccessMessage: "restored successfully",
}
withIntegrationOverride(t, "stubrestore", runner)
stderr := captureStderr(t, func() {
if err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{
Name: "stubrestore",
Restore: true,
}); err != nil {
t.Fatalf("LaunchIntegration returned error: %v", err)
}
})
if !runner.restored {
t.Fatal("expected restore to run")
}
if !strings.Contains(stderr, "restored successfully") {
t.Fatalf("expected restore success in stderr, got %q", stderr)
}
}
func TestLaunchIntegration_ManagedAutodiscoveryForceConfigureRerunsSetup(t *testing.T) {
tmpDir := t.TempDir()
setLaunchTestHome(t, tmpDir)
withInteractiveSession(t, true)
withLauncherHooks(t)
runner := &launcherManagedAutodiscoveryRunner{
autodiscoveryConfigured: true,
}
withIntegrationOverride(t, "stubmanaged", runner)
if err := config.SaveIntegration("stubmanaged", []string{"Ollama Cloud"}); err != nil {
t.Fatalf("failed to save managed integration config: %v", err)
}
if err := config.MarkIntegrationOnboarded("stubmanaged"); err != nil {
t.Fatalf("failed to mark integration onboarded: %v", err)
}
if err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{
Name: "stubmanaged",
ForceConfigure: true,
}); err != nil {
t.Fatalf("LaunchIntegration returned error: %v", err)
}
if runner.autodiscoveryConfigures != 1 {
t.Fatalf("expected forced autodiscovery configure to rerun setup, got %d configures", runner.autodiscoveryConfigures)
}
if runner.ranModel != "Ollama Cloud" {
t.Fatalf("expected launch to run autodiscovery label, got %q", runner.ranModel)
}
}
func TestLaunchIntegration_CloudAutodiscoveryUsesSignInHook(t *testing.T) {
tmpDir := t.TempDir()
setLaunchTestHome(t, tmpDir)
withInteractiveSession(t, true)
withLauncherHooks(t)
runner := &launcherManagedAutodiscoveryRunner{usesCloud: true}
withIntegrationOverride(t, "stubmanaged", runner)
signInCalled := false
DefaultSignIn = func(modelName, signInURL string) (string, error) {
signInCalled = true
if modelName != "Ollama Cloud" {
t.Fatalf("sign-in model = %q, want Ollama Cloud", modelName)
}
if signInURL != "https://example.com/signin" {
t.Fatalf("sign-in URL = %q, want test URL", signInURL)
}
return "test-user", nil
}
DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) {
return true, nil
}
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/status":
w.WriteHeader(http.StatusNotFound)
fmt.Fprint(w, `{"error":"not found"}`)
case "/api/me":
w.WriteHeader(http.StatusUnauthorized)
fmt.Fprint(w, `{"error":"unauthorized","signin_url":"https://example.com/signin"}`)
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
t.Setenv("OLLAMA_HOST", srv.URL)
if err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{Name: "stubmanaged"}); err != nil {
t.Fatalf("LaunchIntegration returned error: %v", err)
}
if !signInCalled {
t.Fatal("expected cloud autodiscovery launch to use the sign-in hook")
}
if runner.autodiscoveryConfigures != 1 {
t.Fatalf("expected one autodiscovery configure, got %d", runner.autodiscoveryConfigures)
}
if runner.ranModel != "Ollama Cloud" {
t.Fatalf("expected launch to run autodiscovery label, got %q", runner.ranModel)
}
}
func TestBuildLauncherIntegrationState_CloudAutodiscoveryRequiresSignedIn(t *testing.T) {
tmpDir := t.TempDir()
setLaunchTestHome(t, tmpDir)
withLauncherHooks(t)
runner := &launcherManagedAutodiscoveryRunner{
autodiscoveryConfigured: true,
usesCloud: true,
}
withIntegrationOverride(t, "stubmanaged", runner)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/status":
w.WriteHeader(http.StatusNotFound)
fmt.Fprint(w, `{"error":"not found"}`)
case "/api/me":
w.WriteHeader(http.StatusUnauthorized)
fmt.Fprint(w, `{"error":"unauthorized","signin_url":"https://example.com/signin"}`)
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
t.Setenv("OLLAMA_HOST", srv.URL)
launchClient, err := newLauncherClient(defaultLaunchPolicy(true, false))
if err != nil {
t.Fatal(err)
}
state, err := launchClient.buildLauncherIntegrationState(context.Background(), IntegrationInfo{
Name: "stubmanaged",
DisplayName: "Stub Managed",
})
if err != nil {
t.Fatalf("buildLauncherIntegrationState returned error: %v", err)
}
if state.CurrentModel != "Ollama Cloud" {
t.Fatalf("current model = %q, want Ollama Cloud", state.CurrentModel)
}
if state.ModelUsable {
t.Fatal("expected cloud autodiscovery config to be unusable while signed out")
}
}
func TestLaunchIntegration_ManagedAutodiscoveryRejectsModelOverride(t *testing.T) {
tmpDir := t.TempDir()
setLaunchTestHome(t, tmpDir)
withInteractiveSession(t, true)
withLauncherHooks(t)
runner := &launcherManagedAutodiscoveryRunner{}
withIntegrationOverride(t, "stubmanaged", runner)
err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{
Name: "stubmanaged",
ModelOverride: "qwen3.5:cloud",
})
if err == nil || !strings.Contains(err.Error(), "discovers models automatically") {
t.Fatalf("LaunchIntegration error = %v, want automatic discovery guidance", err)
}
if runner.autodiscoveryConfigures != 0 {
t.Fatalf("expected no configure after model override rejection, got %d", runner.autodiscoveryConfigures)
}
}
func TestLaunchIntegration_ManagedSingleIntegrationHeadlessNeedsInteractiveOnboarding(t *testing.T) {
tmpDir := t.TempDir()
setLaunchTestHome(t, tmpDir)
@@ -652,6 +1154,8 @@ func TestBuildLauncherState_InstalledAndCloudDisabled(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`)
case "/api/status":
@@ -704,6 +1208,8 @@ func TestBuildLauncherState_MigratesLegacyOpenclawAliasConfig(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`)
default:
@@ -749,6 +1255,8 @@ func TestBuildLauncherState_ToleratesInventoryFailure(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
w.WriteHeader(http.StatusInternalServerError)
fmt.Fprint(w, `{"error":"temporary failure"}`)
@@ -795,6 +1303,8 @@ func TestResolveRunModel_UsesSavedModelWithoutSelector(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`)
case "/api/show":
@@ -840,6 +1350,8 @@ func TestResolveRunModel_HeadlessYesAutoPicksLastModel(t *testing.T) {
modelPulled := false
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`)
case "/api/show":
@@ -902,6 +1414,8 @@ func TestResolveRunModel_UsesRequestPolicy(t *testing.T) {
modelPulled := false
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`)
case "/api/show":
@@ -961,6 +1475,8 @@ func TestResolveRunModel_ForcePickerAlwaysUsesSelector(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"llama3.2"},{"name":"qwen3:8b"}]}`)
case "/api/show":
@@ -1011,6 +1527,8 @@ func TestResolveRunModel_ForcePicker_DoesNotReorderByLastModel(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"qwen3.5"},{"name":"gemma4"}]}`)
case "/api/show":
@@ -1061,6 +1579,8 @@ func TestResolveRunModel_UsesSignInHookForCloudModel(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[]}`)
case "/api/status":
@@ -1118,6 +1638,8 @@ func TestLaunchIntegration_EditorForceConfigure(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"llama3.2"},{"name":"qwen3:8b"}]}`)
case "/api/show":
@@ -1187,6 +1709,8 @@ func TestLaunchIntegration_EditorForceConfigure_FloatsCheckedModelsInPicker(t *t
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"qwen3.5:cloud","remote_model":"qwen3.5"},{"name":"qwen3.5"}]}`)
case "/api/show":
@@ -1216,8 +1740,9 @@ func TestLaunchIntegration_EditorForceConfigure_FloatsCheckedModelsInPicker(t *t
if len(gotItems) == 0 {
t.Fatal("expected multi selector to receive items")
}
if gotItems[0] != "qwen3.5:cloud" {
t.Fatalf("expected checked models floated to top with qwen3.5:cloud first, got %v", gotItems)
wantItems := recommendedNames()
if diff := cmp.Diff(wantItems, gotItems); diff != "" {
t.Fatalf("expected fixed recommended order in selector items (-want +got):\n%s", diff)
}
if len(gotPreChecked) < 2 {
t.Fatalf("expected prechecked models to be preserved, got %v", gotPreChecked)
@@ -1304,6 +1829,8 @@ func TestLaunchIntegration_EditorCloudDisabledFallsBackToSelector(t *testing.T)
switch r.URL.Path {
case "/api/status":
fmt.Fprint(w, `{"cloud":{"disabled":true,"source":"config"}}`)
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`)
case "/api/show":
@@ -1351,6 +1878,8 @@ func TestLaunchIntegration_EditorConfigureMultiSkipsMissingLocalAndPersistsAccep
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"glm-5:cloud","remote_model":"glm-5"}]}`)
case "/api/status":
@@ -1433,6 +1962,8 @@ func TestLaunchIntegration_EditorConfigureMultiSkipsUnauthedCloudAndPersistsAcce
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"llama3.2"},{"name":"glm-5:cloud","remote_model":"glm-5"}]}`)
case "/api/status":
@@ -1518,6 +2049,8 @@ func TestLaunchIntegration_EditorConfigureMultiRemovesReselectedFailingModel(t *
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"glm-5:cloud","remote_model":"glm-5"},{"name":"llama3.2"}]}`)
case "/api/status":
@@ -1605,6 +2138,8 @@ func TestLaunchIntegration_EditorConfigureMultiAllFailuresKeepsExistingAndSkipsL
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[]}`)
case "/api/show":
@@ -1911,6 +2446,8 @@ func TestLaunchIntegration_ConfigureOnlyDoesNotRequireInstalledBinary(t *testing
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`)
case "/api/show":
@@ -1954,6 +2491,8 @@ func TestLaunchIntegration_ClaudeSavesPrimaryModel(t *testing.T) {
var aliasSyncCalled bool
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[]}`)
case "/api/status":
@@ -2013,6 +2552,8 @@ func TestLaunchIntegration_ClaudeForceConfigureReprompts(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"qwen3:8b"}]}`)
case "/api/show":
@@ -2070,6 +2611,8 @@ func TestLaunchIntegration_ClaudeForceConfigureMissingSelectionDoesNotSave(t *te
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`)
case "/api/show":
@@ -2196,6 +2739,8 @@ func TestLaunchIntegration_ConfigureOnlyPrompt(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`)
case "/api/show":
@@ -2420,6 +2965,8 @@ func TestLaunchIntegration_HeadlessSelectorFlowFailsWithoutPrompt(t *testing.T)
pullCalled := false
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/experimental/model-recommendations":
fmt.Fprint(w, `{"recommendations":[]}`)
case "/api/tags":
fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`)
case "/api/show":
+126 -35
View File
@@ -4,34 +4,43 @@ import (
"context"
"errors"
"fmt"
"math"
"net/http"
"os"
"os/exec"
"runtime"
"slices"
"strings"
"sync"
"time"
"github.com/ollama/ollama/api"
"github.com/ollama/ollama/cmd/config"
"github.com/ollama/ollama/cmd/internal/fileutil"
"github.com/ollama/ollama/format"
internalcloud "github.com/ollama/ollama/internal/cloud"
"github.com/ollama/ollama/internal/modelref"
"github.com/ollama/ollama/progress"
)
var recommendedModels = []ModelItem{
{Name: "kimi-k2.5:cloud", Description: "Multimodal reasoning with subagents", Recommended: true},
{Name: "qwen3.5:cloud", Description: "Reasoning, coding, and agentic tool use with vision", Recommended: true},
{Name: "glm-5.1:cloud", Description: "Reasoning and code generation", Recommended: true},
{Name: "minimax-m2.7:cloud", Description: "Fast, efficient coding and real-world productivity", Recommended: true},
{Name: "gemma4", Description: "Reasoning and code generation locally", Recommended: true},
{Name: "qwen3.5", Description: "Reasoning, coding, and visual understanding locally", Recommended: true},
{Name: "kimi-k2.6:cloud", Description: "State-of-the-art coding, long-horizon execution, and multimodal agent swarm capability", Recommended: true, ContextLength: 262_144, MaxOutputTokens: 262_144},
{Name: "qwen3.5:cloud", Description: "Reasoning, coding, and agentic tool use with vision", Recommended: true, ContextLength: 262_144, MaxOutputTokens: 32_768},
{Name: "glm-5.1:cloud", Description: "Reasoning and code generation", Recommended: true, ContextLength: 202_752, MaxOutputTokens: 131_072},
{Name: "minimax-m2.7:cloud", Description: "Fast, efficient coding and real-world productivity", Recommended: true, ContextLength: 204_800, MaxOutputTokens: 128_000},
{Name: "gemma4", Description: "Reasoning and code generation locally", Recommended: true, VRAMBytes: 12 * format.GigaByte},
{Name: "qwen3.5", Description: "Reasoning, coding, and visual understanding locally", Recommended: true, VRAMBytes: 14 * format.GigaByte},
}
var recommendedVRAM = map[string]string{
"gemma4": "~16GB",
"qwen3.5": "~11GB",
func displayVRAM(vramBytes int64) string {
if vramBytes <= 0 {
return ""
}
gb := float64(vramBytes) / format.GigaByte
if gb == math.Trunc(gb) {
return fmt.Sprintf("~%.0fGB", gb)
}
return fmt.Sprintf("~%.1fGB", gb)
}
// cloudModelLimit holds context and output token limits for a cloud model.
@@ -40,10 +49,10 @@ type cloudModelLimit struct {
Output int
}
// cloudModelLimits maps cloud model base names to their token limits.
// extraCloudModelLimits maps cloud model base names to token limits for models
// that are not already covered by recommendedModels fallback entries.
// TODO(parthsareen): grab context/output limits from model info instead of hardcoding
var cloudModelLimits = map[string]cloudModelLimit{
"minimax-m2.7": {Context: 204_800, Output: 128_000},
var extraCloudModelLimits = map[string]cloudModelLimit{
"cogito-2.1:671b": {Context: 163_840, Output: 65_536},
"deepseek-v3.1:671b": {Context: 163_840, Output: 163_840},
"deepseek-v3.2": {Context: 163_840, Output: 65_536},
@@ -56,6 +65,7 @@ var cloudModelLimits = map[string]cloudModelLimit{
"gpt-oss:20b": {Context: 131_072, Output: 131_072},
"kimi-k2:1t": {Context: 262_144, Output: 262_144},
"kimi-k2.5": {Context: 262_144, Output: 262_144},
"kimi-k2.6": {Context: 262_144, Output: 262_144},
"kimi-k2-thinking": {Context: 262_144, Output: 262_144},
"nemotron-3-nano:30b": {Context: 1_048_576, Output: 131_072},
"qwen3-coder:480b": {Context: 262_144, Output: 65_536},
@@ -64,11 +74,24 @@ var cloudModelLimits = map[string]cloudModelLimit{
"qwen3.5": {Context: 262_144, Output: 32_768},
}
var cloudModelLimits = mergeCloudModelLimits(cloudModelLimitsFromRecommendations(recommendedModels), extraCloudModelLimits)
var (
dynamicCloudModelLimitsMu sync.RWMutex
dynamicCloudModelLimits = map[string]cloudModelLimit{}
)
// lookupCloudModelLimit returns the token limits for a cloud model.
// It normalizes explicit cloud source suffixes before checking the shared limit map.
func lookupCloudModelLimit(name string) (cloudModelLimit, bool) {
base, stripped := modelref.StripCloudSourceTag(name)
if stripped {
dynamicCloudModelLimitsMu.RLock()
l, ok := dynamicCloudModelLimits[base]
dynamicCloudModelLimitsMu.RUnlock()
if ok {
return l, true
}
if l, ok := cloudModelLimits[base]; ok {
return l, true
}
@@ -76,6 +99,49 @@ func lookupCloudModelLimit(name string) (cloudModelLimit, bool) {
return cloudModelLimit{}, false
}
func setDynamicCloudModelLimits(limits map[string]cloudModelLimit) {
dynamicCloudModelLimitsMu.Lock()
defer dynamicCloudModelLimitsMu.Unlock()
if limits == nil {
dynamicCloudModelLimits = map[string]cloudModelLimit{}
return
}
cp := make(map[string]cloudModelLimit, len(limits))
for k, v := range limits {
cp[k] = v
}
dynamicCloudModelLimits = cp
}
func cloudModelLimitsFromRecommendations(recommendations []ModelItem) map[string]cloudModelLimit {
limits := make(map[string]cloudModelLimit, len(recommendations))
for _, rec := range recommendations {
if !isCloudModelName(rec.Name) || rec.ContextLength <= 0 || rec.MaxOutputTokens <= 0 {
continue
}
base, stripped := modelref.StripCloudSourceTag(rec.Name)
if !stripped || base == "" {
continue
}
limits[base] = cloudModelLimit{
Context: rec.ContextLength,
Output: rec.MaxOutputTokens,
}
}
return limits
}
func mergeCloudModelLimits(base map[string]cloudModelLimit, overlay map[string]cloudModelLimit) map[string]cloudModelLimit {
out := make(map[string]cloudModelLimit, len(base)+len(overlay))
for name, limit := range base {
out[name] = limit
}
for name, limit := range overlay {
out[name] = limit
}
return out
}
// missingModelPolicy controls how model-not-found errors should be handled.
type missingModelPolicy int
@@ -115,6 +181,10 @@ func ensureAuth(ctx context.Context, client *api.Client, cloudModels map[string]
if len(selectedCloudModels) == 0 {
return nil
}
return ensureCloudAuth(ctx, client, strings.Join(selectedCloudModels, ", "))
}
func ensureCloudAuth(ctx context.Context, client *api.Client, modelList string) error {
if disabled, known := cloudStatusDisabled(ctx, client); known && disabled {
return errors.New(internalcloud.DisabledError("remote inference is unavailable"))
}
@@ -126,11 +196,12 @@ func ensureAuth(ctx context.Context, client *api.Client, cloudModels map[string]
var aErr api.AuthorizationError
if !errors.As(err, &aErr) || aErr.SigninURL == "" {
return err
if err != nil {
return err
}
return fmt.Errorf("%s requires sign in", modelList)
}
modelList := strings.Join(selectedCloudModels, ", ")
if DefaultSignIn != nil {
_, err := DefaultSignIn(modelList, aErr.SigninURL)
if errors.Is(err, ErrCancelled) {
@@ -244,13 +315,35 @@ func prepareEditorIntegration(name string, runner Runner, editor Editor, models
return nil
}
func prepareManagedSingleIntegration(name string, runner Runner, managed ManagedSingleModel, model string) error {
func prepareManagedSingleIntegration(name string, runner Runner, managed ManagedSingleModel, model string, models []string) error {
if ok, err := confirmConfigEdit(runner, managed.Paths()); err != nil {
return err
} else if !ok {
return errCancelled
}
if err := managed.Configure(model); err != nil {
models = dedupeModelList(append([]string{model}, models...))
var err error
if withModels, ok := managed.(ManagedModelListConfigurer); ok {
err = withModels.ConfigureWithModels(model, models)
} else {
err = managed.Configure(model)
}
if err != nil {
return fmt.Errorf("setup failed: %w", err)
}
if err := config.SaveIntegration(name, []string{model}); err != nil {
return fmt.Errorf("failed to save: %w", err)
}
return nil
}
func prepareManagedAutodiscoveryIntegration(name string, runner Runner, autodiscovery ManagedAutodiscoveryIntegration, model string) error {
if ok, err := confirmConfigEdit(runner, autodiscovery.Paths()); err != nil {
return err
} else if !ok {
return errCancelled
}
if err := autodiscovery.ConfigureAutodiscovery(); err != nil {
return fmt.Errorf("setup failed: %w", err)
}
if err := config.SaveIntegration(name, []string{model}); err != nil {
@@ -275,13 +368,17 @@ func confirmConfigEdit(runner Runner, paths []string) (bool, error) {
// buildModelList merges existing models with recommendations for selection UIs.
func buildModelList(existing []modelInfo, preChecked []string, current string) (items []ModelItem, orderedChecked []string, existingModels, cloudModels map[string]bool) {
return buildModelListWithRecommendations(existing, recommendedModels, preChecked, current)
}
func buildModelListWithRecommendations(existing []modelInfo, recommendations []ModelItem, preChecked []string, current string) (items []ModelItem, orderedChecked []string, existingModels, cloudModels map[string]bool) {
existingModels = make(map[string]bool)
cloudModels = make(map[string]bool)
recommended := make(map[string]bool)
var hasLocalModel, hasCloudModel bool
recDesc := make(map[string]string)
for _, rec := range recommendedModels {
for _, rec := range recommendations {
recommended[rec.Name] = true
recDesc[rec.Name] = rec.Description
}
@@ -300,7 +397,7 @@ func buildModelList(existing []modelInfo, preChecked []string, current string) (
items = append(items, item)
}
for _, rec := range recommendedModels {
for _, rec := range recommendations {
if existingModels[rec.Name] || existingModels[rec.Name+":latest"] {
continue
}
@@ -346,7 +443,7 @@ func buildModelList(existing []modelInfo, preChecked []string, current string) (
if items[i].Description != "" {
parts = append(parts, items[i].Description)
}
if vram := recommendedVRAM[items[i].Name]; vram != "" {
if vram := displayVRAM(items[i].VRAMBytes); vram != "" {
parts = append(parts, vram)
}
parts = append(parts, "(not downloaded)")
@@ -355,23 +452,17 @@ func buildModelList(existing []modelInfo, preChecked []string, current string) (
}
recRank := make(map[string]int)
for i, rec := range recommendedModels {
for i, rec := range recommendations {
recRank[rec.Name] = i + 1
}
if hasLocalModel || hasCloudModel {
// Keep the Recommended section pinned to recommendation order. Checked
// and default-model priority only apply within the More section.
slices.SortStableFunc(items, func(a, b ModelItem) int {
ac, bc := checked[a.Name], checked[b.Name]
aNew, bNew := notInstalled[a.Name], notInstalled[b.Name]
aRec, bRec := recRank[a.Name] > 0, recRank[b.Name] > 0
aCloud, bCloud := cloudModels[a.Name], cloudModels[b.Name]
if ac != bc {
if ac {
return -1
}
return 1
}
if aRec != bRec {
if aRec {
return -1
@@ -379,14 +470,14 @@ func buildModelList(existing []modelInfo, preChecked []string, current string) (
return 1
}
if aRec && bRec {
if aCloud != bCloud {
if aCloud {
return -1
}
return 1
}
return recRank[a.Name] - recRank[b.Name]
}
if ac != bc {
if ac {
return -1
}
return 1
}
// Among checked non-recommended items - put the default first
if ac && !aRec && current != "" {
aCurrent := a.Name == current
+203 -183
View File
@@ -14,8 +14,6 @@ import (
"strings"
"time"
"golang.org/x/mod/semver"
"github.com/ollama/ollama/api"
"github.com/ollama/ollama/cmd/internal/fileutil"
"github.com/ollama/ollama/envconfig"
@@ -30,6 +28,8 @@ var openclawModelShowTimeout = 5 * time.Second
// openclawFreshInstall is set to true when ensureOpenclawInstalled performs an install
var openclawFreshInstall bool
var openclawCanInstallDaemon = canInstallDaemon
type Openclaw struct{}
func (c *Openclaw) String() string { return "OpenClaw" }
@@ -60,6 +60,7 @@ func (c *Openclaw) Run(model string, args []string) error {
// the newest wizard flags (e.g. --auth-choice ollama).
if !openclawFreshInstall {
update := exec.Command(bin, "update")
update.Env = openclawInstallEnv()
update.Stdout = os.Stdout
update.Stderr = os.Stderr
_ = update.Run() // best-effort; continue even if update fails
@@ -75,19 +76,18 @@ func (c *Openclaw) Run(model string, args []string) error {
"--auth-choice", "ollama",
"--custom-base-url", envconfig.Host().String(),
"--custom-model-id", model,
// Launch owns the first real gateway startup immediately after onboarding,
// so don't let OpenClaw fail the whole first-run flow on a transient
// daemon health probe.
"--skip-health",
"--skip-channels",
"--skip-skills",
}
if canInstallDaemon() {
if openclawCanInstallDaemon() {
onboardArgs = append(onboardArgs, "--install-daemon")
} else {
// When we can't install a daemon (e.g. no systemd, sudo dropped
// XDG_RUNTIME_DIR, or container environment), skip the gateway
// health check so non-interactive onboarding completes. The
// gateway is started as a foreground child process after onboarding.
onboardArgs = append(onboardArgs, "--skip-health")
}
cmd := exec.Command(bin, onboardArgs...)
cmd.Env = openclawInstallEnv()
cmd.Stdin = os.Stdin
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
@@ -98,13 +98,23 @@ func (c *Openclaw) Run(model string, args []string) error {
patchDeviceScopes()
}
if ensureWebSearchPlugin() {
registerWebSearchPlugin()
}
configureOllamaWebSearch()
// When extra args are passed through, run exactly what the user asked for
// after setup and skip the built-in gateway+TUI convenience flow.
if len(args) > 0 {
cleanup := func() {}
if shouldEnsureGatewayForArgs(args) {
cleanupFn, _, _, err := c.ensureGatewayReady(bin)
if err != nil {
return windowsHint(err)
}
if cleanupFn != nil {
cleanup = cleanupFn
}
}
defer cleanup()
cmd := exec.Command(bin, args...)
cmd.Env = openclawEnv()
cmd.Stdin = os.Stdin
@@ -125,41 +135,11 @@ func (c *Openclaw) Run(model string, args []string) error {
fmt.Fprintf(os.Stderr, "\n%sStarting your assistant — this may take a moment...%s\n\n", ansiGray, ansiReset)
token, port := c.gatewayInfo()
addr := fmt.Sprintf("localhost:%d", port)
// If the gateway is already running (e.g. via the daemon), restart it
// so it picks up any config changes (model, provider, etc.).
if portOpen(addr) {
restart := exec.Command(bin, "daemon", "restart")
restart.Env = openclawEnv()
if err := restart.Run(); err != nil {
fmt.Fprintf(os.Stderr, "%s Warning: daemon restart failed: %v%s\n", ansiYellow, err, ansiReset)
}
if !waitForPort(addr, 10*time.Second) {
fmt.Fprintf(os.Stderr, "%s Warning: gateway did not come back after restart%s\n", ansiYellow, ansiReset)
}
}
// If the gateway isn't running, start it as a background child process.
if !portOpen(addr) {
gw := exec.Command(bin, "gateway", "run", "--force")
gw.Env = openclawEnv()
if err := gw.Start(); err != nil {
return windowsHint(fmt.Errorf("failed to start gateway: %w", err))
}
defer func() {
if gw.Process != nil {
_ = gw.Process.Kill()
_ = gw.Wait()
}
}()
}
fmt.Fprintf(os.Stderr, "%sStarting gateway...%s\n", ansiGray, ansiReset)
if !waitForPort(addr, 30*time.Second) {
return windowsHint(fmt.Errorf("gateway did not start on %s", addr))
cleanup, token, port, err := c.ensureGatewayReady(bin)
if err != nil {
return windowsHint(err)
}
defer cleanup()
printOpenclawReady(bin, token, port, firstLaunch)
@@ -179,6 +159,66 @@ func (c *Openclaw) Run(model string, args []string) error {
return nil
}
func shouldEnsureGatewayForArgs(args []string) bool {
return len(args) > 0 && args[0] == "tui"
}
func (c *Openclaw) ensureGatewayReady(bin string) (func(), string, int, error) {
token, port := c.gatewayInfo()
addr := fmt.Sprintf("127.0.0.1:%d", port)
// If the gateway is already running (e.g. via the daemon), restart it
// so it picks up any config changes (model, provider, etc.).
if portOpen(addr) {
restart := exec.Command(bin, "daemon", "restart")
restart.Env = openclawEnv()
if err := restart.Run(); err != nil {
fmt.Fprintf(os.Stderr, "%s Warning: daemon restart failed: %v%s\n", ansiYellow, err, ansiReset)
}
if !waitForPort(addr, 10*time.Second) {
fmt.Fprintf(os.Stderr, "%s Warning: gateway did not come back after restart%s\n", ansiYellow, ansiReset)
}
}
// If the daemon is installed but not currently listening, try to bring it
// up before falling back to a foreground child process.
if openclawCanInstallDaemon() && !portOpen(addr) {
start := exec.Command(bin, "daemon", "start")
start.Env = openclawEnv()
if err := start.Run(); err != nil {
fmt.Fprintf(os.Stderr, "%s Warning: daemon start failed: %v%s\n", ansiYellow, err, ansiReset)
} else if waitForPort(addr, 10*time.Second) {
fmt.Fprintf(os.Stderr, "%sStarting gateway...%s\n", ansiGray, ansiReset)
return func() {}, token, port, nil
}
}
cleanup := func() {}
// If the gateway still isn't running, start it as a background child process.
if !portOpen(addr) {
gw := exec.Command(bin, "gateway", "run", "--force")
gw.Env = openclawEnv()
if err := gw.Start(); err != nil {
return nil, "", 0, fmt.Errorf("failed to start gateway: %w", err)
}
cleanup = func() {
if gw.Process != nil {
_ = gw.Process.Kill()
_ = gw.Wait()
}
}
}
fmt.Fprintf(os.Stderr, "%sStarting gateway...%s\n", ansiGray, ansiReset)
if !waitForPort(addr, 30*time.Second) {
cleanup()
return nil, "", 0, fmt.Errorf("gateway did not start on %s", addr)
}
return cleanup, token, port, nil
}
// runChannelSetupPreflight prompts users to connect a messaging channel before
// starting the built-in gateway+TUI flow. In interactive sessions, it loops
// until a channel is configured, unless the user chooses "Set up later".
@@ -301,7 +341,7 @@ func (c *Openclaw) gatewayInfo() (token string, port int) {
}
func printOpenclawReady(bin, token string, port int, firstLaunch bool) {
u := fmt.Sprintf("http://localhost:%d", port)
u := fmt.Sprintf("http://127.0.0.1:%d", port)
if token != "" {
u += "/#token=" + url.QueryEscape(token)
}
@@ -339,9 +379,30 @@ func openclawEnv() []string {
env = append(env, e)
}
}
if _, ok := os.LookupEnv("OPENCLAW_PLUGIN_STAGE_DIR"); !ok {
if dir := openclawPluginStageDir(); dir != "" {
env = append(env, "OPENCLAW_PLUGIN_STAGE_DIR="+dir)
}
}
return env
}
func openclawInstallEnv() []string {
env := openclawEnv()
if _, ok := os.LookupEnv("OPENCLAW_EAGER_BUNDLED_PLUGIN_DEPS"); !ok {
env = append(env, "OPENCLAW_EAGER_BUNDLED_PLUGIN_DEPS=1")
}
return env
}
func openclawPluginStageDir() string {
home, err := os.UserHomeDir()
if err != nil {
return ""
}
return filepath.Join(home, ".openclaw", "plugin-runtime-deps")
}
// portOpen checks if a TCP port is currently accepting connections.
func portOpen(addr string) bool {
conn, err := net.DialTimeout("tcp", addr, 500*time.Millisecond)
@@ -565,6 +626,7 @@ func ensureOpenclawInstalled() (string, error) {
fmt.Fprintf(os.Stderr, "\nInstalling OpenClaw...\n")
cmd := exec.Command("npm", "install", "-g", "openclaw@latest")
cmd.Env = openclawInstallEnv()
cmd.Stdin = os.Stdin
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
@@ -738,89 +800,13 @@ func clearSessionModelOverride(primary string) {
_ = os.WriteFile(path, out, 0o600)
}
const (
webSearchNpmPackage = "@ollama/openclaw-web-search"
webSearchMinVersion = "0.2.1"
)
// ensureWebSearchPlugin installs the openclaw-web-search extension into the
// user-level extensions directory (~/.openclaw/extensions/) if it isn't already
// present, or re-installs if the installed version is older than webSearchMinVersion.
// Returns true if the extension is available.
func ensureWebSearchPlugin() bool {
home, err := os.UserHomeDir()
if err != nil {
return false
}
pluginDir := filepath.Join(home, ".openclaw", "extensions", "openclaw-web-search")
if webSearchPluginUpToDate(pluginDir) {
return true
}
npmBin, err := exec.LookPath("npm")
if err != nil {
return false
}
if err := os.MkdirAll(pluginDir, 0o755); err != nil {
return false
}
// Download the tarball via `npm pack`, extract it flat into the plugin dir.
pack := exec.Command(npmBin, "pack", webSearchNpmPackage, "--pack-destination", pluginDir)
out, err := pack.Output()
if err != nil {
fmt.Fprintf(os.Stderr, "%s Warning: could not download web search plugin: %v%s\n", ansiYellow, err, ansiReset)
return false
}
tgzName := strings.TrimSpace(string(out))
tgzPath := filepath.Join(pluginDir, tgzName)
defer os.Remove(tgzPath)
tar := exec.Command("tar", "xzf", tgzPath, "--strip-components=1", "-C", pluginDir)
if err := tar.Run(); err != nil {
fmt.Fprintf(os.Stderr, "%s Warning: could not extract web search plugin: %v%s\n", ansiYellow, err, ansiReset)
return false
}
fmt.Fprintf(os.Stderr, "%s ✓ Installed Ollama web search %s\n", ansiGreen, ansiReset)
return true
}
// webSearchPluginUpToDate returns true if the plugin is installed and its
// package.json version is >= webSearchMinVersion.
func webSearchPluginUpToDate(pluginDir string) bool {
data, err := os.ReadFile(filepath.Join(pluginDir, "package.json"))
if err != nil {
return false
}
var pkg struct {
Version string `json:"version"`
}
if json.Unmarshal(data, &pkg) != nil || pkg.Version == "" {
return false
}
return !versionLessThan(pkg.Version, webSearchMinVersion)
}
// versionLessThan compares two semver version strings (major.minor.patch).
// Inputs may omit the "v" prefix; it is added automatically for semver.Compare.
func versionLessThan(a, b string) bool {
if !strings.HasPrefix(a, "v") {
a = "v" + a
}
if !strings.HasPrefix(b, "v") {
b = "v" + b
}
return semver.Compare(a, b) < 0
}
// registerWebSearchPlugin adds plugins.entries.openclaw-web-search to the OpenClaw
// config so the gateway activates it on next start. Best-effort; silently returns
// on any error.
func registerWebSearchPlugin() {
// configureOllamaWebSearch keeps launch-managed OpenClaw installs on the
// bundled Ollama web_search provider. Older launch builds installed an
// external openclaw-web-search plugin that added custom ollama_web_search and
// ollama_web_fetch tools. Current OpenClaw versions ship Ollama web_search as
// the bundled "ollama" plugin instead, so we migrate stale config and ensure
// fresh installs select the bundled provider.
func configureOllamaWebSearch() {
home, err := os.UserHomeDir()
if err != nil {
return
@@ -835,6 +821,8 @@ func registerWebSearchPlugin() {
return
}
stalePluginConfigured := false
plugins, _ := config["plugins"].(map[string]any)
if plugins == nil {
plugins = make(map[string]any)
@@ -843,68 +831,100 @@ func registerWebSearchPlugin() {
if entries == nil {
entries = make(map[string]any)
}
entries["openclaw-web-search"] = map[string]any{"enabled": true}
plugins["entries"] = entries
// Pin trust so the gateway doesn't warn about untracked plugins.
allow, _ := plugins["allow"].([]any)
hasAllow := false
for _, v := range allow {
if s, ok := v.(string); ok && s == "openclaw-web-search" {
hasAllow = true
break
}
}
if !hasAllow {
allow = append(allow, "openclaw-web-search")
}
plugins["allow"] = allow
// Record install provenance so the loader can verify the plugin origin.
installs, _ := plugins["installs"].(map[string]any)
if installs == nil {
installs = make(map[string]any)
}
pluginDir := filepath.Join(home, ".openclaw", "extensions", "openclaw-web-search")
installs["openclaw-web-search"] = map[string]any{
"source": "npm",
"spec": webSearchNpmPackage,
"installPath": pluginDir,
}
plugins["installs"] = installs
config["plugins"] = plugins
// Add plugin tools to tools.alsoAllow so they survive the coding profile's
// policy pipeline (which has an explicit allow list of core tools only).
tools, _ := config["tools"].(map[string]any)
if tools == nil {
tools = make(map[string]any)
}
alsoAllow, _ := tools["alsoAllow"].([]any)
needed := []string{"ollama_web_search", "ollama_web_fetch"}
have := make(map[string]bool, len(alsoAllow))
for _, v := range alsoAllow {
if s, ok := v.(string); ok {
have[s] = true
}
}
for _, name := range needed {
if !have[name] {
alsoAllow = append(alsoAllow, name)
}
}
tools["alsoAllow"] = alsoAllow
// Disable built-in web search/fetch since our plugin replaces them.
web, _ := tools["web"].(map[string]any)
if web == nil {
web = make(map[string]any)
}
web["search"] = map[string]any{"enabled": false}
web["fetch"] = map[string]any{"enabled": false}
search, _ := web["search"].(map[string]any)
if search == nil {
search = make(map[string]any)
}
fetch, _ := web["fetch"].(map[string]any)
if fetch == nil {
fetch = make(map[string]any)
}
alsoAllow, _ := tools["alsoAllow"].([]any)
var filteredAlsoAllow []any
for _, v := range alsoAllow {
s, ok := v.(string)
if !ok {
filteredAlsoAllow = append(filteredAlsoAllow, v)
continue
}
if s == "ollama_web_search" || s == "ollama_web_fetch" {
stalePluginConfigured = true
continue
}
filteredAlsoAllow = append(filteredAlsoAllow, v)
}
if len(filteredAlsoAllow) > 0 {
tools["alsoAllow"] = filteredAlsoAllow
} else {
delete(tools, "alsoAllow")
}
if _, ok := entries["openclaw-web-search"]; ok {
delete(entries, "openclaw-web-search")
stalePluginConfigured = true
}
ollamaEntry, _ := entries["ollama"].(map[string]any)
if ollamaEntry == nil {
ollamaEntry = make(map[string]any)
}
ollamaEntry["enabled"] = true
entries["ollama"] = ollamaEntry
plugins["entries"] = entries
if allow, ok := plugins["allow"].([]any); ok {
var nextAllow []any
hasOllama := false
for _, v := range allow {
s, ok := v.(string)
if ok && s == "openclaw-web-search" {
stalePluginConfigured = true
continue
}
if ok && s == "ollama" {
hasOllama = true
}
nextAllow = append(nextAllow, v)
}
if !hasOllama {
nextAllow = append(nextAllow, "ollama")
}
plugins["allow"] = nextAllow
}
if installs, ok := plugins["installs"].(map[string]any); ok {
if _, exists := installs["openclaw-web-search"]; exists {
delete(installs, "openclaw-web-search")
stalePluginConfigured = true
}
if len(installs) > 0 {
plugins["installs"] = installs
} else {
delete(plugins, "installs")
}
}
if stalePluginConfigured || search["provider"] == nil {
search["provider"] = "ollama"
}
if stalePluginConfigured {
fetch["enabled"] = true
}
search["enabled"] = true
web["search"] = search
if len(fetch) > 0 {
web["fetch"] = fetch
}
tools["web"] = web
config["plugins"] = plugins
config["tools"] = tools
out, err := json.MarshalIndent(config, "", " ")
+444 -135
View File
@@ -251,6 +251,359 @@ func TestOpenclawRun_SetupLaterContinuesToGatewayAndTUI(t *testing.T) {
}
}
func TestOpenclawRun_FirstLaunchOnboardUsesLaunchManagedHealthFlow(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses a POSIX shell test binary")
}
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
t.Setenv("PATH", tmpDir)
bin := filepath.Join(tmpDir, "openclaw")
script := fmt.Sprintf(`#!/bin/sh
printf '%%s\n' "$*" >> "$HOME/invocations.log"
if [ "$1" = "onboard" ]; then
/usr/bin/env | /usr/bin/sort > "$HOME/onboard-env.log"
/bin/mkdir -p "$HOME/.openclaw"
/bin/cat > "$HOME/.openclaw/openclaw.json" <<'EOF'
{"wizard":{"lastRunAt":"2026-01-01T00:00:00Z"},"gateway":{"port":18789,"mode":"local"}}
EOF
fi
exit 0
`)
if err := os.WriteFile(bin, []byte(script), 0o755); err != nil {
t.Fatal(err)
}
oldConfirmPrompt := DefaultConfirmPrompt
DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) {
if prompt != "I understand the risks. Continue?" {
t.Fatalf("unexpected prompt: %q", prompt)
}
return true, nil
}
defer func() { DefaultConfirmPrompt = oldConfirmPrompt }()
c := &Openclaw{}
if err := c.Run("llama3.2", []string{"status"}); err != nil {
t.Fatalf("Run() error = %v", err)
}
data, err := os.ReadFile(filepath.Join(tmpDir, "invocations.log"))
if err != nil {
t.Fatal(err)
}
lines := strings.Split(strings.TrimSpace(string(data)), "\n")
if len(lines) < 2 {
t.Fatalf("expected onboard + passthrough invocations, got %v", lines)
}
onboardInvocation := ""
for _, line := range lines {
if strings.HasPrefix(line, "onboard ") {
onboardInvocation = line
break
}
}
if onboardInvocation == "" {
t.Fatalf("expected onboard invocation, got %v", lines)
}
if !strings.Contains(onboardInvocation, "--skip-health") {
t.Fatalf("expected onboard invocation to include --skip-health, got %q", onboardInvocation)
}
envData, err := os.ReadFile(filepath.Join(tmpDir, "onboard-env.log"))
if err != nil {
t.Fatal(err)
}
env := envSliceToMap(strings.Split(strings.TrimSpace(string(envData)), "\n"))
if env["OPENCLAW_EAGER_BUNDLED_PLUGIN_DEPS"] != "1" {
t.Fatalf("OPENCLAW_EAGER_BUNDLED_PLUGIN_DEPS = %q, want %q", env["OPENCLAW_EAGER_BUNDLED_PLUGIN_DEPS"], "1")
}
if env["OPENCLAW_PLUGIN_STAGE_DIR"] != filepath.Join(tmpDir, ".openclaw", "plugin-runtime-deps") {
t.Fatalf("OPENCLAW_PLUGIN_STAGE_DIR = %q, want %q", env["OPENCLAW_PLUGIN_STAGE_DIR"], filepath.Join(tmpDir, ".openclaw", "plugin-runtime-deps"))
}
}
func TestOpenclawRun_FirstLaunchTUIArgsEnsureGatewayBeforePassthrough(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses a POSIX shell test binary")
}
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
t.Setenv("PATH", tmpDir)
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
defer ln.Close()
port := ln.Addr().(*net.TCPAddr).Port
bin := filepath.Join(tmpDir, "openclaw")
script := fmt.Sprintf(`#!/bin/sh
printf '%%s\n' "$*" >> "$HOME/invocations.log"
if [ "$1" = "onboard" ]; then
/bin/mkdir -p "$HOME/.openclaw"
/bin/cat > "$HOME/.openclaw/openclaw.json" <<'EOF'
{"wizard":{"lastRunAt":"2026-01-01T00:00:00Z"},"gateway":{"port":%d,"mode":"local"}}
EOF
fi
exit 0
`, port)
if err := os.WriteFile(bin, []byte(script), 0o755); err != nil {
t.Fatal(err)
}
oldConfirmPrompt := DefaultConfirmPrompt
DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) {
if prompt != "I understand the risks. Continue?" {
t.Fatalf("unexpected prompt: %q", prompt)
}
return true, nil
}
defer func() { DefaultConfirmPrompt = oldConfirmPrompt }()
c := &Openclaw{}
if err := c.Run("llama3.2", []string{"tui"}); err != nil {
t.Fatalf("Run() error = %v", err)
}
data, err := os.ReadFile(filepath.Join(tmpDir, "invocations.log"))
if err != nil {
t.Fatal(err)
}
lines := strings.Split(strings.TrimSpace(string(data)), "\n")
if len(lines) < 3 {
t.Fatalf("expected at least 3 invocations (update, onboard, daemon restart, tui), got %v", lines)
}
onboardIdx, daemonRestartIdx, tuiIdx := -1, -1, -1
for i, line := range lines {
if onboardIdx == -1 && strings.HasPrefix(line, "onboard ") {
onboardIdx = i
}
if daemonRestartIdx == -1 && line == "daemon restart" {
daemonRestartIdx = i
}
if tuiIdx == -1 && line == "tui" {
tuiIdx = i
}
}
if onboardIdx == -1 {
t.Fatalf("expected an onboarding invocation, got %v", lines)
}
if daemonRestartIdx == -1 {
t.Fatalf("expected a daemon restart before tui, got %v", lines)
}
if tuiIdx == -1 {
t.Fatalf("expected a tui invocation, got %v", lines)
}
if !(onboardIdx < daemonRestartIdx && daemonRestartIdx < tuiIdx) {
t.Fatalf("expected onboarding, then daemon restart, then tui; got %v", lines)
}
}
func TestOpenclawEnsureGatewayReady_UsesDaemonStartFallback(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses a POSIX shell test binary")
}
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
t.Setenv("PATH", tmpDir)
portProbe, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
port := portProbe.Addr().(*net.TCPAddr).Port
_ = portProbe.Close()
configDir := filepath.Join(tmpDir, ".openclaw")
if err := os.MkdirAll(configDir, 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(configDir, "openclaw.json"), []byte(fmt.Sprintf(`{
"wizard": {"lastRunAt": "2026-01-01T00:00:00Z"},
"gateway": {"port": %d, "mode": "local"}
}`, port)), 0o644); err != nil {
t.Fatal(err)
}
bin := filepath.Join(tmpDir, "openclaw")
if err := os.WriteFile(bin, []byte("#!/bin/sh\nprintf '%s\\n' \"$*\" >> \"$HOME/invocations.log\"\n"), 0o755); err != nil {
t.Fatal(err)
}
oldCanInstallDaemon := openclawCanInstallDaemon
openclawCanInstallDaemon = func() bool { return true }
defer func() { openclawCanInstallDaemon = oldCanInstallDaemon }()
triggeredBy := make(chan string, 1)
listenerReady := make(chan net.Listener, 1)
go func() {
invocationsPath := filepath.Join(tmpDir, "invocations.log")
deadline := time.Now().Add(5 * time.Second)
for time.Now().Before(deadline) {
data, err := os.ReadFile(invocationsPath)
if err == nil {
lines := strings.Split(strings.TrimSpace(string(data)), "\n")
for _, line := range lines {
if line != "daemon start" && line != "gateway run --force" {
continue
}
ln, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port))
if err != nil {
return
}
go func() {
for {
conn, err := ln.Accept()
if err != nil {
return
}
_ = conn.Close()
}
}()
triggeredBy <- line
listenerReady <- ln
return
}
}
time.Sleep(10 * time.Millisecond)
}
}()
c := &Openclaw{}
cleanup, _, gotPort, err := c.ensureGatewayReady(bin)
if err != nil {
t.Fatalf("ensureGatewayReady() error = %v", err)
}
defer cleanup()
if gotPort != port {
t.Fatalf("ensureGatewayReady() port = %d, want %d", gotPort, port)
}
var ln net.Listener
select {
case which := <-triggeredBy:
if which != "daemon start" {
t.Fatalf("expected daemon start fallback, got %q", which)
}
case <-time.After(2 * time.Second):
t.Fatal("timed out waiting for gateway startup trigger")
}
select {
case ln = <-listenerReady:
defer ln.Close()
case <-time.After(2 * time.Second):
t.Fatal("timed out waiting for test listener")
}
data, err := os.ReadFile(filepath.Join(tmpDir, "invocations.log"))
if err != nil {
t.Fatal(err)
}
lines := strings.Split(strings.TrimSpace(string(data)), "\n")
if len(lines) == 0 || lines[0] != "daemon start" {
t.Fatalf("expected daemon start invocation, got %v", lines)
}
for _, line := range lines {
if line == "gateway run --force" {
t.Fatalf("did not expect gateway run fallback when daemon start succeeds, got %v", lines)
}
}
}
func TestOpenclawEnv_StagesBundledPluginRuntimeDeps(t *testing.T) {
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
t.Setenv("OPENAI_API_KEY", "should-be-cleared")
env := envSliceToMap(openclawEnv())
if env["OPENCLAW_PLUGIN_STAGE_DIR"] != filepath.Join(tmpDir, ".openclaw", "plugin-runtime-deps") {
t.Fatalf("OPENCLAW_PLUGIN_STAGE_DIR = %q, want %q", env["OPENCLAW_PLUGIN_STAGE_DIR"], filepath.Join(tmpDir, ".openclaw", "plugin-runtime-deps"))
}
if _, ok := env["OPENAI_API_KEY"]; ok {
t.Fatal("expected OPENAI_API_KEY to be cleared from openclaw environment")
}
}
func TestOpenclawInstallEnv_PreservesExplicitStageDirAndAddsEagerDeps(t *testing.T) {
t.Setenv("OPENCLAW_PLUGIN_STAGE_DIR", "/tmp/custom-stage")
env := envSliceToMap(openclawInstallEnv())
if env["OPENCLAW_PLUGIN_STAGE_DIR"] != "/tmp/custom-stage" {
t.Fatalf("OPENCLAW_PLUGIN_STAGE_DIR = %q, want %q", env["OPENCLAW_PLUGIN_STAGE_DIR"], "/tmp/custom-stage")
}
if env["OPENCLAW_EAGER_BUNDLED_PLUGIN_DEPS"] != "1" {
t.Fatalf("OPENCLAW_EAGER_BUNDLED_PLUGIN_DEPS = %q, want %q", env["OPENCLAW_EAGER_BUNDLED_PLUGIN_DEPS"], "1")
}
}
func TestEnsureOpenclawInstalled_UsesBundledPluginInstallEnv(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses a POSIX shell test binary")
}
tmpDir := t.TempDir()
setTestHome(t, tmpDir)
t.Setenv("PATH", tmpDir)
writeScript := func(path, content string) {
t.Helper()
if err := os.WriteFile(path, []byte(content), 0o755); err != nil {
t.Fatal(err)
}
}
openclawPath := filepath.Join(tmpDir, "openclaw")
npmScript := fmt.Sprintf(`#!/bin/sh
/usr/bin/env | /usr/bin/sort > "$HOME/npm-env.log"
/bin/cat > %q <<'EOF'
#!/bin/sh
exit 0
EOF
/bin/chmod +x %q
exit 0
`, openclawPath, openclawPath)
writeScript(filepath.Join(tmpDir, "npm"), npmScript)
writeScript(filepath.Join(tmpDir, "git"), "#!/bin/sh\nexit 0\n")
oldConfirmPrompt := DefaultConfirmPrompt
DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) {
if prompt != "OpenClaw is not installed. Install with npm?" {
t.Fatalf("unexpected prompt: %q", prompt)
}
return true, nil
}
defer func() { DefaultConfirmPrompt = oldConfirmPrompt }()
openclawFreshInstall = false
bin, err := ensureOpenclawInstalled()
if err != nil {
t.Fatalf("ensureOpenclawInstalled() error = %v", err)
}
if bin != "openclaw" {
t.Fatalf("ensureOpenclawInstalled() bin = %q, want %q", bin, "openclaw")
}
envData, err := os.ReadFile(filepath.Join(tmpDir, "npm-env.log"))
if err != nil {
t.Fatal(err)
}
env := envSliceToMap(strings.Split(strings.TrimSpace(string(envData)), "\n"))
if env["OPENCLAW_EAGER_BUNDLED_PLUGIN_DEPS"] != "1" {
t.Fatalf("OPENCLAW_EAGER_BUNDLED_PLUGIN_DEPS = %q, want %q", env["OPENCLAW_EAGER_BUNDLED_PLUGIN_DEPS"], "1")
}
if env["OPENCLAW_PLUGIN_STAGE_DIR"] != filepath.Join(tmpDir, ".openclaw", "plugin-runtime-deps") {
t.Fatalf("OPENCLAW_PLUGIN_STAGE_DIR = %q, want %q", env["OPENCLAW_PLUGIN_STAGE_DIR"], filepath.Join(tmpDir, ".openclaw", "plugin-runtime-deps"))
}
}
func TestOpenclawEdit(t *testing.T) {
c := &Openclaw{}
tmpDir := t.TempDir()
@@ -1227,6 +1580,18 @@ func TestOpenclawChannelsConfigured(t *testing.T) {
})
}
func envSliceToMap(entries []string) map[string]string {
env := make(map[string]string, len(entries))
for _, entry := range entries {
key, value, ok := strings.Cut(entry, "=")
if !ok {
continue
}
env[key] = value
}
return env
}
func TestOpenclawChannelSetupPreflight(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses a POSIX shell test binary")
@@ -1770,7 +2135,7 @@ func TestPrintOpenclawReady(t *testing.T) {
buf.ReadFrom(r)
output := buf.String()
if !strings.Contains(output, "localhost:9999") {
if !strings.Contains(output, "127.0.0.1:9999") {
t.Errorf("expected port 9999 in output, got:\n%s", output)
}
if strings.Contains(output, "#token=") {
@@ -2242,95 +2607,7 @@ func TestIntegrationOnboarded(t *testing.T) {
})
}
func TestVersionLessThan(t *testing.T) {
tests := []struct {
a, b string
want bool
}{
{"0.1.7", "0.2.1", true},
{"0.2.0", "0.2.1", true},
{"0.2.1", "0.2.1", false},
{"0.2.2", "0.2.1", false},
{"1.0.0", "0.2.1", false},
{"0.2.1", "1.0.0", true},
{"v0.1.7", "0.2.1", true},
{"0.2.1", "v0.2.1", false},
}
for _, tt := range tests {
t.Run(tt.a+"_vs_"+tt.b, func(t *testing.T) {
if got := versionLessThan(tt.a, tt.b); got != tt.want {
t.Errorf("versionLessThan(%q, %q) = %v, want %v", tt.a, tt.b, got, tt.want)
}
})
}
}
func TestWebSearchPluginUpToDate(t *testing.T) {
t.Run("missing directory", func(t *testing.T) {
if webSearchPluginUpToDate(filepath.Join(t.TempDir(), "nonexistent")) {
t.Error("expected false for missing directory")
}
})
t.Run("missing package.json", func(t *testing.T) {
dir := t.TempDir()
if webSearchPluginUpToDate(dir) {
t.Error("expected false for missing package.json")
}
})
t.Run("old version", func(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "package.json"), []byte(`{"version":"0.1.7"}`), 0o644); err != nil {
t.Fatal(err)
}
if webSearchPluginUpToDate(dir) {
t.Error("expected false for old version 0.1.7")
}
})
t.Run("exact minimum version", func(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "package.json"), []byte(`{"version":"0.2.1"}`), 0o644); err != nil {
t.Fatal(err)
}
if !webSearchPluginUpToDate(dir) {
t.Error("expected true for exact minimum version 0.2.1")
}
})
t.Run("newer version", func(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "package.json"), []byte(`{"version":"1.0.0"}`), 0o644); err != nil {
t.Fatal(err)
}
if !webSearchPluginUpToDate(dir) {
t.Error("expected true for newer version 1.0.0")
}
})
t.Run("invalid json", func(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "package.json"), []byte(`not json`), 0o644); err != nil {
t.Fatal(err)
}
if webSearchPluginUpToDate(dir) {
t.Error("expected false for invalid json")
}
})
t.Run("empty version", func(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "package.json"), []byte(`{"version":""}`), 0o644); err != nil {
t.Fatal(err)
}
if webSearchPluginUpToDate(dir) {
t.Error("expected false for empty version")
}
})
}
func TestRegisterWebSearchPlugin(t *testing.T) {
func TestConfigureOllamaWebSearch(t *testing.T) {
home := t.TempDir()
setTestHome(t, home)
@@ -2345,7 +2622,7 @@ func TestRegisterWebSearchPlugin(t *testing.T) {
t.Fatal(err)
}
registerWebSearchPlugin()
configureOllamaWebSearch()
data, err := os.ReadFile(configPath)
if err != nil {
@@ -2361,40 +2638,30 @@ func TestRegisterWebSearchPlugin(t *testing.T) {
t.Fatal("plugins section missing")
}
// Check entries
entries, _ := plugins["entries"].(map[string]any)
entry, _ := entries["openclaw-web-search"].(map[string]any)
entry, _ := entries["ollama"].(map[string]any)
if enabled, _ := entry["enabled"].(bool); !enabled {
t.Error("expected entries.openclaw-web-search.enabled = true")
t.Error("expected entries.ollama.enabled = true")
}
if _, ok := entries["openclaw-web-search"]; ok {
t.Error("expected stale openclaw-web-search entry to be absent")
}
// Check allow list
allow, _ := plugins["allow"].([]any)
found := false
for _, v := range allow {
if s, ok := v.(string); ok && s == "openclaw-web-search" {
found = true
}
if _, ok := plugins["allow"]; ok {
t.Error("did not expect plugins.allow to be created when no allowlist exists")
}
if !found {
t.Error("expected plugins.allow to contain openclaw-web-search")
if _, ok := plugins["installs"]; ok {
t.Error("did not expect plugins.installs to be created")
}
// Check install provenance
installs, _ := plugins["installs"].(map[string]any)
record, _ := installs["openclaw-web-search"].(map[string]any)
if record == nil {
t.Fatal("expected plugins.installs.openclaw-web-search")
tools, _ := config["tools"].(map[string]any)
web, _ := tools["web"].(map[string]any)
search, _ := web["search"].(map[string]any)
if got, _ := search["provider"].(string); got != "ollama" {
t.Errorf("search provider = %q, want %q", got, "ollama")
}
if source, _ := record["source"].(string); source != "npm" {
t.Errorf("install source = %q, want %q", source, "npm")
}
if spec, _ := record["spec"].(string); spec != webSearchNpmPackage {
t.Errorf("install spec = %q, want %q", spec, webSearchNpmPackage)
}
expectedPath := filepath.Join(home, ".openclaw", "extensions", "openclaw-web-search")
if installPath, _ := record["installPath"].(string); installPath != expectedPath {
t.Errorf("installPath = %q, want %q", installPath, expectedPath)
if enabled, _ := search["enabled"].(bool); !enabled {
t.Error("expected tools.web.search.enabled = true")
}
})
@@ -2403,8 +2670,8 @@ func TestRegisterWebSearchPlugin(t *testing.T) {
t.Fatal(err)
}
registerWebSearchPlugin()
registerWebSearchPlugin()
configureOllamaWebSearch()
configureOllamaWebSearch()
data, err := os.ReadFile(configPath)
if err != nil {
@@ -2416,30 +2683,39 @@ func TestRegisterWebSearchPlugin(t *testing.T) {
}
plugins, _ := config["plugins"].(map[string]any)
allow, _ := plugins["allow"].([]any)
count := 0
for _, v := range allow {
if s, ok := v.(string); ok && s == "openclaw-web-search" {
count++
}
entries, _ := plugins["entries"].(map[string]any)
if len(entries) != 1 {
t.Fatalf("expected only bundled ollama entry, got %v", entries)
}
if count != 1 {
t.Errorf("expected exactly 1 openclaw-web-search in allow, got %d", count)
if _, ok := entries["ollama"]; !ok {
t.Fatalf("expected entries.ollama to exist, got %v", entries)
}
})
t.Run("preserves existing config", func(t *testing.T) {
t.Run("migrates stale plugin config and preserves unrelated settings", func(t *testing.T) {
initial := map[string]any{
"plugins": map[string]any{
"allow": []any{"some-other-plugin"},
"allow": []any{"some-other-plugin", "openclaw-web-search"},
"entries": map[string]any{
"some-other-plugin": map[string]any{"enabled": true},
"some-other-plugin": map[string]any{"enabled": true},
"openclaw-web-search": map[string]any{"enabled": true},
},
"installs": map[string]any{
"some-other-plugin": map[string]any{
"source": "npm",
"installPath": "/some/path",
},
"openclaw-web-search": map[string]any{
"source": "npm",
"installPath": "/old/path",
},
},
},
"tools": map[string]any{
"alsoAllow": []any{"ollama_web_search", "ollama_web_fetch", "browser"},
"web": map[string]any{
"search": map[string]any{"enabled": false},
"fetch": map[string]any{"enabled": false},
},
},
"customField": "preserved",
@@ -2449,7 +2725,7 @@ func TestRegisterWebSearchPlugin(t *testing.T) {
t.Fatal(err)
}
registerWebSearchPlugin()
configureOllamaWebSearch()
out, err := os.ReadFile(configPath)
if err != nil {
@@ -2469,28 +2745,61 @@ func TestRegisterWebSearchPlugin(t *testing.T) {
if entries["some-other-plugin"] == nil {
t.Error("existing plugin entry was lost")
}
if entries["openclaw-web-search"] != nil {
t.Error("stale openclaw-web-search entry should be removed")
}
if ollamaEntry, _ := entries["ollama"].(map[string]any); ollamaEntry == nil {
t.Fatal("expected bundled ollama entry to be enabled")
}
installs, _ := plugins["installs"].(map[string]any)
if installs["some-other-plugin"] == nil {
t.Error("existing install record was lost")
}
if installs["openclaw-web-search"] != nil {
t.Error("stale openclaw-web-search install record should be removed")
}
allow, _ := plugins["allow"].([]any)
hasOther, hasWebSearch := false, false
hasOther, hasStalePlugin, hasOllama := false, false, false
for _, v := range allow {
s, _ := v.(string)
if s == "some-other-plugin" {
hasOther = true
}
if s == "openclaw-web-search" {
hasWebSearch = true
hasStalePlugin = true
}
if s == "ollama" {
hasOllama = true
}
}
if !hasOther {
t.Error("existing allow entry was lost")
}
if !hasWebSearch {
t.Error("openclaw-web-search not added to allow")
if hasStalePlugin {
t.Error("stale openclaw-web-search allow entry should be removed")
}
if !hasOllama {
t.Error("expected plugins.allow to contain bundled ollama plugin")
}
tools, _ := config["tools"].(map[string]any)
alsoAllow, _ := tools["alsoAllow"].([]any)
if len(alsoAllow) != 1 || alsoAllow[0] != "browser" {
t.Errorf("expected stale custom web tools to be removed, got %v", alsoAllow)
}
web, _ := tools["web"].(map[string]any)
search, _ := web["search"].(map[string]any)
fetch, _ := web["fetch"].(map[string]any)
if got, _ := search["provider"].(string); got != "ollama" {
t.Errorf("search provider = %q, want %q", got, "ollama")
}
if enabled, _ := search["enabled"].(bool); !enabled {
t.Error("expected migrated tools.web.search.enabled = true")
}
if enabled, _ := fetch["enabled"].(bool); !enabled {
t.Error("expected migrated tools.web.fetch.enabled = true")
}
})
}
+51
View File
@@ -0,0 +1,51 @@
package launch
import (
"fmt"
"os"
"os/exec"
"runtime"
"github.com/ollama/ollama/envconfig"
)
// Poolside implements Runner for Poolside's CLI.
type Poolside struct{}
var poolsideGOOS = runtime.GOOS
func (p *Poolside) String() string { return "Pool" }
func poolsideUnsupportedError() error {
return fmt.Errorf("Warning: Poolside is not currently supported on Windows")
}
func (p *Poolside) args(model string, extra []string) []string {
var args []string
if model != "" {
args = append(args, "-m", model)
}
args = append(args, extra...)
return args
}
func (p *Poolside) Run(model string, args []string) error {
if poolsideGOOS == "windows" {
return poolsideUnsupportedError()
}
bin, err := exec.LookPath("pool")
if err != nil {
return fmt.Errorf("pool is not installed")
}
cmd := exec.Command(bin, p.args(model, args)...)
cmd.Stdin = os.Stdin
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
cmd.Env = append(os.Environ(),
"POOLSIDE_STANDALONE_BASE_URL="+envconfig.Host().String()+"/v1",
"POOLSIDE_API_KEY=ollama",
)
return cmd.Run()
}
+88
View File
@@ -0,0 +1,88 @@
package launch
import (
"os"
"path/filepath"
"runtime"
"slices"
"strings"
"testing"
)
func TestPoolsideArgs(t *testing.T) {
p := &Poolside{}
tests := []struct {
name string
model string
extra []string
want []string
}{
{name: "with model", model: "qwen3.5", want: []string{"-m", "qwen3.5"}},
{name: "without model", extra: []string{"session"}, want: []string{"session"}},
{name: "with model and extra args", model: "llama3.2", extra: []string{"--foo", "bar"}, want: []string{"-m", "llama3.2", "--foo", "bar"}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := p.args(tt.model, tt.extra)
if !slices.Equal(got, tt.want) {
t.Fatalf("args(%q, %v) = %v, want %v", tt.model, tt.extra, got, tt.want)
}
})
}
}
func TestPoolsideRunSetsOllamaEnv(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses POSIX shell fake binary")
}
tmpDir := t.TempDir()
logPath := filepath.Join(tmpDir, "pool.log")
poolPath := filepath.Join(tmpDir, "pool")
script := "#!/bin/sh\n" +
"printf 'base=%s\\nkey=%s\\nargs=%s\\n' \"$POOLSIDE_STANDALONE_BASE_URL\" \"$POOLSIDE_API_KEY\" \"$*\" > \"" + logPath + "\"\n"
if err := os.WriteFile(poolPath, []byte(script), 0o755); err != nil {
t.Fatalf("failed to write fake pool binary: %v", err)
}
t.Setenv("PATH", tmpDir)
t.Setenv("OLLAMA_HOST", "http://127.0.0.1:11434")
p := &Poolside{}
if err := p.Run("qwen3.5", []string{"session"}); err != nil {
t.Fatalf("Run returned error: %v", err)
}
data, err := os.ReadFile(logPath)
if err != nil {
t.Fatalf("failed to read pool log: %v", err)
}
got := string(data)
if !strings.Contains(got, "base=http://127.0.0.1:11434/v1") {
t.Fatalf("expected Poolside base URL override in log, got:\n%s", got)
}
if !strings.Contains(got, "key=ollama") {
t.Fatalf("expected Poolside API key override in log, got:\n%s", got)
}
if !strings.Contains(got, "args=-m qwen3.5 session") {
t.Fatalf("expected model and extra args in log, got:\n%s", got)
}
}
func TestPoolsideRunWindowsUnsupported(t *testing.T) {
prev := poolsideGOOS
poolsideGOOS = "windows"
t.Cleanup(func() { poolsideGOOS = prev })
p := &Poolside{}
err := p.Run("kimi-k2.6:cloud", nil)
if err == nil {
t.Fatal("expected Windows unsupported error")
}
if !strings.Contains(err.Error(), "not currently supported on Windows") {
t.Fatalf("expected Windows warning, got %v", err)
}
}
+71 -1
View File
@@ -33,7 +33,7 @@ type IntegrationInfo struct {
Description string
}
var launcherIntegrationOrder = []string{"openclaw", "claude", "opencode", "hermes", "codex", "droid", "pi"}
var launcherIntegrationOrder = []string{"claude-desktop", "claude", "openclaw", "hermes", "opencode", "codex", "copilot", "droid", "pi", "pool"}
var integrationSpecs = []*IntegrationSpec{
{
@@ -48,6 +48,18 @@ var integrationSpecs = []*IntegrationSpec{
URL: "https://code.claude.com/docs/en/quickstart",
},
},
{
Name: "claude-desktop",
Runner: &ClaudeDesktop{},
Aliases: []string{"claude-app"},
Description: "Claude Desktop with Ollama Cloud",
Install: IntegrationInstallSpec{
CheckInstalled: func() bool {
return claudeDesktopInstalled()
},
URL: "https://claude.com/download",
},
},
{
Name: "cline",
Runner: &Cline{},
@@ -74,6 +86,36 @@ var integrationSpecs = []*IntegrationSpec{
Command: []string{"npm", "install", "-g", "@openai/codex"},
},
},
{
Name: "kimi",
Runner: &Kimi{},
Description: "Moonshot's coding agent for terminal and IDEs",
Hidden: true,
Install: IntegrationInstallSpec{
CheckInstalled: func() bool {
_, err := exec.LookPath("kimi")
return err == nil
},
EnsureInstalled: func() error {
_, err := ensureKimiInstalled()
return err
},
URL: "https://moonshotai.github.io/kimi-cli/en/guides/getting-started.html",
},
},
{
Name: "copilot",
Runner: &Copilot{},
Aliases: []string{"copilot-cli"},
Description: "GitHub's AI coding agent for the terminal",
Install: IntegrationInstallSpec{
CheckInstalled: func() bool {
_, err := (&Copilot{}).findPath()
return err == nil
},
URL: "https://github.com/features/copilot/cli/",
},
},
{
Name: "droid",
Runner: &Droid{},
@@ -136,6 +178,18 @@ var integrationSpecs = []*IntegrationSpec{
Command: []string{"npm", "install", "-g", "@mariozechner/pi-coding-agent@latest"},
},
},
{
Name: "pool",
Runner: &Poolside{},
Description: "Poolside's software agent for enterprise development",
Install: IntegrationInstallSpec{
CheckInstalled: func() bool {
_, err := exec.LookPath("pool")
return err == nil
},
URL: "https://github.com/poolsideai/pool",
},
},
{
Name: "hermes",
Runner: &Hermes{},
@@ -255,6 +309,12 @@ func ListVisibleIntegrationSpecs() []IntegrationSpec {
if spec.Hidden {
continue
}
if supported, ok := spec.Runner.(SupportedIntegration); ok && supported.Supported() != nil {
continue
}
if spec.Name == "pool" && poolsideGOOS == "windows" {
continue
}
visible = append(visible, *spec)
}
@@ -369,6 +429,16 @@ func EnsureIntegrationInstalled(name string, runner Runner) error {
return fmt.Errorf("%s is not installed", runner)
}
if supported, ok := runner.(SupportedIntegration); ok {
if err := supported.Supported(); err != nil {
return err
}
}
if integration.spec.Name == "pool" && poolsideGOOS == "windows" {
return poolsideUnsupportedError()
}
if integration.installed {
return nil
}
+24
View File
@@ -45,10 +45,30 @@ func TestEditorRunsDoNotRewriteConfig(t *testing.T) {
return filepath.Join(home, ".pi", "agent", "models.json")
},
},
{
name: "pool",
binary: "pool",
runner: &Poolside{},
checkPath: func(home string) string {
return filepath.Join(home, ".poolside", "config")
},
},
{
name: "kimi",
binary: "kimi",
runner: &Kimi{},
checkPath: func(home string) string {
return filepath.Join(home, ".kimi", "config.toml")
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if tt.name == "pool" && poolsideGOOS == "windows" {
t.Skip("Poolside is intentionally unsupported on Windows")
}
home := t.TempDir()
setTestHome(t, home)
@@ -57,6 +77,10 @@ func TestEditorRunsDoNotRewriteConfig(t *testing.T) {
if tt.name == "pi" {
writeFakeBinary(t, binDir, "npm")
}
if tt.name == "kimi" {
writeFakeBinary(t, binDir, "curl")
writeFakeBinary(t, binDir, "bash")
}
t.Setenv("PATH", binDir)
configPath := tt.checkPath(home)
-1
View File
@@ -25,7 +25,6 @@ func TestRenderSignIn_ContainsURL(t *testing.T) {
}
}
func TestRenderSignIn_ContainsSpinner(t *testing.T) {
got := renderSignIn("test:cloud", "https://example.com", 0, 80)
if !strings.Contains(got, "Waiting for sign in to complete") {
+1 -1
View File
@@ -45,7 +45,7 @@ type menuItem struct {
isOthers bool
}
const pinnedIntegrationCount = 3
const pinnedIntegrationCount = 4
var runModelMenuItem = menuItem{
title: "Chat with a model",
+46 -8
View File
@@ -44,6 +44,13 @@ func launcherTestState() *launch.LauncherState {
Selectable: true,
Changeable: true,
},
"claude-desktop": {
Name: "claude-desktop",
DisplayName: "Claude Desktop",
Description: "Claude Desktop with Ollama Cloud",
Selectable: true,
Changeable: true,
},
"hermes": {
Name: "hermes",
DisplayName: "Hermes Agent",
@@ -97,18 +104,45 @@ func compareStrings(got, want []string) string {
return cmp.Diff(want, got)
}
func expectedCollapsedSequence(state *launch.LauncherState) []string {
sequence := []string{"run"}
for _, item := range pinnedIntegrationItems(state) {
sequence = append(sequence, item.integration)
}
if len(otherIntegrationItems(state)) > 0 {
sequence = append(sequence, "more")
}
return sequence
}
func expectedExpandedSequence(state *launch.LauncherState) []string {
sequence := []string{"run"}
for _, item := range pinnedIntegrationItems(state) {
sequence = append(sequence, item.integration)
}
for _, item := range otherIntegrationItems(state) {
sequence = append(sequence, item.integration)
}
return sequence
}
func TestMenuRendersPinnedItemsAndMore(t *testing.T) {
menu := newModel(launcherTestState())
state := launcherTestState()
menu := newModel(state)
view := menu.View()
for _, want := range []string{"Chat with a model", "Launch OpenClaw", "Launch Claude Code", "Launch OpenCode", "More..."} {
for _, want := range []string{"Chat with a model", "Launch Claude Code", "Launch OpenClaw", "Launch Hermes Agent", "More..."} {
if !strings.Contains(view, want) {
t.Fatalf("expected menu view to contain %q\n%s", want, view)
}
}
if strings.Contains(view, "Launch Codex") {
t.Fatalf("expected Codex to be under More, not pinned\n%s", view)
if findMenuCursorByIntegration(menu.items, "claude-desktop") == -1 {
if strings.Contains(view, "Launch Claude Desktop") {
t.Fatalf("expected Claude Desktop to be hidden on unsupported platforms\n%s", view)
}
} else if !strings.Contains(view, "Launch Claude Desktop") {
t.Fatalf("expected menu view to contain Claude Desktop\n%s", view)
}
wantOrder := []string{"run", "openclaw", "claude", "opencode", "more"}
wantOrder := expectedCollapsedSequence(state)
if diff := compareStrings(integrationSequence(menu.items), wantOrder); diff != "" {
t.Fatalf("unexpected pinned order: %s", diff)
}
@@ -116,20 +150,24 @@ func TestMenuRendersPinnedItemsAndMore(t *testing.T) {
func TestMenuExpandsOthersFromLastSelection(t *testing.T) {
state := launcherTestState()
state.LastSelection = "codex"
overflow := otherIntegrationItems(state)
if len(overflow) == 0 {
t.Fatal("expected at least one overflow integration")
}
state.LastSelection = overflow[0].integration
menu := newModel(state)
if !menu.showOthers {
t.Fatal("expected others section to expand when last selection is in the overflow list")
}
view := menu.View()
if !strings.Contains(view, "Launch Codex") {
if !strings.Contains(view, overflow[0].title) {
t.Fatalf("expected expanded view to contain overflow integration\n%s", view)
}
if strings.Contains(view, "More...") {
t.Fatalf("expected expanded view to replace More... item\n%s", view)
}
wantOrder := []string{"run", "openclaw", "claude", "opencode", "hermes", "codex", "droid", "pi"}
wantOrder := expectedExpandedSequence(state)
if diff := compareStrings(integrationSequence(menu.items), wantOrder); diff != "" {
t.Fatalf("unexpected expanded order: %s", diff)
}
+8
View File
@@ -316,6 +316,8 @@ func LoadModelMetadata(fsys fs.FS) (ModelKV, *Tokenizer, error) {
conv = &deepseek2Model{}
case "Glm4MoeLiteForCausalLM":
conv = &glm4MoeLiteModel{}
case "LagunaForCausalLM":
conv = &lagunaModel{}
case "GlmOcrForConditionalGeneration":
conv = &glmOcrModel{}
case "Lfm2ForCausalLM", "Lfm2MoeForCausalLM":
@@ -324,6 +326,8 @@ func LoadModelMetadata(fsys fs.FS) (ModelKV, *Tokenizer, error) {
conv = &lfm2VLTextModel{}
case "Qwen3NextForCausalLM", "Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration":
conv = &qwen3NextModel{}
case "NemotronH_Nano_VL_V2", "NemotronH_Nano_Omni_Reasoning_V3":
conv = &nemotronHNanoVLModel{}
case "NemotronHForCausalLM":
conv = &nemotronHModel{}
default:
@@ -387,6 +391,10 @@ func ConvertModel(fsys fs.FS, f *os.File) error {
}
func writeFile(f *os.File, kv KV, ts []*ggml.Tensor) error {
for k, v := range sourceTensorKV(ts) {
kv[k] = v
}
for i := range ts {
ts[i].Shape = slices.Clone(ts[i].Shape)
slices.Reverse(ts[i].Shape)
+604
View File
@@ -0,0 +1,604 @@
package convert
import (
"cmp"
"encoding/json"
"fmt"
iofs "io/fs"
"math"
"strings"
"github.com/ollama/ollama/fs/ggml"
)
type lagunaModel struct {
ModelParameters
NumHiddenLayers uint32 `json:"num_hidden_layers"`
HiddenSize uint32 `json:"hidden_size"`
IntermediateSize uint32 `json:"intermediate_size"`
NumAttentionHeads uint32 `json:"num_attention_heads"`
NumKeyValueHeads uint32 `json:"num_key_value_heads"`
HeadDim uint32 `json:"head_dim"`
RMSNormEPS float32 `json:"rms_norm_eps"`
MaxPositionEmbeddings uint32 `json:"max_position_embeddings"`
SlidingWindow uint32 `json:"sliding_window"`
PartialRotaryFactor float32 `json:"partial_rotary_factor"`
Gating lagunaGatingMode `json:"gating"`
QKNormType string `json:"qk_norm_type"`
LayerTypes []string `json:"layer_types"`
NumAttentionHeadsPerLayer []uint32 `json:"num_attention_heads_per_layer"`
NumExperts uint32 `json:"num_experts"`
NumExpertsPerTok uint32 `json:"num_experts_per_tok"`
MoEIntermediateSize uint32 `json:"moe_intermediate_size"`
SharedExpertIntermediateSize uint32 `json:"shared_expert_intermediate_size"`
NormTopKProb bool `json:"norm_topk_prob"`
MoeRoutedScalingFactor float32 `json:"moe_routed_scaling_factor"`
MoERouterUseSigmoid bool `json:"moe_router_use_sigmoid"`
MoEApplyRouterWeightOnInput bool `json:"moe_apply_router_weight_on_input"`
DecoderSparseStep uint32 `json:"decoder_sparse_step"`
MLPOnlyLayers []uint32 `json:"mlp_only_layers"`
MLPLayerTypes []string `json:"mlp_layer_types"`
RopeParameters lagunaRopeParameters `json:"rope_parameters"`
SwaRopeParameters lagunaRopeParameters `json:"swa_rope_parameters"`
SwaAttentionSinkEnabled bool `json:"swa_attention_sink_enabled"`
}
type lagunaGatingMode string
type lagunaRopeParameters struct {
RopeTheta float32 `json:"rope_theta"`
RopeType string `json:"rope_type"`
Type string `json:"type"`
Factor float32 `json:"factor"`
OriginalMaxPositionEmbeddings uint32 `json:"original_max_position_embeddings"`
BetaSlow float32 `json:"beta_slow"`
BetaFast float32 `json:"beta_fast"`
AttentionFactor float32 `json:"attention_factor"`
PartialRotaryFactor float32 `json:"partial_rotary_factor"`
}
type lagunaRopeConfig struct {
flat lagunaRopeParameters
full lagunaRopeParameters
sliding lagunaRopeParameters
nested bool
}
func (g *lagunaGatingMode) UnmarshalJSON(b []byte) error {
var s string
if err := json.Unmarshal(b, &s); err == nil {
*g = lagunaGatingMode(s)
return nil
}
var enabled bool
if err := json.Unmarshal(b, &enabled); err == nil {
if enabled {
*g = "true"
} else {
*g = "false"
}
return nil
}
if string(b) == "null" {
return nil
}
return fmt.Errorf("unsupported Laguna gating JSON value %s", string(b))
}
func (g lagunaGatingMode) perHead() bool {
return strings.EqualFold(string(g), "per-head") || strings.EqualFold(string(g), "true")
}
func (r *lagunaRopeConfig) UnmarshalJSON(b []byte) error {
if string(b) == "null" {
return nil
}
var probe map[string]json.RawMessage
if err := json.Unmarshal(b, &probe); err != nil {
return err
}
if len(probe) == 0 {
return nil
}
if raw, ok := probe["full_attention"]; ok {
r.nested = true
if err := json.Unmarshal(raw, &r.full); err != nil {
return err
}
if raw = probe["sliding_attention"]; raw != nil {
if err := json.Unmarshal(raw, &r.sliding); err != nil {
return err
}
}
return nil
}
if raw, ok := probe["global_attention"]; ok {
r.nested = true
if err := json.Unmarshal(raw, &r.full); err != nil {
return err
}
if raw = probe["sliding_attention"]; raw != nil {
if err := json.Unmarshal(raw, &r.sliding); err != nil {
return err
}
}
return nil
}
return json.Unmarshal(b, &r.flat)
}
func (r lagunaRopeConfig) fullParams() lagunaRopeParameters {
if r.nested {
return r.full
}
return r.flat
}
func (r lagunaRopeConfig) slidingParams() (lagunaRopeParameters, bool) {
if !r.nested {
return lagunaRopeParameters{}, false
}
return r.sliding, true
}
func (r lagunaRopeParameters) ropeType() string {
return cmp.Or(r.RopeType, r.Type)
}
func (r lagunaRopeParameters) withDefaultPartialRotaryFactor(v float32) lagunaRopeParameters {
if r.PartialRotaryFactor == 0 {
r.PartialRotaryFactor = v
}
return r
}
func (r lagunaRopeParameters) empty() bool {
return r == (lagunaRopeParameters{})
}
type rawLagunaModel struct {
ModelParameters
NumHiddenLayers uint32 `json:"num_hidden_layers"`
HiddenSize uint32 `json:"hidden_size"`
IntermediateSize uint32 `json:"intermediate_size"`
NumAttentionHeads uint32 `json:"num_attention_heads"`
NumKeyValueHeads uint32 `json:"num_key_value_heads"`
HeadDim uint32 `json:"head_dim"`
RMSNormEPS float32 `json:"rms_norm_eps"`
MaxPositionEmbeddings uint32 `json:"max_position_embeddings"`
SlidingWindow uint32 `json:"sliding_window"`
PartialRotaryFactor float32 `json:"partial_rotary_factor"`
Gating lagunaGatingMode `json:"gating"`
QKNormType string `json:"qk_norm_type"`
LayerTypes []string `json:"layer_types"`
NumAttentionHeadsPerLayer []uint32 `json:"num_attention_heads_per_layer"`
NumExperts uint32 `json:"num_experts"`
NumExpertsPerTok uint32 `json:"num_experts_per_tok"`
MoEIntermediateSize uint32 `json:"moe_intermediate_size"`
SharedExpertIntermediateSize uint32 `json:"shared_expert_intermediate_size"`
NormTopKProb *bool `json:"norm_topk_prob"`
MoeRoutedScalingFactor float32 `json:"moe_routed_scaling_factor"`
MoERouterUseSigmoid *bool `json:"moe_router_use_sigmoid"`
MoEApplyRouterWeightOnInput bool `json:"moe_apply_router_weight_on_input"`
DecoderSparseStep uint32 `json:"decoder_sparse_step"`
MLPOnlyLayers []uint32 `json:"mlp_only_layers"`
MLPLayerTypes []string `json:"mlp_layer_types"`
RopeParameters lagunaRopeConfig `json:"rope_parameters"`
SwaRopeParameters lagunaRopeParameters `json:"swa_rope_parameters"`
SwaAttentionSinkEnabled bool `json:"swa_attention_sink_enabled"`
}
func (p *lagunaModel) UnmarshalJSON(b []byte) error {
var raw rawLagunaModel
if err := json.Unmarshal(b, &raw); err != nil {
return err
}
mlpOnlyLayers, err := lagunaDenseLayers(raw.MLPOnlyLayers, raw.MLPLayerTypes)
if err != nil {
return err
}
fullRope := raw.RopeParameters.fullParams().withDefaultPartialRotaryFactor(cmp.Or(raw.PartialRotaryFactor, float32(1)))
swaRope := raw.SwaRopeParameters
if nestedSwa, ok := raw.RopeParameters.slidingParams(); ok && !nestedSwa.empty() {
swaRope = nestedSwa
}
swaRope = swaRope.withDefaultPartialRotaryFactor(cmp.Or(fullRope.PartialRotaryFactor, float32(1)))
*p = lagunaModel{
ModelParameters: raw.ModelParameters,
NumHiddenLayers: raw.NumHiddenLayers,
HiddenSize: raw.HiddenSize,
IntermediateSize: raw.IntermediateSize,
NumAttentionHeads: raw.NumAttentionHeads,
NumKeyValueHeads: raw.NumKeyValueHeads,
HeadDim: raw.HeadDim,
RMSNormEPS: raw.RMSNormEPS,
MaxPositionEmbeddings: raw.MaxPositionEmbeddings,
SlidingWindow: raw.SlidingWindow,
PartialRotaryFactor: cmp.Or(raw.PartialRotaryFactor, fullRope.PartialRotaryFactor),
Gating: raw.Gating,
QKNormType: cmp.Or(raw.QKNormType, "rmsnorm"),
LayerTypes: raw.LayerTypes,
NumAttentionHeadsPerLayer: raw.NumAttentionHeadsPerLayer,
NumExperts: raw.NumExperts,
NumExpertsPerTok: raw.NumExpertsPerTok,
MoEIntermediateSize: raw.MoEIntermediateSize,
SharedExpertIntermediateSize: raw.SharedExpertIntermediateSize,
NormTopKProb: defaultBool(raw.NormTopKProb, true),
MoeRoutedScalingFactor: raw.MoeRoutedScalingFactor,
MoERouterUseSigmoid: defaultBool(raw.MoERouterUseSigmoid, true),
MoEApplyRouterWeightOnInput: raw.MoEApplyRouterWeightOnInput,
DecoderSparseStep: raw.DecoderSparseStep,
MLPOnlyLayers: mlpOnlyLayers,
MLPLayerTypes: raw.MLPLayerTypes,
RopeParameters: fullRope,
SwaRopeParameters: swaRope,
SwaAttentionSinkEnabled: raw.SwaAttentionSinkEnabled,
}
return nil
}
func defaultBool(v *bool, fallback bool) bool {
if v == nil {
return fallback
}
return *v
}
const (
lagunaGatingFuncSoftmax uint32 = 1
lagunaGatingFuncSigmoid uint32 = 2
lagunaLayerTypeGlobal uint32 = 0
lagunaLayerTypeSliding uint32 = 1
)
func (p *lagunaModel) KV(t *Tokenizer) KV {
kv := p.ModelParameters.KV(t)
kv["general.architecture"] = "laguna"
// Laguna's chat template and built-in renderer both emit the leading
// special token explicitly. Auto-prepending BOS here would duplicate it.
kv["tokenizer.ggml.add_bos_token"] = false
kv["tokenizer.ggml.pre"] = "laguna"
// Laguna does not need tokenizer.chat_template at runtime: Ollama create
// sets the Laguna renderer/parser from the architecture, and the renderer
// owns prompt formatting.
delete(kv, "tokenizer.chat_template")
kv["laguna.block_count"] = p.NumHiddenLayers
kv["laguna.context_length"] = p.MaxPositionEmbeddings
kv["laguna.embedding_length"] = p.HiddenSize
kv["laguna.feed_forward_length"] = p.IntermediateSize
if len(p.NumAttentionHeadsPerLayer) == int(p.NumHiddenLayers) {
kv["laguna.attention.head_count"] = p.NumAttentionHeadsPerLayer
} else {
kv["laguna.attention.head_count"] = p.NumAttentionHeads
}
kv["laguna.attention.head_count_kv"] = p.NumKeyValueHeads
kv["laguna.attention.key_length"] = p.HeadDim
kv["laguna.attention.value_length"] = p.HeadDim
kv["laguna.attention.layer_norm_rms_epsilon"] = p.RMSNormEPS
kv["laguna.attention.sliding_window"] = p.SlidingWindow
kv["laguna.attention.sink_enabled"] = p.SwaAttentionSinkEnabled
if len(p.LayerTypes) > 0 {
encoded := make([]uint32, len(p.LayerTypes))
slidingPattern := make([]bool, len(p.LayerTypes))
for i, layerType := range p.LayerTypes {
if lagunaLayerIsSliding(layerType) {
encoded[i] = lagunaLayerTypeSliding
slidingPattern[i] = true
} else {
encoded[i] = lagunaLayerTypeGlobal
}
}
kv["laguna.attention.layer_types"] = encoded
kv["laguna.attention.sliding_window_pattern"] = slidingPattern
}
if p.Gating.perHead() {
kv["laguna.attention.gating_type"] = uint32(1)
} else {
kv["laguna.attention.gating_type"] = uint32(0)
}
kv["laguna.attention.qk_norm"] = p.QKNormType == "rmsnorm"
kv["laguna.expert_count"] = p.NumExperts
kv["laguna.expert_used_count"] = p.NumExpertsPerTok
kv["laguna.expert_feed_forward_length"] = p.MoEIntermediateSize
kv["laguna.expert_shared_feed_forward_length"] = p.SharedExpertIntermediateSize
kv["laguna.expert_shared_count"] = uint32(1)
kv["laguna.expert_weights_norm"] = p.NormTopKProb
kv["laguna.expert_weights_scale"] = p.MoeRoutedScalingFactor
kv["laguna.expert_gating_func"] = lagunaMoeGatingFunc(p.MoERouterUseSigmoid)
kv["laguna.decoder_sparse_step"] = cmp.Or(p.DecoderSparseStep, uint32(1))
if leading, ok := lagunaLeadingDensePrefix(p.MLPOnlyLayers); ok {
kv["laguna.leading_dense_block_count"] = leading
}
if len(p.MLPOnlyLayers) > 0 {
kv["laguna.dense_layers"] = p.MLPOnlyLayers
}
ropeType := p.RopeParameters.ropeType()
kv["laguna.rope.freq_base"] = cmp.Or(p.RopeParameters.RopeTheta, float32(10000))
kv["laguna.rope.scaling.type"] = ropeType
ropeFactor := cmp.Or(p.RopeParameters.Factor, float32(1))
kv["laguna.rope.scaling.factor"] = ropeFactor
kv["laguna.rope.scaling.original_context_length"] = p.RopeParameters.OriginalMaxPositionEmbeddings
kv["laguna.rope.scaling.beta_fast"] = p.RopeParameters.BetaFast
kv["laguna.rope.scaling.beta_slow"] = p.RopeParameters.BetaSlow
kv["laguna.rope.scaling.attn_factor"] = lagunaAttentionFactor(ropeType, ropeFactor, p.RopeParameters.AttentionFactor)
kv["laguna.rope.partial_rotary_factor"] = cmp.Or(p.PartialRotaryFactor, float32(1))
swaRopeType := p.SwaRopeParameters.ropeType()
kv["laguna.rope.swa.freq_base"] = cmp.Or(p.SwaRopeParameters.RopeTheta, float32(10000))
kv["laguna.rope.swa.scaling.type"] = cmp.Or(swaRopeType, "linear")
kv["laguna.rope.swa.scaling.factor"] = cmp.Or(p.SwaRopeParameters.Factor, float32(1))
kv["laguna.rope.swa.partial_rotary_factor"] = cmp.Or(p.SwaRopeParameters.PartialRotaryFactor, float32(1))
headDim := p.HeadDim
if headDim == 0 && p.NumAttentionHeads > 0 {
headDim = p.HiddenSize / p.NumAttentionHeads
}
kv["laguna.rope.dimension_count"] = lagunaRopeDim(headDim, cmp.Or(p.PartialRotaryFactor, float32(1)))
kv["laguna.rope.swa.dimension_count"] = lagunaRopeDim(headDim, cmp.Or(p.SwaRopeParameters.PartialRotaryFactor, float32(1)))
return kv
}
func (p *lagunaModel) parseMore(_ iofs.FS) error {
return p.validate()
}
func (p *lagunaModel) validate() error {
if p.NumHiddenLayers == 0 {
return fmt.Errorf("laguna: num_hidden_layers must be set")
}
if p.HiddenSize == 0 {
return fmt.Errorf("laguna: hidden_size must be set")
}
if p.HeadDim == 0 {
return fmt.Errorf("laguna: head_dim must be set")
}
if p.NumKeyValueHeads == 0 {
return fmt.Errorf("laguna: num_key_value_heads must be set")
}
if p.SwaAttentionSinkEnabled {
return fmt.Errorf("laguna: unsupported swa_attention_sink_enabled=true")
}
if !p.Gating.perHead() {
return fmt.Errorf("laguna: unsupported attention gating %q: only gating=\"per-head\" is supported", p.Gating)
}
if p.QKNormType != "rmsnorm" {
return fmt.Errorf("laguna: unsupported qk_norm_type %q: only rmsnorm is supported", p.QKNormType)
}
if !p.MoERouterUseSigmoid {
return fmt.Errorf("laguna: unsupported moe_router_use_sigmoid=false")
}
if p.MoEApplyRouterWeightOnInput {
return fmt.Errorf("laguna: unsupported moe_apply_router_weight_on_input=true")
}
if p.DecoderSparseStep != 0 && p.DecoderSparseStep != 1 {
return fmt.Errorf("laguna: unsupported decoder_sparse_step=%d: only 1 is supported", p.DecoderSparseStep)
}
if len(p.MLPOnlyLayers) != 1 || p.MLPOnlyLayers[0] != 0 {
return fmt.Errorf("laguna: unsupported mlp_only_layers=%v: only [0] is supported", p.MLPOnlyLayers)
}
if p.NumExperts == 0 {
return fmt.Errorf("laguna: num_experts must be set")
}
if p.NumExpertsPerTok == 0 {
return fmt.Errorf("laguna: num_experts_per_tok must be set")
}
if p.MoEIntermediateSize == 0 {
return fmt.Errorf("laguna: moe_intermediate_size must be set")
}
if p.SharedExpertIntermediateSize == 0 {
return fmt.Errorf("laguna: shared_expert_intermediate_size must be set")
}
if len(p.LayerTypes) > 0 && len(p.LayerTypes) != int(p.NumHiddenLayers) {
return fmt.Errorf("laguna: layer_types has %d entries, expected %d", len(p.LayerTypes), p.NumHiddenLayers)
}
for i, layerType := range p.LayerTypes {
if !lagunaLayerIsGlobal(layerType) && !lagunaLayerIsSliding(layerType) {
return fmt.Errorf("laguna: unsupported layer_types[%d]=%q", i, layerType)
}
}
if len(p.NumAttentionHeadsPerLayer) > 0 && len(p.NumAttentionHeadsPerLayer) != int(p.NumHiddenLayers) {
return fmt.Errorf("laguna: num_attention_heads_per_layer has %d entries, expected %d", len(p.NumAttentionHeadsPerLayer), p.NumHiddenLayers)
}
if len(p.NumAttentionHeadsPerLayer) == 0 && p.NumAttentionHeads == 0 {
return fmt.Errorf("laguna: num_attention_heads or num_attention_heads_per_layer must be set")
}
for i, heads := range p.NumAttentionHeadsPerLayer {
if heads == 0 {
return fmt.Errorf("laguna: num_attention_heads_per_layer[%d] must be non-zero", i)
}
}
return nil
}
func (p *lagunaModel) numHeadsForLayer(layer uint32) uint32 {
if len(p.NumAttentionHeadsPerLayer) > int(layer) && p.NumAttentionHeadsPerLayer[layer] > 0 {
return p.NumAttentionHeadsPerLayer[layer]
}
return p.NumAttentionHeads
}
func (p *lagunaModel) layerUsesMoE(layer uint32) bool {
for _, denseLayer := range p.MLPOnlyLayers {
if denseLayer == layer {
return false
}
}
step := cmp.Or(p.DecoderSparseStep, uint32(1))
return p.NumExperts > 0 && (layer+1)%step == 0
}
func (p *lagunaModel) Replacements() []string {
return []string{
"lm_head", "output",
"model.embed_tokens", "token_embd",
"model.norm", "output_norm",
"model.layers", "blk",
"input_layernorm", "attn_norm",
"post_attention_layernorm", "ffn_norm",
"self_attn.q_proj", "attn_q",
"self_attn.k_proj", "attn_k",
"self_attn.v_proj", "attn_v",
"self_attn.o_proj", "attn_output",
"self_attn.g_proj", "attn_g",
"self_attn.q_norm", "attn_q_norm",
"self_attn.k_norm", "attn_k_norm",
"mlp.gate_proj", "ffn_gate",
"mlp.up_proj", "ffn_up",
"mlp.down_proj", "ffn_down",
"mlp.gate.weight", "ffn_gate_inp.weight",
"mlp.experts.e_score_correction_bias", "exp_probs_b.bias",
"mlp.shared_expert.gate_proj", "ffn_gate_shexp",
"mlp.shared_expert.up_proj", "ffn_up_shexp",
"mlp.shared_expert.down_proj", "ffn_down_shexp",
"mlp.experts.*.gate_proj", "ffn_gate_exps",
"mlp.experts.*.up_proj", "ffn_up_exps",
"mlp.experts.*.down_proj", "ffn_down_exps",
}
}
func (p *lagunaModel) Tensors(ts []Tensor) []*ggml.Tensor {
// Current Laguna drops store routed MoE experts as separate per-expert
// tensors. GGUF stores each projection as one stacked tensor. If future
// drops change expert naming or layout, update these patterns with a
// focused conversion test using the new tensor names.
merges := make([]merge, 0, p.NumHiddenLayers*3)
for i := range p.NumHiddenLayers {
merges = append(merges,
merge{
fmt.Sprintf("blk.%d.mlp.experts.*.gate_proj.weight", i),
fmt.Sprintf("blk.%d.ffn_gate_exps.weight", i),
},
merge{
fmt.Sprintf("blk.%d.mlp.experts.*.up_proj.weight", i),
fmt.Sprintf("blk.%d.ffn_up_exps.weight", i),
},
merge{
fmt.Sprintf("blk.%d.mlp.experts.*.down_proj.weight", i),
fmt.Sprintf("blk.%d.ffn_down_exps.weight", i),
},
)
}
out, rest := mergeTensors(ts, merges...)
for _, t := range rest {
out = append(out, &ggml.Tensor{
Name: t.Name(),
Kind: t.Kind(),
Shape: t.Shape(),
WriterTo: t,
})
}
return out
}
func (p *lagunaModel) specialTokenTypes() []string {
return []string{"bos", "eos", "pad", "unk"}
}
func lagunaLayerIsSliding(layerType string) bool {
return strings.EqualFold(layerType, "sliding_attention")
}
func lagunaLayerIsGlobal(layerType string) bool {
return strings.EqualFold(layerType, "full_attention") || strings.EqualFold(layerType, "global_attention")
}
func lagunaLeadingDensePrefix(layers []uint32) (uint32, bool) {
for i, v := range layers {
if v != uint32(i) {
return 0, false
}
}
return uint32(len(layers)), true
}
func lagunaDenseLayers(mlpOnlyLayers []uint32, mlpLayerTypes []string) ([]uint32, error) {
if len(mlpOnlyLayers) > 0 {
return mlpOnlyLayers, nil
}
if len(mlpLayerTypes) == 0 {
return nil, nil
}
denseLayers := make([]uint32, 0, len(mlpLayerTypes))
for i, layerType := range mlpLayerTypes {
switch {
case strings.EqualFold(layerType, "dense"):
denseLayers = append(denseLayers, uint32(i))
case strings.EqualFold(layerType, "sparse"):
default:
return nil, fmt.Errorf("laguna: unsupported mlp_layer_types[%d]=%q", i, layerType)
}
}
return denseLayers, nil
}
func lagunaMoeGatingFunc(useSigmoid bool) uint32 {
if useSigmoid {
return lagunaGatingFuncSigmoid
}
return lagunaGatingFuncSoftmax
}
func lagunaAttentionFactor(ropeType string, scaleFactor, attentionFactor float32) float32 {
if attentionFactor != 0 {
return attentionFactor
}
if strings.EqualFold(ropeType, "yarn") && scaleFactor > 1 {
return float32(0.1*math.Log(float64(scaleFactor)) + 1)
}
return 1
}
func lagunaRopeDim(headDim uint32, partialRotaryFactor float32) uint32 {
if headDim == 0 {
return 0
}
dim := uint32(float32(headDim) * partialRotaryFactor)
if dim == 0 || dim > headDim {
dim = headDim
}
if dim%2 != 0 {
dim--
}
if dim == 0 {
return headDim
}
return dim
}
var (
_ ModelConverter = (*lagunaModel)(nil)
_ moreParser = (*lagunaModel)(nil)
)
+450
View File
@@ -0,0 +1,450 @@
package convert
import (
"encoding/json"
"fmt"
"io"
"math"
"strings"
"testing"
"github.com/google/go-cmp/cmp"
"github.com/ollama/ollama/fs/ggml"
)
type lagunaTestTensor struct {
tensorBase
}
func newLagunaTestTensor(name string, shape ...uint64) Tensor {
return &lagunaTestTensor{tensorBase: tensorBase{name: name, shape: shape}}
}
func (t *lagunaTestTensor) WriteTo(io.Writer) (int64, error) {
return 0, nil
}
func (t *lagunaTestTensor) Clone() Tensor {
return &lagunaTestTensor{tensorBase: tensorBase{
name: t.name,
shape: append([]uint64(nil), t.shape...),
}}
}
func TestLagunaReplacements(t *testing.T) {
p := lagunaModel{}
r := strings.NewReplacer(p.Replacements()...)
tests := []struct {
name string
in string
want string
}{
{"embed", "model.embed_tokens.weight", "token_embd.weight"},
{"final_norm", "model.norm.weight", "output_norm.weight"},
{"lm_head", "lm_head.weight", "output.weight"},
{"block prefix", "model.layers.7.input_layernorm.weight", "blk.7.attn_norm.weight"},
{"q", "model.layers.3.self_attn.q_proj.weight", "blk.3.attn_q.weight"},
{"k", "model.layers.3.self_attn.k_proj.weight", "blk.3.attn_k.weight"},
{"v", "model.layers.3.self_attn.v_proj.weight", "blk.3.attn_v.weight"},
{"o", "model.layers.3.self_attn.o_proj.weight", "blk.3.attn_output.weight"},
{"g", "model.layers.3.self_attn.g_proj.weight", "blk.3.attn_g.weight"},
{"q_norm", "model.layers.3.self_attn.q_norm.weight", "blk.3.attn_q_norm.weight"},
{"k_norm", "model.layers.3.self_attn.k_norm.weight", "blk.3.attn_k_norm.weight"},
{"post_attn_norm", "model.layers.3.post_attention_layernorm.weight", "blk.3.ffn_norm.weight"},
{"dense gate", "model.layers.0.mlp.gate_proj.weight", "blk.0.ffn_gate.weight"},
{"dense up", "model.layers.0.mlp.up_proj.weight", "blk.0.ffn_up.weight"},
{"dense down", "model.layers.0.mlp.down_proj.weight", "blk.0.ffn_down.weight"},
{"shexp gate", "model.layers.5.mlp.shared_expert.gate_proj.weight", "blk.5.ffn_gate_shexp.weight"},
{"shexp up", "model.layers.5.mlp.shared_expert.up_proj.weight", "blk.5.ffn_up_shexp.weight"},
{"shexp down", "model.layers.5.mlp.shared_expert.down_proj.weight", "blk.5.ffn_down_shexp.weight"},
{"router", "model.layers.5.mlp.gate.weight", "blk.5.ffn_gate_inp.weight"},
{"score bias", "model.layers.5.mlp.experts.e_score_correction_bias", "blk.5.exp_probs_b.bias"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
if got := r.Replace(tc.in); got != tc.want {
t.Errorf("Replace(%q) = %q, want %q", tc.in, got, tc.want)
}
})
}
}
func TestLagunaValidateRejectsUnsupportedVariants(t *testing.T) {
base := validLagunaTestModel()
tests := []struct {
name string
edit func(*lagunaModel)
want string
}{
{
name: "per-element gating",
edit: func(m *lagunaModel) {
m.Gating = "per-element"
},
want: "unsupported attention gating",
},
{
name: "attention sinks",
edit: func(m *lagunaModel) {
m.SwaAttentionSinkEnabled = true
},
want: "swa_attention_sink_enabled=true",
},
{
name: "qk norm disabled",
edit: func(m *lagunaModel) {
m.QKNormType = "none"
},
want: "unsupported qk_norm_type",
},
{
name: "softmax moe",
edit: func(m *lagunaModel) {
m.MoERouterUseSigmoid = false
},
want: "moe_router_use_sigmoid=false",
},
{
name: "router weight on input",
edit: func(m *lagunaModel) {
m.MoEApplyRouterWeightOnInput = true
},
want: "moe_apply_router_weight_on_input=true",
},
{
name: "unknown layer type",
edit: func(m *lagunaModel) {
m.LayerTypes[1] = "local_attention"
},
want: "unsupported layer_types[1]",
},
{
name: "nonstandard dense layout",
edit: func(m *lagunaModel) {
m.MLPOnlyLayers = []uint32{0, 3}
},
want: "unsupported mlp_only_layers",
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
m := base
m.LayerTypes = append([]string(nil), base.LayerTypes...)
m.NumAttentionHeadsPerLayer = append([]uint32(nil), base.NumAttentionHeadsPerLayer...)
m.MLPOnlyLayers = append([]uint32(nil), base.MLPOnlyLayers...)
tc.edit(&m)
err := m.validate()
if err == nil || !strings.Contains(err.Error(), tc.want) {
t.Fatalf("validate() error = %v, want substring %q", err, tc.want)
}
})
}
}
func TestLagunaGAConfigNormalizesBoolGatingAndNestedRope(t *testing.T) {
var m lagunaModel
if err := json.Unmarshal([]byte(`{
"architectures": ["LagunaForCausalLM"],
"num_hidden_layers": 1,
"hidden_size": 8,
"num_attention_heads": 2,
"num_key_value_heads": 1,
"head_dim": 4,
"gating": true,
"num_experts": 2,
"num_experts_per_tok": 1,
"moe_intermediate_size": 4,
"shared_expert_intermediate_size": 4,
"decoder_sparse_step": 1,
"mlp_layer_types": ["dense"],
"rope_parameters": {
"full_attention": {
"rope_theta": 500000,
"rope_type": "yarn",
"factor": 32,
"original_max_position_embeddings": 4096,
"beta_fast": 64,
"beta_slow": 1,
"attention_factor": 1,
"partial_rotary_factor": 0.5
},
"sliding_attention": {
"rope_theta": 10000,
"rope_type": "default",
"partial_rotary_factor": 1
}
}
}`), &m); err != nil {
t.Fatalf("json.Unmarshal() error = %v", err)
}
if err := m.validate(); err != nil {
t.Fatalf("validate() error = %v", err)
}
if m.Gating != "true" {
t.Fatalf("Gating = %q, want raw true marker", m.Gating)
}
if !m.Gating.perHead() {
t.Fatal("expected bool gating to normalize as per-head support")
}
if m.QKNormType != "rmsnorm" {
t.Fatalf("QKNormType = %q, want rmsnorm default", m.QKNormType)
}
if !m.MoERouterUseSigmoid {
t.Fatal("MoERouterUseSigmoid should default true")
}
if !m.NormTopKProb {
t.Fatal("NormTopKProb should default true")
}
if diff := cmp.Diff(m.MLPOnlyLayers, []uint32{0}); diff != "" {
t.Fatalf("MLPOnlyLayers mismatch (-got +want):\n%s", diff)
}
if m.RopeParameters.RopeTheta != 500000 || m.RopeParameters.PartialRotaryFactor != 0.5 {
t.Fatalf("full rope = %#v, want theta=500000 partial=0.5", m.RopeParameters)
}
if m.SwaRopeParameters.RopeTheta != 10000 || m.SwaRopeParameters.PartialRotaryFactor != 1 {
t.Fatalf("swa rope = %#v, want theta=10000 partial=1", m.SwaRopeParameters)
}
}
func validLagunaTestModel() lagunaModel {
return lagunaModel{
ModelParameters: ModelParameters{
VocabSize: 32,
},
NumHiddenLayers: 2,
HiddenSize: 8,
IntermediateSize: 16,
NumAttentionHeads: 2,
NumKeyValueHeads: 1,
HeadDim: 4,
RMSNormEPS: 1e-6,
MaxPositionEmbeddings: 4096,
SlidingWindow: 512,
Gating: "per-head",
QKNormType: "rmsnorm",
LayerTypes: []string{"global_attention", "sliding_attention"},
NumAttentionHeadsPerLayer: []uint32{2, 2},
NumExperts: 2,
NumExpertsPerTok: 1,
MoEIntermediateSize: 4,
SharedExpertIntermediateSize: 4,
NormTopKProb: true,
MoeRoutedScalingFactor: 2.5,
MoERouterUseSigmoid: true,
DecoderSparseStep: 1,
MLPOnlyLayers: []uint32{0},
}
}
func validLagunaTestTensors(m lagunaModel) []Tensor {
ts := []Tensor{
newLagunaTestTensor("token_embd.weight", uint64(m.VocabSize), uint64(m.HiddenSize)),
newLagunaTestTensor("output_norm.weight", uint64(m.HiddenSize)),
}
for layer := range m.NumHiddenLayers {
prefix := fmt.Sprintf("blk.%d", layer)
heads := uint64(m.numHeadsForLayer(layer))
attnWidth := heads * uint64(m.HeadDim)
kvWidth := uint64(m.NumKeyValueHeads * m.HeadDim)
ts = append(ts,
newLagunaTestTensor(prefix+".attn_norm.weight", uint64(m.HiddenSize)),
newLagunaTestTensor(prefix+".ffn_norm.weight", uint64(m.HiddenSize)),
newLagunaTestTensor(prefix+".attn_q.weight", attnWidth, uint64(m.HiddenSize)),
newLagunaTestTensor(prefix+".attn_k.weight", kvWidth, uint64(m.HiddenSize)),
newLagunaTestTensor(prefix+".attn_v.weight", kvWidth, uint64(m.HiddenSize)),
newLagunaTestTensor(prefix+".attn_output.weight", uint64(m.HiddenSize), attnWidth),
newLagunaTestTensor(prefix+".attn_g.weight", heads, uint64(m.HiddenSize)),
newLagunaTestTensor(prefix+".attn_q_norm.weight", uint64(m.HeadDim)),
newLagunaTestTensor(prefix+".attn_k_norm.weight", uint64(m.HeadDim)),
)
if m.layerUsesMoE(layer) {
ts = append(ts,
newLagunaTestTensor(prefix+".ffn_gate_inp.weight", uint64(m.NumExperts), uint64(m.HiddenSize)),
newLagunaTestTensor(prefix+".exp_probs_b.bias", uint64(m.NumExperts)),
newLagunaTestTensor(prefix+".ffn_gate_shexp.weight", uint64(m.SharedExpertIntermediateSize), uint64(m.HiddenSize)),
newLagunaTestTensor(prefix+".ffn_up_shexp.weight", uint64(m.SharedExpertIntermediateSize), uint64(m.HiddenSize)),
newLagunaTestTensor(prefix+".ffn_down_shexp.weight", uint64(m.HiddenSize), uint64(m.SharedExpertIntermediateSize)),
)
for expert := range m.NumExperts {
ts = append(ts,
newLagunaTestTensor(fmt.Sprintf("%s.mlp.experts.%d.gate_proj.weight", prefix, expert), uint64(m.MoEIntermediateSize), uint64(m.HiddenSize)),
newLagunaTestTensor(fmt.Sprintf("%s.mlp.experts.%d.up_proj.weight", prefix, expert), uint64(m.MoEIntermediateSize), uint64(m.HiddenSize)),
newLagunaTestTensor(fmt.Sprintf("%s.mlp.experts.%d.down_proj.weight", prefix, expert), uint64(m.HiddenSize), uint64(m.MoEIntermediateSize)),
)
}
} else {
ts = append(ts,
newLagunaTestTensor(prefix+".ffn_gate.weight", uint64(m.IntermediateSize), uint64(m.HiddenSize)),
newLagunaTestTensor(prefix+".ffn_up.weight", uint64(m.IntermediateSize), uint64(m.HiddenSize)),
newLagunaTestTensor(prefix+".ffn_down.weight", uint64(m.HiddenSize), uint64(m.IntermediateSize)),
)
}
}
return ts
}
func TestLagunaTensorsMergeRoutedExperts(t *testing.T) {
m := validLagunaTestModel()
out := m.Tensors(validLagunaTestTensors(m))
tensors := make(map[string]*ggml.Tensor, len(out))
for _, t := range out {
tensors[t.Name] = t
}
tests := map[string][]uint64{
"blk.1.ffn_gate_exps.weight": {uint64(m.NumExperts), uint64(m.MoEIntermediateSize), uint64(m.HiddenSize)},
"blk.1.ffn_up_exps.weight": {uint64(m.NumExperts), uint64(m.MoEIntermediateSize), uint64(m.HiddenSize)},
"blk.1.ffn_down_exps.weight": {uint64(m.NumExperts), uint64(m.HiddenSize), uint64(m.MoEIntermediateSize)},
}
for name, wantShape := range tests {
tensor, ok := tensors[name]
if !ok {
t.Fatalf("missing merged tensor %q", name)
}
if diff := cmp.Diff(wantShape, tensor.Shape); diff != "" {
t.Fatalf("%s shape mismatch (-want +got):\n%s", name, diff)
}
}
for expert := range m.NumExperts {
name := fmt.Sprintf("blk.1.mlp.experts.%d.gate_proj.weight", expert)
if _, ok := tensors[name]; ok {
t.Fatalf("unexpected unmerged expert tensor %q", name)
}
}
}
func TestLagunaKVShape(t *testing.T) {
m := lagunaModel{
NumHiddenLayers: 4,
HiddenSize: 128,
IntermediateSize: 256,
NumAttentionHeads: 8,
NumKeyValueHeads: 4,
HeadDim: 16,
RMSNormEPS: 1e-6,
MaxPositionEmbeddings: 4096,
SlidingWindow: 512,
PartialRotaryFactor: 0.5,
Gating: "per-head",
QKNormType: "rmsnorm",
LayerTypes: []string{"full_attention", "sliding_attention", "sliding_attention", "sliding_attention"},
NumAttentionHeadsPerLayer: []uint32{8, 16, 16, 16},
NumExperts: 32,
NumExpertsPerTok: 4,
MoEIntermediateSize: 64,
SharedExpertIntermediateSize: 64,
NormTopKProb: true,
MoeRoutedScalingFactor: 2.5,
MoERouterUseSigmoid: true,
DecoderSparseStep: 1,
MLPOnlyLayers: []uint32{0},
}
m.RopeParameters.RopeTheta = 500000
m.RopeParameters.RopeType = "yarn"
m.RopeParameters.Factor = 32
m.RopeParameters.OriginalMaxPositionEmbeddings = 4096
m.RopeParameters.BetaFast = 64
m.RopeParameters.BetaSlow = 1
m.SwaRopeParameters.RopeTheta = 10000
m.SwaRopeParameters.RopeType = "linear"
m.SwaRopeParameters.Factor = 1
m.SwaRopeParameters.PartialRotaryFactor = 1
kv := m.KV(&Tokenizer{Vocabulary: &Vocabulary{}, Template: "{% include 'chat_template.jinja' %}"})
required := []string{
"general.architecture",
"tokenizer.ggml.pre",
"laguna.block_count",
"laguna.context_length",
"laguna.embedding_length",
"laguna.feed_forward_length",
"laguna.attention.head_count",
"laguna.attention.head_count_kv",
"laguna.attention.key_length",
"laguna.attention.value_length",
"laguna.attention.layer_norm_rms_epsilon",
"laguna.attention.sliding_window",
"laguna.attention.layer_types",
"laguna.attention.sliding_window_pattern",
"laguna.attention.gating_type",
"laguna.attention.qk_norm",
"laguna.expert_count",
"laguna.expert_used_count",
"laguna.expert_feed_forward_length",
"laguna.expert_shared_feed_forward_length",
"laguna.expert_shared_count",
"laguna.expert_weights_norm",
"laguna.expert_weights_scale",
"laguna.expert_gating_func",
"laguna.leading_dense_block_count",
"laguna.dense_layers",
"laguna.rope.freq_base",
"laguna.rope.scaling.type",
"laguna.rope.scaling.factor",
"laguna.rope.partial_rotary_factor",
"laguna.rope.swa.freq_base",
"laguna.rope.swa.scaling.type",
"laguna.rope.dimension_count",
"laguna.rope.swa.dimension_count",
}
for _, k := range required {
if _, ok := kv[k]; !ok {
t.Errorf("missing required KV: %s", k)
}
}
if got := kv["general.architecture"]; got != "laguna" {
t.Errorf("architecture = %v, want laguna", got)
}
if got := kv["tokenizer.ggml.add_bos_token"]; got != false {
t.Errorf("tokenizer.ggml.add_bos_token = %v, want false", got)
}
if _, ok := kv["tokenizer.chat_template"]; ok {
t.Fatal("tokenizer.chat_template should be omitted for Laguna")
}
if got := kv["laguna.expert_gating_func"]; got != lagunaGatingFuncSigmoid {
t.Errorf("expert_gating_func = %v, want sigmoid(%d)", got, lagunaGatingFuncSigmoid)
}
if got := kv["laguna.leading_dense_block_count"]; got != uint32(1) {
t.Errorf("leading_dense_block_count = %v, want 1", got)
}
if got := kv["laguna.rope.dimension_count"]; got != uint32(8) {
t.Errorf("rope.dimension_count = %v, want 8", got)
}
if got := kv["laguna.rope.swa.dimension_count"]; got != uint32(16) {
t.Errorf("rope.swa.dimension_count = %v, want 16", got)
}
if got, ok := kv["laguna.attention.layer_types"].([]uint32); !ok || len(got) != 4 || got[0] != 0 || got[1] != 1 || got[2] != 1 || got[3] != 1 {
t.Fatalf("layer_types = %#v, want [0 1 1 1]", kv["laguna.attention.layer_types"])
}
if got, ok := kv["laguna.attention.sliding_window_pattern"].([]bool); !ok || len(got) != 4 || got[0] || !got[1] || !got[2] || !got[3] {
t.Fatalf("sliding_window_pattern = %#v, want [false true true true]", kv["laguna.attention.sliding_window_pattern"])
}
}
func TestLagunaKVYarnAttentionFactorFallback(t *testing.T) {
m := validLagunaTestModel()
m.RopeParameters.RopeType = "yarn"
m.RopeParameters.Factor = 32
kv := m.KV(&Tokenizer{Vocabulary: &Vocabulary{}})
got, ok := kv["laguna.rope.scaling.attn_factor"].(float32)
if !ok {
t.Fatalf("attn_factor type = %T, want float32", kv["laguna.rope.scaling.attn_factor"])
}
want := float32(0.1*math.Log(32) + 1)
if diff := math.Abs(float64(got - want)); diff > 1e-6 {
t.Fatalf("attn_factor = %v, want %v", got, want)
}
}
+409
View File
@@ -3,6 +3,7 @@ package convert
import (
"cmp"
"encoding/json"
"errors"
"fmt"
"io/fs"
"math"
@@ -69,7 +70,415 @@ type nemotronHModel struct {
ExpertGroupUsedCount uint32 `json:"topk_group"`
}
type nemotronHNanoVLModel struct {
ModelParameters
MaxSequenceLength uint32 `json:"max_sequence_length"`
ForceImageSize uint32 `json:"force_image_size"`
DownsampleRatio float32 `json:"downsample_ratio"`
PatchSize uint32 `json:"patch_size"`
UseThumbnail *bool `json:"use_thumbnail"`
ImgContextTokenID uint32 `json:"img_context_token_id"`
ImgContextToken string `json:"img_context_token"`
ImgStartToken string `json:"img_start_token"`
ImgEndToken string `json:"img_end_token"`
VitHiddenSize uint32 `json:"vit_hidden_size"`
ProjectorHidden uint32 `json:"projector_hidden_size"`
SoundContextTokenID uint32 `json:"sound_context_token_id"`
SoundContextToken string `json:"sound_context_token"`
NormMean []float32 `json:"norm_mean"`
NormStd []float32 `json:"norm_std"`
VisionConfig radioConfig `json:"vision_config"`
SoundConfig soundConfig `json:"sound_config"`
LLMConfig nemotronHModel `json:"llm_config"`
Preprocessor struct {
ImageSize uint32 `json:"image_size"`
PatchSize uint32 `json:"patch_size"`
DownsampleRatio float32 `json:"downsample_ratio"`
MaxNumTiles uint32 `json:"max_num_tiles"`
UseThumbnail *bool `json:"use_thumbnail"`
NormMean []float32 `json:"norm_mean"`
NormStd []float32 `json:"norm_std"`
}
}
type soundConfig struct {
ModelType string `json:"model_type"`
HiddenSize uint32 `json:"hidden_size"`
NumAttentionHeads uint32 `json:"num_attention_heads"`
NumHiddenLayers uint32 `json:"num_hidden_layers"`
IntermediateSize uint32 `json:"intermediate_size"`
ConvKernelSize uint32 `json:"conv_kernel_size"`
SubsamplingConvChannels uint32 `json:"subsampling_conv_channels"`
SubsamplingConvKernelSize uint32 `json:"subsampling_conv_kernel_size"`
SubsamplingConvStride uint32 `json:"subsampling_conv_stride"`
SubsamplingFactor uint32 `json:"subsampling_factor"`
NumMelBins uint32 `json:"num_mel_bins"`
ProjectionHiddenSize uint32 `json:"projection_hidden_size"`
SamplingRate uint32 `json:"sampling_rate"`
ScaleInput bool `json:"scale_input"`
}
type radioConfig struct {
Version string `json:"version"`
PatchSize uint32 `json:"patch_size"`
MaxResolution uint32 `json:"max_resolution"`
MinNumPatches uint32 `json:"min_num_patches"`
MaxNumPatches uint32 `json:"max_num_patches"`
SeparateVideoEmbedder bool `json:"separate_video_embedder"`
Args struct {
MinNumPatches uint32 `json:"min_num_patches"`
MaxNumPatches uint32 `json:"max_num_patches"`
} `json:"args"`
}
var _ ModelConverter = (*nemotronHModel)(nil)
var _ ModelConverter = (*nemotronHNanoVLModel)(nil)
func (n *nemotronHNanoVLModel) parseMore(fsys fs.FS) error {
if n.MaxSequenceLength > 0 {
n.LLMConfig.MaxPositionEmbeddings = n.MaxSequenceLength
}
if err := n.LLMConfig.parseMore(fsys); err != nil {
return err
}
if bts, err := fs.ReadFile(fsys, "preprocessor_config.json"); err == nil {
if err := json.Unmarshal(bts, &n.Preprocessor); err != nil {
return fmt.Errorf("nemotron_h_omni: parse preprocessor_config.json: %w", err)
}
} else if !errors.Is(err, fs.ErrNotExist) {
return err
}
if version := strings.TrimSpace(n.VisionConfig.Version); version != "" && version != "radio_v2.5-h" {
return fmt.Errorf("nemotron_h_omni: unsupported RADIO version %q", version)
}
if patchSize := n.visionPatchSize(); patchSize != 16 {
return fmt.Errorf("nemotron_h_omni: unsupported vision patch_size=%d", patchSize)
}
if scale := n.visionProjectorScaleFactor(); scale != 2 {
return fmt.Errorf("nemotron_h_omni: unsupported vision projector scale factor=%d", scale)
}
if n.SoundConfig.NumHiddenLayers > 0 {
if modelType := strings.TrimSpace(n.SoundConfig.ModelType); modelType != "" && modelType != "parakeet" {
return fmt.Errorf("nemotron_h_omni: unsupported sound model_type %q", modelType)
}
if n.soundHiddenSize() == 0 {
return fmt.Errorf("nemotron_h_omni: sound hidden_size must be set")
}
if n.soundAttentionHeads() == 0 {
return fmt.Errorf("nemotron_h_omni: sound num_attention_heads must be set")
}
if n.soundSubsamplingFactor() != 8 {
return fmt.Errorf("nemotron_h_omni: unsupported sound subsampling_factor=%d", n.soundSubsamplingFactor())
}
if n.soundMelBins() != 128 {
return fmt.Errorf("nemotron_h_omni: unsupported sound num_mel_bins=%d", n.soundMelBins())
}
}
return nil
}
func (n *nemotronHNanoVLModel) KV(t *Tokenizer) KV {
kv := n.LLMConfig.KV(t)
kv["general.architecture"] = "nemotron_h_omni"
kv["vision.block_count"] = n.visionBlockCount()
kv["vision.embedding_length"] = n.visionEmbeddingLength()
kv["vision.feed_forward_length"] = n.visionFeedForwardLength()
kv["vision.attention.head_count"] = n.visionAttentionHeads()
kv["vision.attention.layer_norm_epsilon"] = float32(1e-6)
kv["vision.patch_size"] = n.visionPatchSize()
kv["vision.image_size"] = n.visionImageSize()
kv["vision.max_tiles"] = n.visionMaxTiles()
kv["vision.use_thumbnail"] = n.visionUseThumbnail()
if minPatches := n.visionMinNumPatches(); minPatches > 0 {
kv["vision.min_num_patches"] = minPatches
}
if maxPatches := n.visionMaxNumPatches(); maxPatches > 0 {
kv["vision.max_num_patches"] = maxPatches
}
kv["vision.num_channels"] = uint32(3)
kv["vision.image_mean"] = slices.Clone(defaultFloat32Slice(n.visionMean(), imageNetStandardMean))
kv["vision.image_std"] = slices.Clone(defaultFloat32Slice(n.visionStd(), imageNetStandardSTD))
kv["vision.projector.scale_factor"] = n.visionProjectorScaleFactor()
setTokenID := func(key string, explicit uint32, token string) {
if explicit > 0 {
kv[key] = explicit
return
}
if t == nil || t.Vocabulary == nil {
return
}
for i, v := range t.Vocabulary.Tokens {
if v == token {
kv[key] = uint32(i)
return
}
}
}
setTokenID("vision.image_token_id", n.ImgContextTokenID, cmp.Or(n.ImgContextToken, "<image>"))
setTokenID("vision.image_start_token_id", 0, cmp.Or(n.ImgStartToken, "<img>"))
setTokenID("vision.image_end_token_id", 0, cmp.Or(n.ImgEndToken, "</img>"))
if n.SoundConfig.NumHiddenLayers > 0 {
kv["audio.block_count"] = n.SoundConfig.NumHiddenLayers
kv["audio.embedding_length"] = n.soundHiddenSize()
kv["audio.feed_forward_length"] = n.soundFeedForwardLength()
kv["audio.attention.head_count"] = n.soundAttentionHeads()
kv["audio.attention.layer_norm_epsilon"] = float32(1e-5)
kv["audio.conv_kernel_size"] = n.soundConvKernelSize()
kv["audio.num_mel_bins"] = n.soundMelBins()
kv["audio.sample_rate"] = n.soundSampleRate()
kv["audio.subsampling_factor"] = n.soundSubsamplingFactor()
kv["audio.subsampling_conv_channels"] = n.soundSubsamplingConvChannels()
kv["audio.subsampling_conv_kernel_size"] = n.soundSubsamplingConvKernelSize()
kv["audio.subsampling_conv_stride"] = n.soundSubsamplingConvStride()
kv["audio.projection_hidden_size"] = n.soundProjectionHiddenSize()
kv["audio.scale_input"] = n.SoundConfig.ScaleInput
setTokenID("audio.sound_token_id", n.SoundContextTokenID, cmp.Or(n.SoundContextToken, "<so_embedding>"))
}
return kv
}
func (n *nemotronHNanoVLModel) Tensors(ts []Tensor) []*ggml.Tensor {
var textTensors []Tensor
var out []*ggml.Tensor
for _, t := range ts {
switch {
case isNemotronHNanoVLOmittedTensor(t.Name()):
continue
case strings.Contains(t.Name(), ".attn_qkv"):
out = append(out, slices.Collect(splitDim(t, 0,
split{Replacer: strings.NewReplacer("attn_qkv", "attn_q")},
split{Replacer: strings.NewReplacer("attn_qkv", "attn_k")},
split{Replacer: strings.NewReplacer("attn_qkv", "attn_v")},
))...)
case t.Name() == "v.position_embd":
shape := t.Shape()
if len(shape) == 3 && shape[0] == 1 {
shape = shape[1:]
}
out = append(out, &ggml.Tensor{
Name: t.Name(),
Kind: t.Kind(),
Shape: shape,
WriterTo: t,
})
case strings.HasPrefix(t.Name(), "a.") || strings.HasPrefix(t.Name(), "v.") || strings.HasPrefix(t.Name(), "mm."):
name := t.Name()
shape := slices.Clone(t.Shape())
if strings.HasPrefix(name, "a.blk.") && strings.Contains(name, ".conv_dw.") && strings.HasSuffix(name, ".weight") && len(shape) == 3 {
t.SetRepacker(squeezeMiddleDim)
shape = []uint64{shape[0], shape[2]}
}
if strings.HasPrefix(name, "a.blk.") && (strings.Contains(name, ".conv_pw1.") || strings.Contains(name, ".conv_pw2.")) && strings.HasSuffix(name, ".weight") && len(shape) == 3 && shape[2] == 1 {
t.SetRepacker(squeezeLastDim)
shape = shape[:2]
}
out = append(out, &ggml.Tensor{
Name: name,
Kind: t.Kind(),
Shape: shape,
WriterTo: t,
})
default:
textTensors = append(textTensors, t)
}
}
return append(n.LLMConfig.Tensors(textTensors), out...)
}
func (n *nemotronHNanoVLModel) Replacements() []string {
return append([]string{
"language_model.", "",
"vision_model.radio_model.model.patch_generator.embedder", "v.patch_embd",
"vision_model.radio_model.model.patch_generator.pos_embed", "v.position_embd",
"vision_model.radio_model.model.patch_generator.cls_token.token", "v.cls_embd",
"vision_model.radio_model.model.blocks", "v.blk",
"attn.qkv", "attn_qkv",
"attn.proj", "attn_out",
"mlp.fc1", "ffn_up",
"mlp.fc2", "ffn_down",
"norm1", "ln1",
"norm2", "ln2",
"mlp1.0", "mm.norm",
"mlp1.1", "mm.1",
"mlp1.3", "mm.2",
"sound_encoder.encoder.feature_extractor.featurizer.fb", "a.feature_extractor.fb",
"sound_encoder.encoder.feature_extractor.featurizer.window", "a.feature_extractor.window",
"sound_encoder.encoder.subsampling.layers.0", "a.subsampling.conv0",
"sound_encoder.encoder.subsampling.layers.2", "a.subsampling.dw1",
"sound_encoder.encoder.subsampling.layers.3", "a.subsampling.pw1",
"sound_encoder.encoder.subsampling.layers.5", "a.subsampling.dw2",
"sound_encoder.encoder.subsampling.layers.6", "a.subsampling.pw2",
"sound_encoder.encoder.subsampling.linear", "a.subsampling.linear",
"sound_encoder.encoder.layers", "a.blk",
"feed_forward1.linear1", "ffn1_up",
"feed_forward1.linear2", "ffn1_down",
"feed_forward2.linear1", "ffn2_up",
"feed_forward2.linear2", "ffn2_down",
"norm_feed_forward1", "ffn1_norm",
"norm_feed_forward2", "ffn2_norm",
"norm_self_att", "attn_norm",
"norm_conv", "conv_norm",
"norm_out", "out_norm",
"self_attn.q_proj", "attn_q",
"self_attn.k_proj", "attn_k",
"self_attn.v_proj", "attn_v",
"self_attn.o_proj", "attn_out",
"self_attn.relative_k_proj", "attn_rel_k",
"self_attn.bias_u", "attn_bias_u",
"self_attn.bias_v", "attn_bias_v",
"conv.pointwise_conv1", "conv_pw1",
"conv.pointwise_conv2", "conv_pw2",
"conv.depthwise_conv", "conv_dw",
"conv.norm", "conv_bn",
"sound_projection.norm", "mm.a.norm",
"sound_projection.linear1", "mm.a.1",
"sound_projection.linear2", "mm.a.2",
}, n.LLMConfig.Replacements()...)
}
func (n *nemotronHNanoVLModel) specialTokenTypes() []string {
return n.LLMConfig.specialTokenTypes()
}
func isNemotronHNanoVLOmittedTensor(name string) bool {
return strings.HasSuffix(name, ".conv_bn.num_batches_tracked") ||
strings.HasPrefix(name, "vision_model.radio_model.input_conditioner.") ||
strings.HasPrefix(name, "vision_model.radio_model.model.patch_generator.video_embedder")
}
func squeezeLastDim(_ string, data []float32, _ []uint64) ([]float32, error) {
return data, nil
}
func (n *nemotronHNanoVLModel) visionImageSize() uint32 {
return cmp.Or(n.ForceImageSize, n.Preprocessor.ImageSize, uint32(512))
}
func (n *nemotronHNanoVLModel) visionPatchSize() uint32 {
return cmp.Or(n.PatchSize, n.Preprocessor.PatchSize, n.VisionConfig.PatchSize, uint32(16))
}
func (n *nemotronHNanoVLModel) visionProjectorScaleFactor() uint32 {
ratio := cmp.Or(n.DownsampleRatio, n.Preprocessor.DownsampleRatio, float32(0.5))
if ratio <= 0 {
return 2
}
return max(uint32(1), uint32(math.Round(1.0/float64(ratio))))
}
func (n *nemotronHNanoVLModel) visionBlockCount() uint32 {
return 32
}
func (n *nemotronHNanoVLModel) visionEmbeddingLength() uint32 {
return cmp.Or(n.VitHiddenSize, uint32(1280))
}
func (n *nemotronHNanoVLModel) visionAttentionHeads() uint32 {
return 16
}
func (n *nemotronHNanoVLModel) visionFeedForwardLength() uint32 {
return 4 * n.visionEmbeddingLength()
}
func (n *nemotronHNanoVLModel) visionMaxTiles() uint32 {
return cmp.Or(n.Preprocessor.MaxNumTiles, uint32(12))
}
func (n *nemotronHNanoVLModel) visionMinNumPatches() uint32 {
return cmp.Or(n.VisionConfig.MinNumPatches, n.VisionConfig.Args.MinNumPatches)
}
func (n *nemotronHNanoVLModel) visionMaxNumPatches() uint32 {
return cmp.Or(n.VisionConfig.MaxNumPatches, n.VisionConfig.Args.MaxNumPatches)
}
func (n *nemotronHNanoVLModel) visionUseThumbnail() bool {
for _, v := range []*bool{n.UseThumbnail, n.Preprocessor.UseThumbnail} {
if v != nil {
return *v
}
}
return true
}
func (n *nemotronHNanoVLModel) visionMean() []float32 {
if len(n.NormMean) > 0 {
return n.NormMean
}
return n.Preprocessor.NormMean
}
func (n *nemotronHNanoVLModel) visionStd() []float32 {
if len(n.NormStd) > 0 {
return n.NormStd
}
return n.Preprocessor.NormStd
}
func (n *nemotronHNanoVLModel) soundHiddenSize() uint32 {
return cmp.Or(n.SoundConfig.HiddenSize, uint32(1024))
}
func (n *nemotronHNanoVLModel) soundAttentionHeads() uint32 {
return cmp.Or(n.SoundConfig.NumAttentionHeads, uint32(8))
}
func (n *nemotronHNanoVLModel) soundFeedForwardLength() uint32 {
return cmp.Or(n.SoundConfig.IntermediateSize, 4*n.soundHiddenSize())
}
func (n *nemotronHNanoVLModel) soundConvKernelSize() uint32 {
return cmp.Or(n.SoundConfig.ConvKernelSize, uint32(9))
}
func (n *nemotronHNanoVLModel) soundMelBins() uint32 {
return cmp.Or(n.SoundConfig.NumMelBins, uint32(128))
}
func (n *nemotronHNanoVLModel) soundSampleRate() uint32 {
return cmp.Or(n.SoundConfig.SamplingRate, uint32(16000))
}
func (n *nemotronHNanoVLModel) soundSubsamplingFactor() uint32 {
return cmp.Or(n.SoundConfig.SubsamplingFactor, uint32(8))
}
func (n *nemotronHNanoVLModel) soundSubsamplingConvChannels() uint32 {
return cmp.Or(n.SoundConfig.SubsamplingConvChannels, uint32(256))
}
func (n *nemotronHNanoVLModel) soundSubsamplingConvKernelSize() uint32 {
return cmp.Or(n.SoundConfig.SubsamplingConvKernelSize, uint32(3))
}
func (n *nemotronHNanoVLModel) soundSubsamplingConvStride() uint32 {
return cmp.Or(n.SoundConfig.SubsamplingConvStride, uint32(2))
}
func (n *nemotronHNanoVLModel) soundProjectionHiddenSize() uint32 {
return cmp.Or(n.SoundConfig.ProjectionHiddenSize, uint32(4096))
}
var (
imageNetStandardMean = []float32{0.48145466, 0.4578275, 0.40821073}
imageNetStandardSTD = []float32{0.26862954, 0.26130258, 0.27577711}
)
func (n *nemotronHModel) parseMore(_ fs.FS) error {
if n.NumHiddenLayers == 0 {
+310
View File
@@ -217,6 +217,316 @@ func TestNemotronHLoadModelMetadata(t *testing.T) {
}
}
func TestNemotronHNanoVLLoadModelMetadata(t *testing.T) {
tempDir := t.TempDir()
config := `{
"architectures": ["NemotronH_Nano_VL_V2"],
"model_type": "NemotronH_Nano_VL_V2",
"max_sequence_length": 131072,
"force_image_size": 512,
"downsample_ratio": 0.5,
"patch_size": 16,
"use_thumbnail": true,
"img_context_token_id": 18,
"img_context_token": "<image>",
"img_start_token": "<img>",
"img_end_token": "</img>",
"sound_context_token_id": 27,
"sound_context_token": "<so_embedding>",
"vit_hidden_size": 1280,
"projector_hidden_size": 20480,
"norm_mean": [0.48145466, 0.4578275, 0.40821073],
"norm_std": [0.26862954, 0.26130258, 0.27577711],
"vision_config": {
"version": "radio_v2.5-h",
"patch_size": 16,
"max_resolution": 2048,
"separate_video_embedder": true
},
"sound_config": {
"model_type": "parakeet",
"hidden_size": 1024,
"num_attention_heads": 8,
"num_hidden_layers": 24,
"intermediate_size": 4096,
"conv_kernel_size": 9,
"subsampling_conv_channels": 256,
"subsampling_conv_kernel_size": 3,
"subsampling_conv_stride": 2,
"subsampling_factor": 8,
"num_mel_bins": 128,
"projection_hidden_size": 4096,
"sampling_rate": 16000
},
"llm_config": {
"architectures": ["NemotronHForCausalLM"],
"model_type": "nemotron_h",
"num_hidden_layers": 4,
"hidden_size": 512,
"max_position_embeddings": 262144,
"num_attention_heads": 8,
"num_key_value_heads": 2,
"head_dim": 64,
"layer_norm_epsilon": 1e-5,
"conv_kernel": 4,
"ssm_state_size": 128,
"mamba_num_heads": 16,
"mamba_head_dim": 32,
"n_groups": 8,
"hybrid_override_pattern": "ME*M",
"n_routed_experts": 16,
"num_experts_per_tok": 4,
"moe_intermediate_size": 256
}
}`
if err := os.WriteFile(filepath.Join(tempDir, "config.json"), []byte(config), 0o644); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(tempDir, "preprocessor_config.json"), []byte(`{
"image_size": 512,
"patch_size": 16,
"downsample_ratio": 0.5,
"max_num_tiles": 12,
"use_thumbnail": true,
"norm_mean": [0.48145466, 0.4578275, 0.40821073],
"norm_std": [0.26862954, 0.26130258, 0.27577711]
}`), 0o644); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(tempDir, "tokenizer.json"), []byte(`{}`), 0o644); err != nil {
t.Fatal(err)
}
conv, tokenizer, err := LoadModelMetadata(os.DirFS(tempDir))
if err != nil {
t.Fatal(err)
}
if _, ok := conv.(*nemotronHNanoVLModel); !ok {
t.Fatalf("unexpected converter type: %T", conv)
}
kv := conv.KV(tokenizer)
if got, want := kv["general.architecture"], "nemotron_h_omni"; got != want {
t.Fatalf("unexpected architecture: got %v want %v", got, want)
}
if got, want := kv["context_length"], uint32(131072); got != want {
t.Fatalf("unexpected context length: got %v want %v", got, want)
}
if got, want := kv["vision.block_count"], uint32(32); got != want {
t.Fatalf("unexpected vision block count: got %v want %v", got, want)
}
if got, want := kv["vision.image_size"], uint32(512); got != want {
t.Fatalf("unexpected vision image size: got %v want %v", got, want)
}
if got, want := kv["vision.projector.scale_factor"], uint32(2); got != want {
t.Fatalf("unexpected projector scale factor: got %v want %v", got, want)
}
if got, want := kv["audio.block_count"], uint32(24); got != want {
t.Fatalf("unexpected audio block count: got %v want %v", got, want)
}
if got, want := kv["audio.sound_token_id"], uint32(27); got != want {
t.Fatalf("unexpected audio token id: got %v want %v", got, want)
}
if got, want := kv["audio.subsampling_factor"], uint32(8); got != want {
t.Fatalf("unexpected audio subsampling factor: got %v want %v", got, want)
}
}
func TestNemotronHNanoOmniReasoningV3LoadModelMetadata(t *testing.T) {
tempDir := t.TempDir()
config := `{
"architectures": ["NemotronH_Nano_Omni_Reasoning_V3"],
"model_type": "NemotronH_Nano_Omni_Reasoning_V3",
"max_sequence_length": 131072,
"force_image_size": 512,
"downsample_ratio": 0.5,
"patch_size": 16,
"img_context_token_id": 18,
"img_context_token": "<image>",
"img_start_token": "<img>",
"img_end_token": "</img>",
"sound_context_token_id": 27,
"sound_context_token": "<so_embedding>",
"vit_hidden_size": 1280,
"projector_hidden_size": 4096,
"vision_config": {
"version": "radio_v2.5-h",
"patch_size": 16,
"min_num_patches": 1024,
"max_num_patches": 13312,
"args": {
"min_num_patches": 1024,
"max_num_patches": 13312
}
},
"sound_config": {
"model_type": "parakeet",
"hidden_size": 1024,
"num_attention_heads": 8,
"num_hidden_layers": 24,
"intermediate_size": 4096,
"conv_kernel_size": 9,
"subsampling_conv_channels": 256,
"subsampling_conv_kernel_size": 3,
"subsampling_conv_stride": 2,
"subsampling_factor": 8,
"num_mel_bins": 128,
"projection_hidden_size": 4096,
"sampling_rate": 16000
},
"llm_config": {
"architectures": ["NemotronHForCausalLM"],
"model_type": "nemotron_h",
"num_hidden_layers": 4,
"hidden_size": 512,
"max_position_embeddings": 262144,
"num_attention_heads": 8,
"num_key_value_heads": 2,
"head_dim": 64,
"layer_norm_epsilon": 1e-5,
"conv_kernel": 4,
"ssm_state_size": 128,
"mamba_num_heads": 16,
"mamba_head_dim": 32,
"n_groups": 8,
"hybrid_override_pattern": "ME*M",
"n_routed_experts": 16,
"num_experts_per_tok": 4,
"moe_intermediate_size": 256
}
}`
if err := os.WriteFile(filepath.Join(tempDir, "config.json"), []byte(config), 0o644); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(tempDir, "tokenizer.json"), []byte(`{}`), 0o644); err != nil {
t.Fatal(err)
}
conv, tokenizer, err := LoadModelMetadata(os.DirFS(tempDir))
if err != nil {
t.Fatal(err)
}
if _, ok := conv.(*nemotronHNanoVLModel); !ok {
t.Fatalf("unexpected converter type: %T", conv)
}
kv := conv.KV(tokenizer)
if got, want := kv["general.architecture"], "nemotron_h_omni"; got != want {
t.Fatalf("unexpected architecture: got %v want %v", got, want)
}
if got, want := kv["vision.block_count"], uint32(32); got != want {
t.Fatalf("unexpected vision block count: got %v want %v", got, want)
}
if got, want := kv["vision.min_num_patches"], uint32(1024); got != want {
t.Fatalf("unexpected vision min patches: got %v want %v", got, want)
}
if got, want := kv["vision.max_num_patches"], uint32(13312); got != want {
t.Fatalf("unexpected vision max patches: got %v want %v", got, want)
}
if got, want := kv["audio.block_count"], uint32(24); got != want {
t.Fatalf("unexpected audio block count: got %v want %v", got, want)
}
if got, want := kv["audio.sound_token_id"], uint32(27); got != want {
t.Fatalf("unexpected audio token id: got %v want %v", got, want)
}
}
func TestNemotronHNanoVLTensorsRetainVisionAndAudio(t *testing.T) {
m := &nemotronHNanoVLModel{
LLMConfig: nemotronHModel{NGroups: 8},
}
in := []Tensor{
&fakeTensor{
name: "blk.0.ssm_a",
shape: []uint64{4},
data: []float32{0, 1, 2, 3},
},
&fakeTensor{name: "v.blk.0.attn_qkv.weight", shape: []uint64{3840, 1280}},
&fakeTensor{name: "v.position_embd", shape: []uint64{1, 16384, 1280}},
&fakeTensor{name: "v.cls_embd", shape: []uint64{10, 1280}},
&fakeTensor{name: "mm.norm.weight", shape: []uint64{5120}},
&fakeTensor{name: "a.feature_extractor.fb", shape: []uint64{1, 128, 257}},
&fakeTensor{name: "a.subsampling.dw1.weight", shape: []uint64{256, 1, 3, 3}},
&fakeTensor{name: "a.blk.0.conv_dw.weight", shape: []uint64{1024, 1, 9}},
&fakeTensor{name: "a.blk.0.conv_pw1.weight", shape: []uint64{2048, 1024, 1}},
&fakeTensor{name: "a.blk.0.conv_bn.num_batches_tracked", shape: []uint64{1}},
&fakeTensor{name: "mm.a.1.weight", shape: []uint64{4096, 1024}},
}
out := m.Tensors(in)
got := map[string][]uint64{}
for _, tns := range out {
got[tns.Name] = tns.Shape
}
for _, name := range []string{
"blk.0.ssm_a",
"v.blk.0.attn_q.weight",
"v.blk.0.attn_k.weight",
"v.blk.0.attn_v.weight",
"v.position_embd",
"v.cls_embd",
"mm.norm.weight",
"a.feature_extractor.fb",
"a.subsampling.dw1.weight",
"a.blk.0.conv_dw.weight",
"a.blk.0.conv_pw1.weight",
"mm.a.1.weight",
} {
if _, ok := got[name]; !ok {
t.Fatalf("expected tensor %q in output", name)
}
}
if gotShape, want := got["blk.0.ssm_a"], []uint64{4, 1}; !slices.Equal(gotShape, want) {
t.Fatalf("unexpected ssm_a shape: got %v want %v", gotShape, want)
}
if gotShape, want := got["v.position_embd"], []uint64{16384, 1280}; !slices.Equal(gotShape, want) {
t.Fatalf("unexpected position embedding shape: got %v want %v", gotShape, want)
}
if gotShape, want := got["a.blk.0.conv_dw.weight"], []uint64{1024, 9}; !slices.Equal(gotShape, want) {
t.Fatalf("unexpected audio conv_dw shape: got %v want %v", gotShape, want)
}
if gotShape, want := got["a.blk.0.conv_pw1.weight"], []uint64{2048, 1024}; !slices.Equal(gotShape, want) {
t.Fatalf("unexpected audio conv_pw1 shape: got %v want %v", gotShape, want)
}
if _, ok := got["a.blk.0.conv_bn.num_batches_tracked"]; ok {
t.Fatal("audio batchnorm num_batches_tracked should be omitted")
}
}
func TestNemotronHNanoVLReplacements(t *testing.T) {
m := &nemotronHNanoVLModel{}
r := strings.NewReplacer(m.Replacements()...)
if got, want := r.Replace("language_model.backbone.layers.1.mixer.fc1_latent_proj.weight"), "blk.1.ffn_latent_in.weight"; got != want {
t.Fatalf("unexpected fc1 replacement: got %q want %q", got, want)
}
if got, want := r.Replace("language_model.lm_head.weight"), "output.weight"; got != want {
t.Fatalf("unexpected lm_head replacement: got %q want %q", got, want)
}
if got, want := r.Replace("vision_model.radio_model.model.blocks.0.attn.qkv.weight"), "v.blk.0.attn_qkv.weight"; got != want {
t.Fatalf("unexpected vision replacement: got %q want %q", got, want)
}
if got, want := r.Replace("mlp1.1.weight"), "mm.1.weight"; got != want {
t.Fatalf("unexpected projector replacement: got %q want %q", got, want)
}
if got, want := r.Replace("sound_encoder.encoder.layers.0.self_attn.q_proj.weight"), "a.blk.0.attn_q.weight"; got != want {
t.Fatalf("unexpected audio q_proj replacement: got %q want %q", got, want)
}
if got, want := r.Replace("sound_encoder.encoder.layers.0.conv.pointwise_conv1.weight"), "a.blk.0.conv_pw1.weight"; got != want {
t.Fatalf("unexpected audio conv replacement: got %q want %q", got, want)
}
if got, want := r.Replace("sound_projection.linear2.weight"), "mm.a.2.weight"; got != want {
t.Fatalf("unexpected audio projector replacement: got %q want %q", got, want)
}
}
func TestNemotronHReplacementsLatentProjections(t *testing.T) {
m := &nemotronHModel{}
r := strings.NewReplacer(m.Replacements()...)
+4 -2
View File
@@ -42,8 +42,10 @@ func (t tensorBase) Kind() uint32 {
strings.HasSuffix(t.name, ".bias") ||
strings.HasSuffix(t.name, ".shortconv.conv.weight") ||
strings.HasSuffix(t.name, ".ssm_conv1d.weight") || // SSM conv kernel must be F32 for Metal
strings.HasPrefix(t.name, "a.conv1d.") || // audio SSCP conv weights must be F32 for im2col
strings.Contains(t.name, ".conv_dw.") || // audio depthwise conv weights must be F32
strings.HasPrefix(t.name, "a.feature_extractor.") || // audio feature-extractor constants are read with BackendGet and must be real F32 values
strings.HasPrefix(t.name, "a.conv1d.") || // audio SSCP conv weights are kept F32 for im2col; this likely slows audio and should be revisited
strings.HasPrefix(t.name, "a.subsampling.") || // audio Parakeet subsampling weights are kept F32 for conv/linear stability; this likely slows audio and should be revisited
strings.Contains(t.name, ".conv_dw.") || // audio depthwise conv weights are kept F32; this likely slows audio and should be revisited
t.name == "token_types.weight" ||
t.name == "v.positional_embedding_vlm" ||
t.name == "v.position_embd.weight" ||
+416 -13
View File
@@ -5,10 +5,12 @@ import (
"bytes"
"encoding/binary"
"encoding/json"
"errors"
"fmt"
"io"
"io/fs"
"maps"
"math"
"slices"
"strings"
@@ -23,6 +25,11 @@ type safetensorMetadata struct {
}
func parseSafetensors(fsys fs.FS, replacer *strings.Replacer, ps ...string) ([]Tensor, error) {
fp8Block, err := safetensorsFP8BlockSize(fsys)
if err != nil {
return nil, err
}
var ts []Tensor
for _, p := range ps {
f, err := fsys.Open(p)
@@ -50,24 +57,47 @@ func parseSafetensors(fsys fs.FS, replacer *strings.Replacer, ps ...string) ([]T
names := make(map[string]struct{}, len(keys))
fp8Scales, err := collectSafetensorsFP8Scales(n, headers)
if err != nil {
return nil, err
}
for _, key := range keys {
if value := headers[key]; value.Type != "" {
if _, ok := fp8Scales.consumed[key]; ok {
continue
}
// Scalar tensors (e.g. clipped linear min/max) are 0-dim in safetensors.
// Promote them to 1-dim so they can be stored in GGUF.
if len(value.Shape) == 0 {
value.Shape = []uint64{1}
}
var scale *safetensorScale
if value.Type == "F8_E4M3" {
if !fp8Block.ok {
return nil, fmt.Errorf("missing fp8 block size metadata for tensor %q", key)
}
scale = fp8Scales.byWeight[key]
if scale == nil {
return nil, fmt.Errorf("missing fp8 scale companion for tensor %q", key)
}
}
ggufName := replacer.Replace(key)
if _, ok := names[ggufName]; ok {
return nil, fmt.Errorf("duplicate tensor name '%s' was found for this model", ggufName)
}
names[ggufName] = struct{}{}
ts = append(ts, safetensor{
fs: fsys,
path: p,
dtype: value.Type,
offset: safetensorsPad(n, value.Offsets[0]),
size: safetensorsPad(n, value.Offsets[1]) - safetensorsPad(n, value.Offsets[0]),
fs: fsys,
path: p,
dtype: value.Type,
offset: safetensorsPad(n, value.Offsets[0]),
size: safetensorsPad(n, value.Offsets[1]) - safetensorsPad(n, value.Offsets[0]),
scale: scale,
fp8Block: fp8Block,
tensorBase: &tensorBase{
name: ggufName,
shape: value.Shape,
@@ -85,12 +115,22 @@ func safetensorsPad(n, offset int64) int64 {
return 8 + n + offset
}
type safetensor struct {
fs fs.FS
path string
type safetensorScale struct {
name string
dtype string
shape []uint64
offset int64
size int64
}
type safetensor struct {
fs fs.FS
path string
dtype string
offset int64
size int64
scale *safetensorScale
fp8Block safetensorFP8BlockSize
*tensorBase
}
@@ -104,17 +144,26 @@ func (st safetensor) Kind() uint32 {
kind != tensorKindFP32 {
kind = tensorKindBF16
}
if st.dtype == "F8_E4M3" && kind != tensorKindFP32 {
kind = tensorKindBF16
}
return kind
}
func (st safetensor) SourceDType() string {
return st.dtype
}
func (st safetensor) Clone() Tensor {
return &safetensor{
fs: st.fs,
path: st.path,
dtype: st.dtype,
offset: st.offset,
size: st.size,
fs: st.fs,
path: st.path,
dtype: st.dtype,
offset: st.offset,
size: st.size,
scale: st.scale.Clone(),
fp8Block: st.fp8Block,
tensorBase: &tensorBase{
name: st.name,
repacker: st.repacker,
@@ -123,6 +172,19 @@ func (st safetensor) Clone() Tensor {
}
}
func (ss *safetensorScale) Clone() *safetensorScale {
if ss == nil {
return nil
}
return &safetensorScale{
name: ss.name,
dtype: ss.dtype,
shape: slices.Clone(ss.shape),
offset: ss.offset,
size: ss.size,
}
}
func (st safetensor) WriteTo(w io.Writer) (int64, error) {
f, err := st.fs.Open(st.path)
if err != nil {
@@ -180,6 +242,16 @@ func (st safetensor) WriteTo(w io.Writer) (int64, error) {
}
f32s = bfloat16.DecodeFloat32(u8s)
case "F8_E4M3":
u8s := make([]uint8, st.size)
if err = binary.Read(br, binary.LittleEndian, u8s); err != nil {
return 0, err
}
f32s, err = st.decodeFP8E4M3(u8s)
if err != nil {
return 0, err
}
default:
return 0, fmt.Errorf("unknown data type: %s", st.dtype)
}
@@ -208,3 +280,334 @@ func (st safetensor) WriteTo(w io.Writer) (int64, error) {
return 0, fmt.Errorf("unknown storage type: %d", st.Kind())
}
}
type safetensorsFP8Scales struct {
byWeight map[string]*safetensorScale
consumed map[string]struct{}
}
func collectSafetensorsFP8Scales(n int64, headers map[string]safetensorMetadata) (safetensorsFP8Scales, error) {
scales := safetensorsFP8Scales{
byWeight: make(map[string]*safetensorScale),
consumed: make(map[string]struct{}),
}
for key, value := range headers {
if value.Type != "F8_E4M3" {
continue
}
scaleKey, scaleValue, ok, err := safetensorsFP8Scale(key, headers)
if err != nil {
return safetensorsFP8Scales{}, err
}
if !ok {
continue
}
if _, ok := scales.consumed[scaleKey]; ok {
return safetensorsFP8Scales{}, fmt.Errorf("fp8 scale companion %q is used by multiple tensors", scaleKey)
}
scales.byWeight[key] = &safetensorScale{
name: scaleKey,
dtype: scaleValue.Type,
shape: slices.Clone(scaleValue.Shape),
offset: safetensorsPad(n, scaleValue.Offsets[0]),
size: safetensorsPad(n, scaleValue.Offsets[1]) - safetensorsPad(n, scaleValue.Offsets[0]),
}
scales.consumed[scaleKey] = struct{}{}
}
return scales, nil
}
func safetensorsFP8Scale(key string, headers map[string]safetensorMetadata) (string, safetensorMetadata, bool, error) {
candidates := safetensorsFP8ScaleCandidates(key)
var scaleKey string
var scaleValue safetensorMetadata
if strings.HasSuffix(key, ".weight") {
// Keep support for compressed-tensors exports that place the scale name
// between the module path and weight suffix.
base := strings.TrimSuffix(key, ".weight")
candidates = appendUnique(candidates, base+".weight_scale")
candidates = appendUnique(candidates, base+".weight_scale_inv")
}
for _, candidate := range candidates {
if value, ok := headers[candidate]; ok && value.Type != "" {
if scaleKey != "" {
return "", safetensorMetadata{}, false, fmt.Errorf("multiple fp8 scale companions for tensor %q: %q and %q", key, scaleKey, candidate)
}
scaleKey = candidate
scaleValue = value
}
}
if scaleKey == "" {
return "", safetensorMetadata{}, false, nil
}
return scaleKey, scaleValue, true, nil
}
func safetensorsFP8ScaleCandidates(key string) []string {
var candidates []string
candidates = appendUnique(candidates, key+"_scale")
candidates = appendUnique(candidates, key+"_scale_inv")
candidates = appendUnique(candidates, key+".scale")
candidates = appendUnique(candidates, key+".scale_inv")
return candidates
}
func appendUnique(values []string, value string) []string {
if !slices.Contains(values, value) {
values = append(values, value)
}
return values
}
type safetensorFP8BlockSize struct {
rows int
cols int
ok bool
}
type safetensorsSourceQuantization struct {
QuantMethod string `json:"quant_method"`
Format string `json:"format"`
WeightBlockSize []int `json:"weight_block_size"`
ConfigGroups map[string]struct {
Format string `json:"format"`
Weights struct {
BlockStructure []int `json:"block_structure"`
NumBits int `json:"num_bits"`
Type string `json:"type"`
} `json:"weights"`
} `json:"config_groups"`
}
type safetensorsModelConfig struct {
Quantization safetensorsSourceQuantization `json:"quantization"`
QuantizationConfig safetensorsSourceQuantization `json:"quantization_config"`
CompressionConfig safetensorsSourceQuantization `json:"compression_config"`
TextConfig struct {
Quantization safetensorsSourceQuantization `json:"quantization"`
QuantizationConfig safetensorsSourceQuantization `json:"quantization_config"`
CompressionConfig safetensorsSourceQuantization `json:"compression_config"`
} `json:"text_config"`
}
func safetensorsFP8BlockSize(fsys fs.FS) (safetensorFP8BlockSize, error) {
bts, err := fs.ReadFile(fsys, "config.json")
if errors.Is(err, fs.ErrNotExist) {
return safetensorFP8BlockSize{}, nil
}
if err != nil {
return safetensorFP8BlockSize{}, err
}
bts = sanitizeNonFiniteJSON(bts)
var cfg safetensorsModelConfig
if err := json.Unmarshal(bts, &cfg); err != nil {
return safetensorFP8BlockSize{}, fmt.Errorf("parse config.json fp8 metadata: %w", err)
}
var blocks []safetensorFP8BlockSize
for _, q := range []safetensorsSourceQuantization{
cfg.Quantization,
cfg.QuantizationConfig,
cfg.CompressionConfig,
cfg.TextConfig.Quantization,
cfg.TextConfig.QuantizationConfig,
cfg.TextConfig.CompressionConfig,
} {
if strings.EqualFold(q.QuantMethod, "fp8") && len(q.WeightBlockSize) == 2 {
block, err := newSafetensorFP8BlockSize(q.WeightBlockSize[0], q.WeightBlockSize[1])
if err != nil {
return safetensorFP8BlockSize{}, err
}
blocks = append(blocks, block)
}
if !strings.EqualFold(q.QuantMethod, "compressed-tensors") && !strings.EqualFold(q.Format, "float-quantized") {
continue
}
for _, group := range q.ConfigGroups {
if !strings.EqualFold(group.Format, "float-quantized") ||
group.Weights.NumBits != 8 ||
!strings.EqualFold(group.Weights.Type, "float") ||
len(group.Weights.BlockStructure) != 2 {
continue
}
block, err := newSafetensorFP8BlockSize(group.Weights.BlockStructure[0], group.Weights.BlockStructure[1])
if err != nil {
return safetensorFP8BlockSize{}, err
}
blocks = append(blocks, block)
}
}
if len(blocks) == 0 {
return safetensorFP8BlockSize{}, nil
}
block := blocks[0]
for _, other := range blocks[1:] {
if other.rows != block.rows || other.cols != block.cols {
return safetensorFP8BlockSize{}, fmt.Errorf("multiple fp8 block sizes in config.json: %dx%d and %dx%d", block.rows, block.cols, other.rows, other.cols)
}
}
return block, nil
}
func newSafetensorFP8BlockSize(rows, cols int) (safetensorFP8BlockSize, error) {
if rows <= 0 || cols <= 0 {
return safetensorFP8BlockSize{}, fmt.Errorf("invalid fp8 block size %dx%d", rows, cols)
}
return safetensorFP8BlockSize{rows: rows, cols: cols, ok: true}, nil
}
func (st safetensor) decodeFP8E4M3(data []byte) ([]float32, error) {
if st.scale == nil {
return nil, fmt.Errorf("missing fp8 scale companion for tensor %q", st.name)
}
if !st.fp8Block.ok {
return nil, fmt.Errorf("missing fp8 block size metadata for tensor %q", st.name)
}
if len(st.shape) != 2 {
return nil, fmt.Errorf("expected 2D fp8 tensor %q, got shape %v", st.name, st.shape)
}
rows, cols := int(st.shape[0]), int(st.shape[1])
if rows < 0 || cols < 0 || rows*cols != len(data) {
return nil, fmt.Errorf("fp8 tensor %q shape %v does not match %d bytes", st.name, st.shape, len(data))
}
scale, err := st.readScale()
if err != nil {
return nil, err
}
if len(st.scale.shape) != 2 {
return nil, fmt.Errorf("expected 2D fp8 scale tensor %q, got shape %v", st.scale.name, st.scale.shape)
}
blockRows := st.fp8Block.rows
blockCols := st.fp8Block.cols
scaleRows, scaleCols := int(st.scale.shape[0]), int(st.scale.shape[1])
expectedRows := (rows + blockRows - 1) / blockRows
expectedCols := (cols + blockCols - 1) / blockCols
if scaleRows != expectedRows || scaleCols != expectedCols {
return nil, fmt.Errorf("unexpected fp8 scale shape %v for tensor %q shape %v; want [%d %d]", st.scale.shape, st.name, st.shape, expectedRows, expectedCols)
}
if len(scale) != scaleRows*scaleCols {
return nil, fmt.Errorf("fp8 scale tensor %q shape %v does not match decoded length %d", st.scale.name, st.scale.shape, len(scale))
}
f32s := make([]float32, len(data))
for r := range rows {
scaleRow := r / blockRows
rowOffset := r * cols
for c := range cols {
f32s[rowOffset+c] = decodeFloat8E4M3FN(data[rowOffset+c]) * scale[scaleRow*scaleCols+c/blockCols]
}
}
return f32s, nil
}
func (st safetensor) readScale() ([]float32, error) {
r, err := st.sectionReader(st.scale.offset, st.scale.size)
if err != nil {
return nil, fmt.Errorf("failed to read fp8 scale tensor %q: %w", st.scale.name, err)
}
if closer, ok := r.(io.Closer); ok {
defer closer.Close()
}
br := bufio.NewReaderSize(r, min(32<<10, int(st.scale.size)))
switch st.scale.dtype {
case "F32":
f32s := make([]float32, st.scale.size/4)
if err := binary.Read(br, binary.LittleEndian, f32s); err != nil {
return nil, err
}
return f32s, nil
case "F16":
u16s := make([]uint16, st.scale.size/2)
if err := binary.Read(br, binary.LittleEndian, u16s); err != nil {
return nil, err
}
f32s := make([]float32, len(u16s))
for i := range u16s {
f32s[i] = float16.Frombits(u16s[i]).Float32()
}
return f32s, nil
case "BF16":
u8s := make([]uint8, st.scale.size)
if err := binary.Read(br, binary.LittleEndian, u8s); err != nil {
return nil, err
}
return bfloat16.DecodeFloat32(u8s), nil
default:
return nil, fmt.Errorf("unsupported fp8 scale dtype %q for tensor %q", st.scale.dtype, st.scale.name)
}
}
func (st safetensor) sectionReader(offset, size int64) (io.Reader, error) {
f, err := st.fs.Open(st.path)
if err != nil {
return nil, err
}
if readerAt, ok := f.(io.ReaderAt); ok {
return &readCloserReader{
Reader: io.NewSectionReader(readerAt, offset, size),
Closer: f,
}, nil
}
if seeker, ok := f.(io.Seeker); ok {
if _, err := seeker.Seek(offset, io.SeekStart); err != nil {
f.Close()
return nil, err
}
return &readCloserReader{
Reader: io.LimitReader(f, size),
Closer: f,
}, nil
}
if _, err := io.CopyN(io.Discard, f, offset); err != nil {
f.Close()
return nil, err
}
return &readCloserReader{
Reader: io.LimitReader(f, size),
Closer: f,
}, nil
}
type readCloserReader struct {
io.Reader
io.Closer
}
func decodeFloat8E4M3FN(v byte) float32 {
sign := float32(1)
if v&0x80 != 0 {
sign = -1
}
exp := int((v >> 3) & 0x0f)
mant := int(v & 0x07)
if exp == 0 {
if mant == 0 {
return 0 * sign
}
return sign * float32(math.Ldexp(float64(mant)/8, -6))
}
if exp == 0x0f && mant == 0x07 {
return float32(math.NaN())
}
return sign * float32(math.Ldexp(1+float64(mant)/8, exp-7))
}
+229
View File
@@ -3,8 +3,10 @@ package convert
import (
"bytes"
"encoding/binary"
"fmt"
"os"
"path/filepath"
"strings"
"testing"
"github.com/d4l3k/go-bfloat16"
@@ -231,6 +233,222 @@ func TestSafetensors(t *testing.T) {
}
}
func TestSafetensorWriteToFP8E4M3(t *testing.T) {
root, err := os.OpenRoot(t.TempDir())
if err != nil {
t.Fatal(err)
}
defer root.Close()
path := filepath.Base(t.Name())
f, err := root.Create(path)
if err != nil {
t.Fatal(err)
}
// E4M3FN encodings for 1.0, 2.0, 0.5, and -1.0.
if _, err := f.Write([]byte{0x38, 0x40, 0x30, 0xb8}); err != nil {
t.Fatal(err)
}
if _, err := f.Write(bfloat16.EncodeFloat32([]float32{2})); err != nil {
t.Fatal(err)
}
if err := f.Close(); err != nil {
t.Fatal(err)
}
st := safetensor{
fs: root.FS(),
path: path,
dtype: "F8_E4M3",
offset: 0,
size: 4,
fp8Block: safetensorFP8BlockSize{rows: 128, cols: 128, ok: true},
scale: &safetensorScale{
name: "linear.weight_scale",
dtype: "BF16",
shape: []uint64{1, 1},
offset: 4,
size: 2,
},
tensorBase: &tensorBase{
name: "linear.weight",
shape: []uint64{2, 2},
},
}
var b bytes.Buffer
if _, err := st.WriteTo(&b); err != nil {
t.Fatal(err)
}
want := bfloat16.EncodeFloat32([]float32{2, 4, 1, -2})
if diff := cmp.Diff(want, b.Bytes()); diff != "" {
t.Errorf("safetensor.WriteTo() mismatch (-want +got):\n%s", diff)
}
}
func TestSafetensorWriteToFP8E4M3UsesConfiguredBlockSize(t *testing.T) {
root, err := os.OpenRoot(t.TempDir())
if err != nil {
t.Fatal(err)
}
defer root.Close()
path := filepath.Base(t.Name())
f, err := root.Create(path)
if err != nil {
t.Fatal(err)
}
if _, err := f.Write(bytes.Repeat([]byte{0x38}, 12)); err != nil {
t.Fatal(err)
}
if _, err := f.Write(bfloat16.EncodeFloat32([]float32{1, 2, 3, 4})); err != nil {
t.Fatal(err)
}
if err := f.Close(); err != nil {
t.Fatal(err)
}
st := safetensor{
fs: root.FS(),
path: path,
dtype: "F8_E4M3",
offset: 0,
size: 12,
fp8Block: safetensorFP8BlockSize{rows: 2, cols: 3, ok: true},
scale: &safetensorScale{
name: "linear.weight_scale",
dtype: "BF16",
shape: []uint64{2, 2},
offset: 12,
size: 8,
},
tensorBase: &tensorBase{
name: "linear.weight",
shape: []uint64{3, 4},
},
}
var b bytes.Buffer
if _, err := st.WriteTo(&b); err != nil {
t.Fatal(err)
}
want := bfloat16.EncodeFloat32([]float32{
1, 1, 1, 2,
1, 1, 1, 2,
3, 3, 3, 4,
})
if diff := cmp.Diff(want, b.Bytes()); diff != "" {
t.Errorf("safetensor.WriteTo() mismatch (-want +got):\n%s", diff)
}
}
func TestParseSafetensorsConsumesFP8ScaleCompanion(t *testing.T) {
tempDir := t.TempDir()
generateSafetensorTestData(t, tempDir, map[string]*tensorData{
"linear.weight": {
Offsets: []int{0, 4},
Type: "F8_E4M3",
Shape: []int{2, 2},
},
"linear.weight_scale": {
Offsets: []int{4, 6},
Type: "BF16",
Shape: []int{1, 1},
},
})
writeFP8BlockConfig(t, tempDir, 128, 128)
tensors, err := parseSafetensors(os.DirFS(tempDir), strings.NewReplacer(), "model-00001-of-00001.safetensors")
if err != nil {
t.Fatal(err)
}
if len(tensors) != 1 {
t.Fatalf("expected one tensor, got %d", len(tensors))
}
if got := tensors[0].Name(); got != "linear.weight" {
t.Fatalf("unexpected tensor name %q", got)
}
if got := tensors[0].Kind(); got != tensorKindBF16 {
t.Fatalf("unexpected fp8 converted kind %d, want %d", got, tensorKindBF16)
}
}
func TestParseSafetensorsRejectsFP8WithoutBlockMetadata(t *testing.T) {
tempDir := t.TempDir()
generateSafetensorTestData(t, tempDir, map[string]*tensorData{
"linear.weight": {
Offsets: []int{0, 4},
Type: "F8_E4M3",
Shape: []int{2, 2},
},
"linear.weight_scale": {
Offsets: []int{4, 6},
Type: "BF16",
Shape: []int{1, 1},
},
})
_, err := parseSafetensors(os.DirFS(tempDir), strings.NewReplacer(), "model-00001-of-00001.safetensors")
if err == nil || !strings.Contains(err.Error(), "missing fp8 block size metadata") {
t.Fatalf("expected missing fp8 block size metadata error, got %v", err)
}
}
func TestParseSafetensorsRejectsAmbiguousFP8ScaleCompanion(t *testing.T) {
tempDir := t.TempDir()
generateSafetensorTestData(t, tempDir, map[string]*tensorData{
"linear.weight": {
Offsets: []int{0, 4},
Type: "F8_E4M3",
Shape: []int{2, 2},
},
"linear.weight_scale": {
Offsets: []int{4, 6},
Type: "BF16",
Shape: []int{1, 1},
},
"linear.weight.scale": {
Offsets: []int{6, 8},
Type: "BF16",
Shape: []int{1, 1},
},
})
writeFP8BlockConfig(t, tempDir, 128, 128)
_, err := parseSafetensors(os.DirFS(tempDir), strings.NewReplacer(), "model-00001-of-00001.safetensors")
if err == nil || !strings.Contains(err.Error(), "multiple fp8 scale companions") {
t.Fatalf("expected ambiguous fp8 scale companion error, got %v", err)
}
}
func writeFP8BlockConfig(t *testing.T, dir string, rows, cols int) {
t.Helper()
config := fmt.Sprintf(`{
"architectures": ["GenericForCausalLM"],
"compression_config": {
"format": "float-quantized",
"config_groups": {
"group_0": {
"format": "float-quantized",
"weights": {
"type": "float",
"num_bits": 8,
"block_structure": [%d, %d]
}
}
}
}
}`, rows, cols)
if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(config), 0o644); err != nil {
t.Fatal(err)
}
}
func TestSafetensorKind(t *testing.T) {
tests := []struct {
name string
@@ -259,6 +477,17 @@ func TestSafetensorKind(t *testing.T) {
},
expected: tensorKindFP16,
},
{
name: "BF16 audio feature extractor constants should return FP32",
st: safetensor{
tensorBase: &tensorBase{
name: "a.feature_extractor.fb",
shape: []uint64{1, 128, 257},
},
dtype: "BF16",
},
expected: tensorKindFP32,
},
{
name: "BF16 dtype with FP32 base kind should return FP32",
st: safetensor{
+52
View File
@@ -5,6 +5,7 @@ import (
"errors"
"io"
"iter"
"maps"
"path"
"slices"
"strconv"
@@ -153,3 +154,54 @@ func (g mergeGroup) WriteTo(w io.Writer) (int64, error) {
return 0, nil
}
func sourceTensorKV(ts []*ggml.Tensor) KV {
sourceFP8 := make(map[string]struct{})
for _, t := range ts {
if writerSourceDType(t.WriterTo) == "F8_E4M3" {
sourceFP8[t.Name] = struct{}{}
}
}
if len(sourceFP8) == 0 {
return nil
}
return KV{
"source_quantization": "hf_fp8",
"source_fp8_tensors": slices.Sorted(maps.Keys(sourceFP8)),
}
}
type sourceDTypeTensor interface {
SourceDType() string
}
func writerSourceDType(w io.WriterTo) string {
switch w := w.(type) {
case sourceDTypeTensor:
return w.SourceDType()
case mergeGroup:
if len(w) == 0 {
return ""
}
dtype := sourceDType(w[0])
if dtype == "" {
return ""
}
for _, t := range w[1:] {
if sourceDType(t) != dtype {
return ""
}
}
return dtype
default:
return ""
}
}
func sourceDType(t Tensor) string {
if t, ok := t.(sourceDTypeTensor); ok {
return t.SourceDType()
}
return ""
}
+51 -5
View File
@@ -21,7 +21,8 @@ type fakeTensor struct {
shape []uint64
data []float32
repacker Repacker
sourceDType string
repacker Repacker
}
func (f fakeTensor) Name() string {
@@ -36,16 +37,21 @@ func (f fakeTensor) Kind() uint32 {
return 0
}
func (f fakeTensor) SourceDType() string {
return f.sourceDType
}
func (f *fakeTensor) SetRepacker(fn Repacker) {
f.repacker = fn
}
func (f fakeTensor) Clone() Tensor {
return &fakeTensor{
name: f.name,
shape: slices.Clone(f.shape),
data: slices.Clone(f.data),
repacker: f.repacker,
name: f.name,
shape: slices.Clone(f.shape),
data: slices.Clone(f.data),
sourceDType: f.sourceDType,
repacker: f.repacker,
}
}
@@ -995,3 +1001,43 @@ func TestMergeOrder(t *testing.T) {
})
}
}
func TestSourceTensorKVRecordsFP8OutputTensors(t *testing.T) {
fp8 := &fakeTensor{name: "linear.weight", shape: []uint64{2, 2}, sourceDType: "F8_E4M3"}
bf16 := &fakeTensor{name: "other.weight", shape: []uint64{2, 2}, sourceDType: "BF16"}
kv := sourceTensorKV([]*ggml.Tensor{
{Name: "blk.0.linear.weight", WriterTo: fp8},
{Name: "blk.0.other.weight", WriterTo: bf16},
})
if got := kv["source_quantization"]; got != "hf_fp8" {
t.Fatalf("source_quantization = %v, want hf_fp8", got)
}
got, ok := kv["source_fp8_tensors"].([]string)
if !ok {
t.Fatalf("source_fp8_tensors = %#v, want []string", kv["source_fp8_tensors"])
}
if diff := cmp.Diff([]string{"blk.0.linear.weight"}, got); diff != "" {
t.Fatalf("source_fp8_tensors mismatch (-want +got):\n%s", diff)
}
}
func TestSourceTensorKVRecordsMergedFP8OutputTensors(t *testing.T) {
fp8A := &fakeTensor{name: "expert.0.weight", shape: []uint64{2, 2}, sourceDType: "F8_E4M3"}
fp8B := &fakeTensor{name: "expert.1.weight", shape: []uint64{2, 2}, sourceDType: "F8_E4M3"}
bf16 := &fakeTensor{name: "expert.2.weight", shape: []uint64{2, 2}, sourceDType: "BF16"}
kv := sourceTensorKV([]*ggml.Tensor{
{Name: "ffn_exps.weight", WriterTo: mergeGroup{fp8A, fp8B}},
{Name: "mixed_exps.weight", WriterTo: mergeGroup{fp8A, bf16}},
})
got, ok := kv["source_fp8_tensors"].([]string)
if !ok {
t.Fatalf("source_fp8_tensors = %#v, want []string", kv["source_fp8_tensors"])
}
if diff := cmp.Diff([]string{"ffn_exps.weight"}, got); diff != "" {
t.Fatalf("source_fp8_tensors mismatch (-want +got):\n%s", diff)
}
}
+2
View File
@@ -103,6 +103,8 @@ func parseTokenizer(fsys fs.FS, specialTokenTypes []string) (*Tokenizer, error)
t.Pre = "qwen2"
case "00431aed57e696b747435f734d1e3b9b1bfd931a121fb5cac7129e97c181e9ba":
t.Pre = "qwen35"
case "b92c0824a58e1d8dc3221cf3e12c433c3a86f57e46d57229993489f0798e7702":
t.Pre = "laguna"
case "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855":
// noop, empty pretokenizer
default:
+79 -28
View File
@@ -112,10 +112,9 @@ func GPUDevices(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.
}
ctx1stPass, cancel := context.WithTimeout(ctx, bootstrapTimeout)
defer cancel()
// For this pass, we retain duplicates in case any are incompatible with some libraries
devices = append(devices, bootstrapDevices(ctx1stPass, dirs, nil)...)
devices = append(devices, bootstrapDevicesWithMetalRetry(ctx1stPass, ctx, bootstrapTimeout, dirs, nil)...)
cancel()
}
// In the second pass, we more deeply initialize the GPUs to weed out devices that
@@ -147,9 +146,9 @@ func GPUDevices(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.
wg.Add(1)
go func(i int) {
defer wg.Done()
extraEnvs := ml.GetVisibleDevicesEnv(devices[i:i+1], true)
extraEnvs := ml.GetDevicesEnv(devices[i:i+1], true)
devices[i].AddInitValidation(extraEnvs)
if len(bootstrapDevices(ctx2ndPass, devices[i].LibraryPath, extraEnvs)) == 0 {
if len(bootstrapDevicesWithMetalRetry(ctx2ndPass, ctx, 30*time.Second, devices[i].LibraryPath, extraEnvs)) == 0 {
slog.Debug("filtering device which didn't fully initialize",
"id", devices[i].ID,
"libdir", devices[i].LibraryPath[len(devices[i].LibraryPath)-1],
@@ -334,7 +333,7 @@ func GPUDevices(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.
// Apply any dev filters to avoid re-discovering unsupported devices, and get IDs correct
// We avoid CUDA filters here to keep ROCm from failing to discover GPUs in a mixed environment
devFilter := ml.GetVisibleDevicesEnv(devices, false)
devFilter := ml.GetDevicesEnv(devices, false)
for dir := range libDirs {
updatedDevices := bootstrapDevices(ctx, []string{ml.LibOllamaPath, dir}, devFilter)
@@ -427,27 +426,84 @@ func (r *bootstrapRunner) HasExited() bool {
return false
}
func bootstrapDevices(ctx context.Context, ollamaLibDirs []string, extraEnvs map[string]string) []ml.DeviceInfo {
var out io.Writer
if envconfig.LogLevel() == logutil.LevelTrace {
out = os.Stderr
func bootstrapDevicesWithMetalRetry(firstAttemptCtx, retryParentCtx context.Context, timeout time.Duration, ollamaLibDirs []string, extraEnvs map[string]string) []ml.DeviceInfo {
runDiscovery := func(ctx context.Context, extraEnvs map[string]string) ([]ml.DeviceInfo, *llm.StatusWriter, int, error) {
start := time.Now()
defer func() {
slog.Debug("bootstrap discovery took", "duration", time.Since(start), "OLLAMA_LIBRARY_PATH", ollamaLibDirs, "extra_envs", extraEnvs)
}()
return bootstrapDevicesWithStatus(ctx, ollamaLibDirs, extraEnvs)
}
start := time.Now()
defer func() {
slog.Debug("bootstrap discovery took", "duration", time.Since(start), "OLLAMA_LIBRARY_PATH", ollamaLibDirs, "extra_envs", extraEnvs)
}()
logutil.Trace("starting runner for device discovery", "libDirs", ollamaLibDirs, "extraEnvs", extraEnvs)
devices, status, exitCode, err := runDiscovery(firstAttemptCtx, extraEnvs)
if err == nil {
recordPersistentRunnerEnv(devices, extraEnvs)
}
if err != nil && llm.ShouldRetryWithMetalTensorDisabled(err, status) && (extraEnvs == nil || extraEnvs["GGML_METAL_TENSOR_DISABLE"] != "1") {
retryEnvs := map[string]string{}
for k, v := range extraEnvs {
retryEnvs[k] = v
}
retryEnvs["GGML_METAL_TENSOR_DISABLE"] = "1"
slog.Warn("retrying GPU discovery with Metal tensor API disabled", "error", err)
retryCtx, cancel := context.WithTimeout(retryParentCtx, timeout)
defer cancel()
devices, status, exitCode, err = runDiscovery(retryCtx, retryEnvs)
if err == nil {
recordPersistentRunnerEnv(devices, retryEnvs)
}
}
if err != nil {
if exitCode >= 0 {
// Expected during bootstrapping while we filter out unsupported GPUs.
logutil.Trace("runner exited", "OLLAMA_LIBRARY_PATH", ollamaLibDirs, "extra_envs", extraEnvs, "code", exitCode, "detail", status.LastError())
} else {
slog.Info("failure during GPU discovery", "OLLAMA_LIBRARY_PATH", ollamaLibDirs, "extra_envs", extraEnvs, "error", err, "detail", status.LastError())
}
}
return devices
}
func recordPersistentRunnerEnv(devices []ml.DeviceInfo, extraEnvs map[string]string) {
if extraEnvs["GGML_METAL_TENSOR_DISABLE"] != "1" {
return
}
for i := range devices {
if devices[i].Library != "Metal" {
continue
}
if devices[i].RunnerEnvOverrides == nil {
devices[i].RunnerEnvOverrides = map[string]string{}
}
devices[i].RunnerEnvOverrides["GGML_METAL_TENSOR_DISABLE"] = "1"
}
}
func bootstrapDevices(ctx context.Context, ollamaLibDirs []string, extraEnvs map[string]string) []ml.DeviceInfo {
devices, _, _, _ := bootstrapDevicesWithStatus(ctx, ollamaLibDirs, extraEnvs)
return devices
}
func bootstrapDevicesWithStatus(ctx context.Context, ollamaLibDirs []string, extraEnvs map[string]string) ([]ml.DeviceInfo, *llm.StatusWriter, int, error) {
var baseOut io.Writer = io.Discard
if envconfig.LogLevel() == logutil.LevelTrace {
baseOut = os.Stderr
}
status := llm.NewStatusWriter(baseOut)
cmd, port, err := llm.StartRunner(
true, // ollama engine
"", // no model
ollamaLibDirs,
out,
status,
extraEnvs,
)
if err != nil {
slog.Debug("failed to start runner to discovery GPUs", "error", err)
return nil
return nil, status, -1, err
}
go func() {
@@ -455,18 +511,13 @@ func bootstrapDevices(ctx context.Context, ollamaLibDirs []string, extraEnvs map
}()
defer cmd.Process.Kill()
devices, err := ml.GetDevicesFromRunner(ctx, &bootstrapRunner{port: port, cmd: cmd})
if err != nil {
if cmd.ProcessState != nil && cmd.ProcessState.ExitCode() >= 0 {
// Expected during bootstrapping while we filter out unsupported AMD GPUs
logutil.Trace("runner exited", "OLLAMA_LIBRARY_PATH", ollamaLibDirs, "extra_envs", extraEnvs, "code", cmd.ProcessState.ExitCode())
} else {
slog.Info("failure during GPU discovery", "OLLAMA_LIBRARY_PATH", ollamaLibDirs, "extra_envs", extraEnvs, "error", err)
}
}
logutil.Trace("runner enumerated devices", "OLLAMA_LIBRARY_PATH", ollamaLibDirs, "devices", devices)
return devices
devices, err := ml.GetDevicesFromRunner(ctx, &bootstrapRunner{port: port, cmd: cmd})
exitCode := -1
if cmd.ProcessState != nil {
exitCode = cmd.ProcessState.ExitCode()
}
return devices, status, exitCode, err
}
func overrideWarnings() {
+26
View File
@@ -4,6 +4,8 @@ import (
"log/slog"
"os"
"testing"
"github.com/ollama/ollama/ml"
)
func init() {
@@ -107,3 +109,27 @@ func TestFilterOverlapByLibrary(t *testing.T) {
})
}
}
func TestRecordPersistentRunnerEnv(t *testing.T) {
devices := []ml.DeviceInfo{
{DeviceID: ml.DeviceID{Library: "Metal", ID: "0"}},
{DeviceID: ml.DeviceID{Library: "CUDA", ID: "1"}},
}
recordPersistentRunnerEnv(devices, map[string]string{
"GGML_METAL_TENSOR_DISABLE": "1",
"CUDA_VISIBLE_DEVICES": "1",
})
if got := devices[0].RunnerEnvOverrides["GGML_METAL_TENSOR_DISABLE"]; got != "1" {
t.Fatalf("Metal RunnerEnvOverrides = %q, want %q", got, "1")
}
if _, ok := devices[0].RunnerEnvOverrides["CUDA_VISIBLE_DEVICES"]; ok {
t.Fatal("unexpected CUDA_VISIBLE_DEVICES in Metal RunnerEnvOverrides")
}
if devices[1].RunnerEnvOverrides != nil {
t.Fatalf("unexpected RunnerEnvOverrides recorded for non-Metal device: %#v", devices[1].RunnerEnvOverrides)
}
}
+4
View File
@@ -2,6 +2,10 @@
title: Structured Outputs
---
<Note>
Ollama's Cloud currently does not support structured outputs.
</Note>
Structured outputs let you enforce a JSON schema on model responses so you can reliably extract structured data, describe images, or keep every reply consistent.
## Generating structured JSON
+9 -2
View File
@@ -75,6 +75,10 @@
{
"source": "/integrations/clawdbot",
"destination": "/integrations/openclaw"
},
{
"source": "/integrations/poolside",
"destination": "/integrations/pool"
}
],
"navigation": {
@@ -111,7 +115,8 @@
"expanded": true,
"pages": [
"/integrations/openclaw",
"/integrations/hermes"
"/integrations/hermes",
"/integrations/claude-desktop"
]
},
{
@@ -120,10 +125,12 @@
"pages": [
"/integrations/claude-code",
"/integrations/codex",
"/integrations/copilot-cli",
"/integrations/opencode",
"/integrations/droid",
"/integrations/goose",
"/integrations/pi"
"/integrations/pi",
"/integrations/pool"
]
},
{
Binary file not shown.

After

Width:  |  Height:  |  Size: 307 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 297 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.4 MiB

+78
View File
@@ -0,0 +1,78 @@
---
title: Claude Desktop
---
Claude Desktop can use Ollama Cloud, including Claude Cowork and Claude Code inside the app.
<img
src="/images/claude-cowork-kimi-k2-6.png"
alt="Claude Cowork using kimi-k2.6 through Ollama Cloud"
className="rounded-xl"
/>
## Requirements
- Claude Desktop for macOS or Windows
- An [Ollama API key](https://ollama.com/settings/keys)
Set the key in your shell before launching:
```shell
export OLLAMA_API_KEY=your_api_key
```
## Quick setup
```shell
ollama launch claude-desktop
```
To bring back the usual Anthropic Claude profile later, run:
```shell
ollama launch claude-desktop --restore
```
## Using Ollama Cloud models
After setup, Claude Desktop discovers your available Ollama Cloud models automatically. For example, `kimi-k2.6` appears inside Claude Cowork.
The same models are also available to Claude Code inside Claude Desktop:
![Claude Code using kimi-k2.6 inside Claude Desktop](/images/claude-code-kimi-k2-6.png)
## Configure without launching
```shell
ollama launch claude-desktop --config
```
## Restore normal Claude
Switch Claude Desktop back to its usual profile:
```shell
ollama launch claude-desktop --restore
```
If Claude Desktop is running, use `--yes` to approve the restart automatically:
```shell
ollama launch claude-desktop --restore --yes
```
## Supported with Ollama
Claude Desktop with Ollama currently supports:
- Ollama Cloud as the third-party inference gateway
- Automatic model discovery from Ollama Cloud
- Claude Cowork with Ollama Cloud models
- Claude Code inside Claude Desktop with the same cloud models
- subagents (tell Claude to have subagents inherit the current model)
Claude Desktop with Ollama does not support yet:
- Web search
- Extensions
+93
View File
@@ -0,0 +1,93 @@
---
title: Copilot CLI
---
GitHub Copilot CLI is GitHub's AI coding agent for the terminal. It can understand your codebase, make edits, run commands, and help you build software faster.
Open models can be used with Copilot CLI through Ollama, enabling you to use models such as `qwen3.5`, `glm-5.1:cloud`, `kimi-k2.5:cloud`.
## Install
Install [Copilot CLI](https://github.com/features/copilot/cli/):
<CodeGroup>
```shell macOS / Linux (Homebrew)
brew install copilot-cli
```
```shell npm (all platforms)
npm install -g @github/copilot
```
```shell macOS / Linux (script)
curl -fsSL https://gh.io/copilot-install | bash
```
```powershell Windows (WinGet)
winget install GitHub.Copilot
```
</CodeGroup>
## Usage with Ollama
### Quick setup
```shell
ollama launch copilot
```
### Run directly with a model
```shell
ollama launch copilot --model kimi-k2.5:cloud
```
## Recommended Models
- `kimi-k2.5:cloud`
- `glm-5:cloud`
- `minimax-m2.7:cloud`
- `qwen3.5:cloud`
- `glm-4.7-flash`
- `qwen3.5`
Cloud models are also available at [ollama.com/search?c=cloud](https://ollama.com/search?c=cloud).
## Non-interactive (headless) mode
Run Copilot CLI without interaction for use in Docker, CI/CD, or scripts:
```shell
ollama launch copilot --model kimi-k2.5:cloud --yes -- -p "how does this repository work?"
```
The `--yes` flag auto-pulls the model, skips selectors, and requires `--model` to be specified. Arguments after `--` are passed directly to Copilot CLI.
## Manual setup
Copilot CLI connects to Ollama using the OpenAI-compatible API via environment variables.
1. Set the environment variables:
```shell
export COPILOT_PROVIDER_BASE_URL=http://localhost:11434/v1
export COPILOT_PROVIDER_API_KEY=
export COPILOT_PROVIDER_WIRE_API=responses
export COPILOT_MODEL=qwen3.5
```
1. Run Copilot CLI:
```shell
copilot
```
Or run with environment variables inline:
```shell
COPILOT_PROVIDER_BASE_URL=http://localhost:11434/v1 COPILOT_PROVIDER_API_KEY= COPILOT_PROVIDER_WIRE_API=responses COPILOT_MODEL=glm-5:cloud copilot
```
**Note:** Copilot requires a large context window. We recommend at least 64k tokens. See the [context length documentation](/context-length) for how to adjust context length in Ollama.
+42 -38
View File
@@ -2,7 +2,9 @@
title: Hermes Agent
---
Hermes Agent is a self-improving AI agent built by Nous Research. It features automatic skill creation, cross-session memory, and connects messaging platforms (Telegram, Discord, Slack, WhatsApp, Signal, Email) to models through a unified gateway.
Hermes Agent is a self-improving AI agent built by Nous Research. It features automatic skill creation, cross-session memory, and 70+ skills that it ships with by default.
![Hermes Agent with Ollama](/images/hermes.png)
## Quick start
@@ -10,25 +12,56 @@ Hermes Agent is a self-improving AI agent built by Nous Research. It features au
ollama launch hermes
```
### Pull a model
Ollama handles everything automatically:
Before running the setup wizard, make sure you have a model available. Hermes will auto-detect models downloaded through Ollama.
1. **Install** — If Hermes isn't installed, Ollama prompts to install it via the Nous Research install script
2. **Model** — Pick a model from the selector (local or cloud)
3. **Onboarding** — Ollama configures the Ollama provider, points Hermes at `http://127.0.0.1:11434/v1`, and sets your model as the primary
4. **Gateway** — Optionally connects a messaging platform (Telegram, Discord, Slack, WhatsApp, Signal, Email) and launches the Hermes chat
<Note>Hermes on Windows requires WSL2. Install it with `wsl --install` and re-run from inside the WSL shell.</Note>
## Recommended models
**Cloud models**:
- `kimi-k2.5:cloud` — Multimodal reasoning with subagents
- `glm-5.1:cloud` — Reasoning and code generation
- `qwen3.5:cloud` — Reasoning, coding, and agentic tool use with vision
- `minimax-m2.7:cloud` — Fast, efficient coding and real-world productivity
**Local models:**
- `gemma4` — Reasoning and code generation locally (~16 GB VRAM)
- `qwen3.6` — Reasoning, coding, and visual understanding locally (~24 GB VRAM)
More models at [ollama.com/search](https://ollama.com/search?c=cloud).
## Connect messaging apps
Link Telegram, Discord, Slack, WhatsApp, Signal, or Email to chat with your models from anywhere:
```bash
ollama pull kimi-k2.5:cloud
hermes gateway setup
```
See [Recommended models](#recommended-models) for more options.
## Reconfigure
### Install
Re-run the full setup wizard at any time:
```bash
hermes setup
```
## Manual setup
If you'd rather drive Hermes's own wizard instead of `ollama launch hermes`, install it directly:
```bash
curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent/main/scripts/install.sh | bash
```
### Set up
After installation, Hermes launches the setup wizard automatically. Choose **Quick setup**:
Hermes launches the setup wizard automatically. Choose **Quick setup**:
```
How would you like to set up Hermes?
@@ -84,32 +117,3 @@ Connect a messaging platform? (Telegram, Discord, etc.)
Launch hermes chat now? [Y/n]: Y
```
## Recommended models
**Cloud models**:
- `kimi-k2.5:cloud` — Multimodal reasoning with subagents
- `qwen3.5:cloud` — Reasoning, coding, and agentic tool use with vision
- `glm-5.1:cloud` — Reasoning and code generation
- `minimax-m2.7:cloud` — Fast, efficient coding and real-world productivity
**Local models:**
- `gemma4` — Reasoning and code generation locally (~16 GB VRAM)
- `qwen3.5` — Reasoning, coding, and visual understanding locally (~11 GB VRAM)
More models at [ollama.com/search](https://ollama.com/models).
## Configure later
Re-run the setup wizard at any time:
```bash
hermes setup
```
To configure just messaging:
```bash
hermes setup gateway
```
+3
View File
@@ -10,10 +10,12 @@ Coding assistants that can read, modify, and execute code in your projects.
- [Claude Code](/integrations/claude-code)
- [Codex](/integrations/codex)
- [Copilot CLI](/integrations/copilot-cli)
- [OpenCode](/integrations/opencode)
- [Droid](/integrations/droid)
- [Goose](/integrations/goose)
- [Pi](/integrations/pi)
- [Pool](/integrations/pool)
## Assistants
@@ -21,6 +23,7 @@ AI assistants that help with everyday tasks.
- [OpenClaw](/integrations/openclaw)
- [Hermes Agent](/integrations/hermes)
- [Claude Desktop](/integrations/claude-desktop)
## IDEs & Editors
+5 -6
View File
@@ -15,7 +15,7 @@ Ollama handles everything automatically:
1. **Install** — If OpenClaw isn't installed, Ollama prompts to install it via npm
2. **Security** — On the first launch, a security notice explains the risks of tool access
3. **Model** — Pick a model from the selector (local or cloud)
4. **Onboarding** — Ollama configures the provider, installs the gateway daemon, sets your model as the primary, and installs the web search and fetch plugin
4. **Onboarding** — Ollama configures the provider, installs the gateway daemon, sets your model as the primary, and enables OpenClaw's bundled Ollama web search
5. **Gateway** — Starts in the background and opens the OpenClaw TUI
<Note>OpenClaw requires a larger context window. It is recommended to use a context window of at least 64k tokens if using local models. See [Context length](/context-length) for more information.</Note>
@@ -24,19 +24,19 @@ Ollama handles everything automatically:
## Web search and fetch
OpenClaw ships with a web search and fetch plugin that gives local or cloud models the ability to search the web and extract readable page content.
OpenClaw ships with a bundled Ollama `web_search` provider that lets local or cloud-backed Ollama setups search the web through the configured Ollama host.
```bash
ollama launch openclaw
```
Web search and fetch is enabled automatically when launching OpenClaw through Ollama. To install the plugin directly:
Ollama web search is enabled automatically when launching OpenClaw through Ollama. To configure it manually:
```bash
openclaw plugins install @ollama/openclaw-web-search
openclaw configure --section web
```
<Note>Web search for local models requires `ollama signin`.</Note>
<Note>Ollama web search for local models requires `ollama signin`.</Note>
## Configure without launching
@@ -93,4 +93,3 @@ Link WhatsApp, Telegram, Slack, Discord, or iMessage to chat with your local mod
```bash
openclaw gateway stop
```
+54
View File
@@ -0,0 +1,54 @@
---
title: Pool
---
Pool is Poolside's software agent for the terminal, built for enterprise development workflows.
## Install
Install [Pool](https://github.com/poolsideai/pool):
## Usage with Ollama
### Quick setup
```shell
ollama launch pool
```
### Run directly with a model
```shell
ollama launch pool --model kimi-k2.6:cloud
```
### Pass arguments through to Pool
Arguments after `--` are passed directly to Pool:
```shell
ollama launch pool -- --help
```
## Manual setup
Pool connects to Ollama using the OpenAI-compatible API via environment variables.
1. Set the environment variables:
```shell
export POOLSIDE_STANDALONE_BASE_URL=http://localhost:11434/v1
export POOLSIDE_API_KEY=ollama
```
2. Run Pool with an Ollama model:
```shell
pool -m kimi-k2.6:cloud
```
Or run with environment variables inline:
```shell
POOLSIDE_STANDALONE_BASE_URL=http://localhost:11434/v1 POOLSIDE_API_KEY=ollama pool -m kimi-k2.6:cloud
```
+3 -2
View File
@@ -283,10 +283,11 @@ func (kv KV) OllamaEngineRequired() bool {
"gemma3n",
"gemma4",
"gptoss", "gpt-oss",
"laguna",
"llama4",
"mistral3",
"mllama",
"nemotron_h", "nemotron_h_moe",
"nemotron_h", "nemotron_h_moe", "nemotron_h_omni",
"nomic-bert",
"olmo3",
"qwen25vl",
@@ -897,7 +898,7 @@ func (f GGML) FlashAttention() bool {
"lfm2",
"lfm2moe",
"mistral3",
"nemotron_h", "nemotron_h_moe",
"nemotron_h", "nemotron_h_moe", "nemotron_h_omni",
"olmo3",
"qwen3", "qwen3moe",
"qwen35", "qwen35moe",
-8
View File
@@ -406,10 +406,6 @@ func TestAPIShowModel(t *testing.T) {
}
func TestAPIGenerateLogprobs(t *testing.T) {
if testModel != "" {
// Logprobs requires runner support (e.g. llama.cpp has it, MLX does not).
t.Skip("logprobs not supported by all runners")
}
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer cancel()
@@ -523,10 +519,6 @@ func TestAPIGenerateLogprobs(t *testing.T) {
}
func TestAPIChatLogprobs(t *testing.T) {
if testModel != "" {
// Logprobs requires runner support (e.g. llama.cpp has it, MLX does not).
t.Skip("logprobs not supported by all runners")
}
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer cancel()
+61 -31
View File
@@ -278,7 +278,7 @@ func NewLlamaServer(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, modelPath st
modelPath,
gpuLibs,
status,
ml.GetVisibleDevicesEnv(gpus, false),
ml.GetDevicesEnv(gpus, false),
)
s := llmServer{
@@ -298,8 +298,8 @@ func NewLlamaServer(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, modelPath st
if err != nil {
var msg string
if s.status != nil && s.status.LastErrMsg != "" {
msg = s.status.LastErrMsg
if s.status != nil && s.status.LastError() != "" {
msg = s.status.LastError()
}
err := fmt.Errorf("error starting runner: %v %s", err, msg)
if llamaModel != nil {
@@ -312,12 +312,12 @@ func NewLlamaServer(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, modelPath st
go func() {
err := s.cmd.Wait()
// Favor a more detailed message over the process exit status
if err != nil && s.status != nil && s.status.LastErrMsg != "" {
if err != nil && s.status != nil && s.status.LastError() != "" {
slog.Error("llama runner terminated", "error", err)
if strings.Contains(s.status.LastErrMsg, "unknown model") {
s.status.LastErrMsg = "this model is not supported by your version of Ollama. You may need to upgrade"
if strings.Contains(s.status.LastError(), "unknown model") {
s.status.SetLastError("this model is not supported by your version of Ollama. You may need to upgrade")
}
s.doneErr = errors.New(s.status.LastErrMsg)
s.doneErr = errors.New(s.status.LastError())
} else {
s.doneErr = err
}
@@ -385,20 +385,9 @@ func StartRunner(ollamaEngine bool, modelPath string, gpuLibs []string, out io.W
cmd.Env = os.Environ()
if out != nil {
stdout, err := cmd.StdoutPipe()
if err != nil {
return nil, 0, fmt.Errorf("failed to spawn server stdout pipe: %w", err)
}
stderr, err := cmd.StderrPipe()
if err != nil {
return nil, 0, fmt.Errorf("failed to spawn server stderr pipe: %w", err)
}
go func() {
io.Copy(out, stdout) //nolint:errcheck
}()
go func() {
io.Copy(out, stderr) //nolint:errcheck
}()
// os/exec serializes Write calls when shared
cmd.Stdout = out
cmd.Stderr = out
}
cmd.SysProcAttr = LlamaServerSysProcAttr
@@ -451,6 +440,38 @@ func StartRunner(ollamaEngine bool, modelPath string, gpuLibs []string, out io.W
return
}
// Workaround possible runtime crash where the probe incorrectly
// enables metal tensor, but fails at runtime
func ShouldRetryWithMetalTensorDisabled(err error, status *StatusWriter) bool {
if runtime.GOOS != "darwin" {
return false
}
var msg strings.Builder
msg.WriteString(strings.ToLower(err.Error()))
if status != nil && status.LastError() != "" {
msg.WriteByte(' ')
msg.WriteString(strings.ToLower(status.LastError()))
}
text := msg.String()
for _, needle := range []string{
"failed to initialize ggml backend device: metal",
"failed to initialize metal backend",
"failed to initialize the metal library",
"failed to allocate context",
"unable to create llama context",
"signal arrived during cgo execution",
"input types must match cooperative tensor types",
} {
if strings.Contains(text, needle) {
return true
}
}
return false
}
func (s *llmServer) ModelPath() string {
return s.modelPath
}
@@ -723,7 +744,7 @@ func (s *llamaServer) Load(ctx context.Context, systemInfo ml.SystemInfo, system
// The llama engine does its memory allocations together with model loading, so we
// need to wait until it is done to ensure that we have accurate memory data before
// loading the next model
// loading the next model.
return uniqueDeviceIDs(s.loadRequest.GPULayers), s.WaitUntilRunning(ctx)
}
@@ -1276,8 +1297,8 @@ func (s *llmServer) getServerStatus(ctx context.Context) (ServerStatus, error) {
// Fail fast if its exited
if s.cmd.ProcessState != nil {
msg := ""
if s.status != nil && s.status.LastErrMsg != "" {
msg = s.status.LastErrMsg
if s.status != nil && s.status.LastError() != "" {
msg = s.status.LastError()
}
if s.cmd.ProcessState.ExitCode() == -1 {
// Most likely a signal killed it, log some more details to try to help troubleshoot
@@ -1371,21 +1392,30 @@ func (s *llmServer) WaitUntilRunning(ctx context.Context) error {
slog.Warn("client connection closed before server finished loading, aborting load")
return fmt.Errorf("timed out waiting for llama runner to start: %w", ctx.Err())
case <-s.done:
return fmt.Errorf("llama runner process has terminated: %w", s.doneErr)
if s.status != nil && s.status.LastError() != "" {
return fmt.Errorf("llama runner process has terminated: %s", s.status.LastError())
}
if s.doneErr != nil {
return fmt.Errorf("llama runner process has terminated: %w", s.doneErr)
}
if s.cmd != nil && s.cmd.ProcessState != nil {
return fmt.Errorf("llama runner process has terminated with exit code %d", s.cmd.ProcessState.ExitCode())
}
return errors.New("llama runner process has terminated")
default:
}
if time.Now().After(stallTimer) {
// timeout
msg := ""
if s.status != nil && s.status.LastErrMsg != "" {
msg = s.status.LastErrMsg
if s.status != nil && s.status.LastError() != "" {
msg = s.status.LastError()
}
return fmt.Errorf("timed out waiting for llama runner to start - progress %0.2f - %s", s.loadProgress, msg)
}
if s.cmd.ProcessState != nil {
msg := ""
if s.status != nil && s.status.LastErrMsg != "" {
msg = s.status.LastErrMsg
if s.status != nil && s.status.LastError() != "" {
msg = s.status.LastError()
}
return fmt.Errorf("llama runner process no longer running: %d %s", s.cmd.ProcessState.ExitCode(), msg)
}
@@ -1695,8 +1725,8 @@ func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn fu
if strings.Contains(err.Error(), "unexpected EOF") || strings.Contains(err.Error(), "forcibly closed") {
s.Close()
var msg string
if s.status != nil && s.status.LastErrMsg != "" {
msg = s.status.LastErrMsg
if s.status != nil && s.status.LastError() != "" {
msg = s.status.LastError()
} else {
msg = err.Error()
}
+31
View File
@@ -0,0 +1,31 @@
package llm
import (
"context"
"strings"
"testing"
)
func TestWaitUntilRunningUsesStatusMessageWhenDoneErrIsNil(t *testing.T) {
done := make(chan struct{})
close(done)
status := &StatusWriter{}
status.SetLastError("llama_init_from_model: failed to initialize the context: failed to initialize Metal backend")
s := &llmServer{
done: done,
status: status,
}
err := s.WaitUntilRunning(context.Background())
if err == nil {
t.Fatal("expected error")
}
if strings.Contains(err.Error(), "%!w(<nil>)") {
t.Fatalf("unexpected wrapped nil error: %q", err)
}
if !strings.Contains(err.Error(), s.status.LastError()) {
t.Fatalf("error %q does not include status message %q", err, s.status.LastError())
}
}
+56 -6
View File
@@ -2,23 +2,67 @@ package llm
import (
"bytes"
"os"
"io"
"strings"
"sync/atomic"
)
// StatusWriter is a writer that captures error messages from the llama runner process
type StatusWriter struct {
LastErrMsg string
out *os.File
out io.Writer
// StartRunner wires both Stdout and Stderr to the same StatusWriter, and
// os/exec serializes Write calls in that case.
lastErrMsg atomic.Value
}
func NewStatusWriter(out *os.File) *StatusWriter {
const maxCapturedErrorBytes = 8 * 1024
func NewStatusWriter(out io.Writer) *StatusWriter {
return &StatusWriter{
out: out,
}
}
func (w *StatusWriter) LastError() string {
if w == nil {
return ""
}
if v := w.lastErrMsg.Load(); v != nil {
return v.(string)
}
return ""
}
func (w *StatusWriter) SetLastError(msg string) {
if w == nil {
return
}
w.lastErrMsg.Store(msg)
}
func (w *StatusWriter) AppendError(msg string) {
if w == nil || msg == "" {
return
}
if current := w.LastError(); current != "" {
msg = current + "\n" + msg
}
if len(msg) > maxCapturedErrorBytes {
msg = msg[len(msg)-maxCapturedErrorBytes:]
if i := strings.IndexByte(msg, '\n'); i >= 0 {
msg = msg[i+1:]
}
}
w.SetLastError(msg)
}
// TODO - regex matching to detect errors like
// libcublasLt.so.11: cannot open shared object file: No such file or directory
// TODO - if we later see error lines split across multiple Write calls in real
// logs, add a small rolling buffer here to capture those fragments.
var errorPrefixes = []string{
"error:",
@@ -29,17 +73,23 @@ var errorPrefixes = []string{
"error loading model",
"GGML_ASSERT",
"Deepseek2 does not support K-shift",
"signal arrived during cgo execution",
"llama_init_from_model:",
}
func (w *StatusWriter) Write(b []byte) (int, error) {
var errMsg string
for _, prefix := range errorPrefixes {
if _, after, ok := bytes.Cut(b, []byte(prefix)); ok {
errMsg = prefix + string(bytes.TrimSpace(after))
line := after
if j := bytes.IndexByte(line, '\n'); j >= 0 {
line = line[:j]
}
errMsg = prefix + string(bytes.TrimRight(line, " \t\r"))
}
}
if errMsg != "" {
w.LastErrMsg = errMsg
w.AppendError(errMsg)
}
return w.out.Write(b)
+44
View File
@@ -0,0 +1,44 @@
package llm
import (
"os"
"testing"
)
func TestStatusWriterCapturesErrorLine(t *testing.T) {
f, err := os.CreateTemp(t.TempDir(), "status-writer")
if err != nil {
t.Fatal(err)
}
defer f.Close()
w := NewStatusWriter(f)
if _, err := w.Write([]byte("llama_init_from_model: failed to initialize the context: failed to initialize Metal backend\n")); err != nil {
t.Fatal(err)
}
if got, want := w.LastError(), "llama_init_from_model: failed to initialize the context: failed to initialize Metal backend"; got != want {
t.Fatalf("LastError = %q, want %q", got, want)
}
}
func TestStatusWriterAccumulatesErrorLines(t *testing.T) {
f, err := os.CreateTemp(t.TempDir(), "status-writer")
if err != nil {
t.Fatal(err)
}
defer f.Close()
w := NewStatusWriter(f)
if _, err := w.Write([]byte("error: failed to initialize the Metal library\n")); err != nil {
t.Fatal(err)
}
if _, err := w.Write([]byte("GGML_ASSERT([rsets->data count] == 0) failed\n")); err != nil {
t.Fatal(err)
}
want := "error: failed to initialize the Metal library\nGGML_ASSERT([rsets->data count] == 0) failed"
if got := w.LastError(); got != want {
t.Fatalf("LastError = %q, want %q", got, want)
}
}
+12
View File
@@ -6,6 +6,7 @@ import (
"fmt"
"io"
"os"
"time"
)
type Layer struct {
@@ -60,6 +61,9 @@ func NewLayer(r io.Reader, mediatype string) (Layer, error) {
return Layer{}, err
}
}
if err := touchLayer(blob); err != nil {
return Layer{}, err
}
return Layer{
MediaType: mediatype,
@@ -83,6 +87,9 @@ func NewLayerFromLayer(digest, mediatype, from string) (Layer, error) {
if err != nil {
return Layer{}, err
}
if err := touchLayer(blob); err != nil {
return Layer{}, err
}
return Layer{
MediaType: mediatype,
@@ -93,6 +100,11 @@ func NewLayerFromLayer(digest, mediatype, from string) (Layer, error) {
}, nil
}
func touchLayer(path string) error {
now := time.Now()
return os.Chtimes(path, now, now)
}
func (l *Layer) Open() (io.ReadSeekCloser, error) {
if l.Digest == "" {
return nil, errors.New("opening layer with empty digest")
+11 -3
View File
@@ -50,8 +50,18 @@ var initDevices = sync.OnceFunc(func() {
backends = make(map[C.ggml_backend_dev_t]C.ggml_backend_t)
for i := range C.ggml_backend_dev_count() {
d := C.ggml_backend_dev_get(i)
t := C.ggml_backend_dev_type(d)
name := C.GoString(C.ggml_backend_dev_name(d))
switch C.ggml_backend_dev_type(d) {
b := C.ggml_backend_dev_init(d, nil)
if b == nil {
slog.Error("failed to initialize ggml backend device", "device", name, "type", t)
panic(fmt.Sprintf("failed to initialize ggml backend device: %s", name))
}
backends[d] = b
switch t {
case C.GGML_BACKEND_DEVICE_TYPE_CPU:
if len(cpus) == 0 {
// only the first cpu device should be used
@@ -63,8 +73,6 @@ var initDevices = sync.OnceFunc(func() {
C.GGML_BACKEND_DEVICE_TYPE_IGPU:
gpus = append(gpus, d)
}
backends[d] = C.ggml_backend_dev_init(d, nil)
}
})
+16 -3
View File
@@ -314,6 +314,11 @@ type DeviceInfo struct {
// Where backends were loaded from
LibraryPath []string
// RunnerEnvOverrides stores exceptional per-device runner environment
// overrides discovered during bootstrap. This is internal server state and
// is not serialized.
RunnerEnvOverrides map[string]string `json:"-"`
}
type SystemInfo struct {
@@ -519,16 +524,24 @@ func (f FlashAttentionType) String() string {
}
// Given the list of GPUs this instantiation is targeted for,
// figure out the visible devices environment variables
// Set mustFilter true to enable filtering of CUDA devices
func GetVisibleDevicesEnv(l []DeviceInfo, mustFilter bool) map[string]string {
// figure out the device environment variables and any recorded
// per-device runner environment overrides. Set mustFilter true to enable
// filtering of CUDA devices.
func GetDevicesEnv(l []DeviceInfo, mustFilter bool) map[string]string {
if len(l) == 0 {
return nil
}
env := map[string]string{}
for _, d := range l {
d.updateVisibleDevicesEnv(env, mustFilter)
for k, v := range d.RunnerEnvOverrides {
if existing, ok := env[k]; ok && existing != v {
slog.Warn("conflicting device environment override", "key", k, "existing", existing, "new", v, "library", d.Library, "id", d.ID)
}
env[k] = v
}
}
return env
}
+60
View File
@@ -0,0 +1,60 @@
package ml
import (
"bytes"
"log/slog"
"strings"
"testing"
)
func TestMergeEnvWithRunnerEnvOverrides(t *testing.T) {
devices := []DeviceInfo{
{
DeviceID: DeviceID{Library: "Metal", ID: "0"},
RunnerEnvOverrides: map[string]string{"GGML_METAL_TENSOR_DISABLE": "1"},
},
{
DeviceID: DeviceID{Library: "CUDA", ID: "3"},
},
}
env := GetDevicesEnv(devices, true)
if got, want := env["GGML_METAL_TENSOR_DISABLE"], "1"; got != want {
t.Fatalf("GGML_METAL_TENSOR_DISABLE = %q, want %q", got, want)
}
if got, want := env["CUDA_VISIBLE_DEVICES"], "3"; got != want {
t.Fatalf("CUDA_VISIBLE_DEVICES = %q, want %q", got, want)
}
}
func TestGetDevicesEnvWarnsOnConflictingOverrides(t *testing.T) {
var logs bytes.Buffer
oldLogger := slog.Default()
slog.SetDefault(slog.New(slog.NewTextHandler(&logs, &slog.HandlerOptions{Level: slog.LevelDebug})))
t.Cleanup(func() {
slog.SetDefault(oldLogger)
})
devices := []DeviceInfo{
{
DeviceID: DeviceID{Library: "Metal", ID: "0"},
RunnerEnvOverrides: map[string]string{"TEST_OVERRIDE": "one"},
},
{
DeviceID: DeviceID{Library: "Metal", ID: "1"},
RunnerEnvOverrides: map[string]string{"TEST_OVERRIDE": "two"},
},
}
env := GetDevicesEnv(devices, false)
if got, want := env["TEST_OVERRIDE"], "two"; got != want {
t.Fatalf("TEST_OVERRIDE = %q, want %q", got, want)
}
if !strings.Contains(logs.String(), "conflicting device environment override") {
t.Fatalf("expected warning log, got %q", logs.String())
}
}
+444
View File
@@ -0,0 +1,444 @@
package laguna
import (
"fmt"
"math"
"github.com/ollama/ollama/fs"
"github.com/ollama/ollama/kvcache"
"github.com/ollama/ollama/ml"
"github.com/ollama/ollama/ml/nn"
"github.com/ollama/ollama/ml/nn/rope"
"github.com/ollama/ollama/model"
"github.com/ollama/ollama/model/input"
"github.com/ollama/ollama/tokenizer"
)
const (
cacheTypeSWA = iota
cacheTypeCausal
)
type Options struct {
hiddenSize int
headDim int
numHeads []int
numKVHeads int
eps float32
slidingWindow int
slidingWindowPattern []bool
fullRopeDim int
fullRopeBase, fullRopeScale float32
fullRopeOriginalContextLength int
fullRopeAttentionFactor float32
fullRopeBetaFast float32
fullRopeBetaSlow float32
swaRopeDim int
swaRopeBase, swaRopeScale float32
numExperts, numExpertsUsed int
normTopKProb bool
routedScalingFactor float32
decoderSparseStep int
denseLayers map[int]bool
}
func (o *Options) numHeadsForLayer(layer int) int {
if layer < len(o.numHeads) && o.numHeads[layer] > 0 {
return o.numHeads[layer]
}
if len(o.numHeads) > 0 && o.numHeads[0] > 0 {
return o.numHeads[0]
}
return 1
}
func (o *Options) layerIsSliding(layer int) bool {
return layer < len(o.slidingWindowPattern) && o.slidingWindowPattern[layer]
}
func (o *Options) layerUsesMoE(layer int) bool {
if o.numExperts == 0 || o.denseLayers[layer] {
return false
}
step := o.decoderSparseStep
if step <= 0 {
step = 1
}
return (layer+1)%step == 0
}
func (o *Options) applyRotaryPositionEmbeddings(ctx ml.Context, layer int, states, positions ml.Tensor) ml.Tensor {
opts := []func(*rope.Options){rope.WithTypeNeoX()}
if o.layerIsSliding(layer) {
return nn.RoPE(ctx, states, positions, o.swaRopeDim, o.swaRopeBase, 1./o.swaRopeScale, opts...)
}
opts = append(opts,
rope.WithOriginalContextLength(o.fullRopeOriginalContextLength),
rope.WithExtrapolationFactor(1),
rope.WithAttentionFactor(o.fullRopeAttentionFactor),
rope.WithBetaFast(o.fullRopeBetaFast),
rope.WithBetaSlow(o.fullRopeBetaSlow),
)
return nn.RoPE(ctx, states, positions, o.fullRopeDim, o.fullRopeBase, 1./o.fullRopeScale, opts...)
}
type Attention struct {
Query *nn.Linear `gguf:"attn_q"`
QueryNorm *nn.RMSNorm `gguf:"attn_q_norm"`
Key *nn.Linear `gguf:"attn_k"`
KeyNorm *nn.RMSNorm `gguf:"attn_k_norm"`
Value *nn.Linear `gguf:"attn_v"`
Gate *nn.Linear `gguf:"attn_g"`
Output *nn.Linear `gguf:"attn_output"`
}
func (sa *Attention) Forward(ctx ml.Context, layer int, hiddenStates, positions ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor {
batchSize := hiddenStates.Dim(1)
numHeads := opts.numHeadsForLayer(layer)
query := sa.Query.Forward(ctx, hiddenStates)
key := sa.Key.Forward(ctx, hiddenStates)
value := sa.Value.Forward(ctx, hiddenStates)
gate := sa.Gate.Forward(ctx, hiddenStates)
query = query.Reshape(ctx, opts.headDim, numHeads, batchSize)
key = key.Reshape(ctx, opts.headDim, opts.numKVHeads, batchSize)
value = value.Reshape(ctx, opts.headDim, opts.numKVHeads, batchSize)
query = sa.QueryNorm.Forward(ctx, query, opts.eps)
key = sa.KeyNorm.Forward(ctx, key, opts.eps)
query = opts.applyRotaryPositionEmbeddings(ctx, layer, query, positions)
key = opts.applyRotaryPositionEmbeddings(ctx, layer, key, positions)
attention := nn.Attention(ctx, query, key, value, 1./math.Sqrt(float64(opts.headDim)), cache)
// Laguna applies the per-head gate softplus in float32, then casts back.
gate = gate.Cast(ctx, ml.DTypeF32).Softplus(ctx).Cast(ctx, attention.DType())
attention = attention.Mul(ctx, gate.Reshape(ctx, 1, numHeads, batchSize))
attention = attention.Reshape(ctx, opts.headDim*numHeads, batchSize)
return sa.Output.Forward(ctx, attention)
}
type MLP interface {
Forward(ml.Context, ml.Tensor, *Options) ml.Tensor
}
type dense struct {
Gate *nn.Linear `gguf:"ffn_gate"`
Up *nn.Linear `gguf:"ffn_up"`
Down *nn.Linear `gguf:"ffn_down"`
}
func (mlp *dense) Forward(ctx ml.Context, hiddenStates ml.Tensor, _ *Options) ml.Tensor {
hiddenStates = mlp.Gate.Forward(ctx, hiddenStates).SILU(ctx, mlp.Up.Forward(ctx, hiddenStates))
return mlp.Down.Forward(ctx, hiddenStates)
}
type sparse struct {
Router *nn.Linear `gguf:"ffn_gate_inp"`
Gate *nn.LinearBatch `gguf:"ffn_gate_exps"`
Up *nn.LinearBatch `gguf:"ffn_up_exps"`
Down *nn.LinearBatch `gguf:"ffn_down_exps"`
SharedExpert *dense `gguf:",suf:_shexp"`
ExpProbsBias ml.Tensor `gguf:"exp_probs_b.bias,alt:exp_probs_b"`
}
func (moe *sparse) topKIndices(ctx ml.Context, scores ml.Tensor, opts *Options) ml.Tensor {
if moe.ExpProbsBias != nil {
scores = scores.Add(ctx, moe.ExpProbsBias)
}
return scores.TopK(ctx, opts.numExpertsUsed)
}
func (moe *sparse) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *Options) ml.Tensor {
residual := hiddenStates
scores := moe.Router.Forward(ctx, hiddenStates).Cast(ctx, ml.DTypeF32).Sigmoid(ctx)
selectedExperts := moe.topKIndices(ctx, scores, opts)
routingWeights := scores.Reshape(ctx, 1, opts.numExperts, hiddenStates.Dim(1)).Rows(ctx, selectedExperts)
if opts.normTopKProb {
routingWeights = routingWeights.Reshape(ctx, opts.numExpertsUsed, hiddenStates.Dim(1))
routingWeights = routingWeights.Div(ctx, routingWeights.SumRows(ctx))
routingWeights = routingWeights.Reshape(ctx, 1, opts.numExpertsUsed, hiddenStates.Dim(1))
}
routingWeights = routingWeights.Scale(ctx, float64(opts.routedScalingFactor))
hiddenStates = hiddenStates.Reshape(ctx, hiddenStates.Dim(0), 1, hiddenStates.Dim(1))
upStates := moe.Up.Forward(ctx, hiddenStates, selectedExperts)
hiddenStates = moe.Gate.Forward(ctx, hiddenStates, selectedExperts).SILU(ctx, upStates)
experts := moe.Down.Forward(ctx, hiddenStates, selectedExperts)
experts = experts.Mul(ctx, routingWeights)
nextStates := experts.View(ctx, 0, experts.Dim(0), experts.Stride(2), experts.Dim(2))
for i := 1; i < opts.numExpertsUsed; i++ {
nextStates = nextStates.Add(ctx, experts.View(ctx, i*experts.Stride(1), experts.Dim(0), experts.Stride(2), experts.Dim(2)))
}
return nextStates.Add(ctx, moe.SharedExpert.Forward(ctx, residual, opts))
}
type Layer struct {
AttentionNorm *nn.RMSNorm `gguf:"attn_norm"`
*Attention
MLPNorm *nn.RMSNorm `gguf:"ffn_norm"`
MLP
}
func (l *Layer) Forward(ctx ml.Context, layer int, hiddenStates, positions, outputs ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor {
residual := hiddenStates
hiddenStates = l.AttentionNorm.Forward(ctx, hiddenStates, opts.eps)
hiddenStates = l.Attention.Forward(ctx, layer, hiddenStates, positions, cache, opts)
if outputs != nil {
hiddenStates = hiddenStates.Rows(ctx, outputs)
residual = residual.Rows(ctx, outputs)
}
hiddenStates = hiddenStates.Add(ctx, residual)
residual = hiddenStates
hiddenStates = l.MLPNorm.Forward(ctx, hiddenStates, opts.eps)
hiddenStates = l.MLP.Forward(ctx, hiddenStates, opts)
return hiddenStates.Add(ctx, residual)
}
type Model struct {
model.Base
tokenizer.Tokenizer
TokenEmbedding *nn.Embedding `gguf:"token_embd"`
Layers []Layer `gguf:"blk"`
OutputNorm *nn.RMSNorm `gguf:"output_norm"`
Output *nn.Linear `gguf:"output,alt:token_embd"`
*Options
}
func New(c fs.Config) (model.Model, error) {
if c.Bool("attention.sink_enabled") {
return nil, fmt.Errorf("laguna: SWA attention sinks are not supported")
}
if c.Uint("attention.gating_type") != 1 {
return nil, fmt.Errorf("laguna: unsupported attention gating type %d", c.Uint("attention.gating_type"))
}
if !c.Bool("attention.qk_norm") {
return nil, fmt.Errorf("laguna: Q/K RMSNorm is required")
}
if gating := c.Uint("expert_gating_func"); gating != 2 {
return nil, fmt.Errorf("laguna: unsupported expert gating function %d", gating)
}
numLayers := int(c.Uint("block_count"))
opts := newOptions(c, numLayers)
layers := make([]Layer, numLayers)
for i := range layers {
if opts.layerUsesMoE(i) {
layers[i].MLP = &sparse{}
} else {
layers[i].MLP = &dense{}
}
}
var pre []string
switch c.String("tokenizer.ggml.pre") {
case "laguna":
pre = []string{
`(?:\r?\n)+(?!\r?\n)`,
`(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
}
default:
return nil, model.ErrUnsupportedTokenizer
}
m := Model{
Tokenizer: tokenizer.NewBytePairEncoding(
&tokenizer.Vocabulary{
Values: c.Strings("tokenizer.ggml.tokens"),
Types: c.Ints("tokenizer.ggml.token_type"),
Merges: c.Strings("tokenizer.ggml.merges"),
AddBOS: c.Bool("tokenizer.ggml.add_bos_token", true),
BOS: []int32{int32(c.Uint("tokenizer.ggml.bos_token_id"))},
AddEOS: c.Bool("tokenizer.ggml.add_eos_token", false),
EOS: append(
[]int32{int32(c.Uint("tokenizer.ggml.eos_token_id"))},
c.Ints("tokenizer.ggml.eos_token_ids")...,
),
},
pre...,
),
Layers: layers,
Options: opts,
}
m.Cache = kvcache.NewWrapperCache(
kvcache.NewSWACache(int32(opts.slidingWindow), m.Shift),
kvcache.NewCausalCache(m.Shift),
)
return &m, nil
}
func newOptions(c fs.Config, numLayers int) *Options {
denseLayers := make(map[int]bool)
for _, layer := range configUints(c, "dense_layers") {
denseLayers[int(layer)] = true
}
for i := range c.Uint("leading_dense_block_count") {
denseLayers[int(i)] = true
}
fullRopeScale := c.Float("rope.scaling.factor", 1)
if fullRopeScale == 0 {
fullRopeScale = 1
}
swaRopeScale := c.Float("rope.swa.scaling.factor", 1)
if swaRopeScale == 0 {
swaRopeScale = 1
}
fullRopeType := c.String("rope.scaling.type")
fullRopeAttentionFactor := lagunaAttentionFactor(fullRopeType, fullRopeScale, c.Float("rope.scaling.attn_factor"))
return &Options{
hiddenSize: int(c.Uint("embedding_length")),
headDim: int(c.Uint("attention.key_length")),
numHeads: expandIntArray(configUints(c, "attention.head_count"), numLayers, c.Uint("attention.head_count", 1)),
numKVHeads: int(c.Uint("attention.head_count_kv")),
eps: c.Float("attention.layer_norm_rms_epsilon", 1e-6),
slidingWindow: int(c.Uint("attention.sliding_window", 512)),
slidingWindowPattern: slidingWindowPattern(c, numLayers),
fullRopeDim: int(c.Uint("rope.dimension_count", c.Uint("attention.key_length"))),
fullRopeBase: c.Float("rope.freq_base", 500000),
fullRopeScale: fullRopeScale,
fullRopeOriginalContextLength: int(c.Uint("rope.scaling.original_context_length", 4096)),
fullRopeAttentionFactor: fullRopeAttentionFactor,
fullRopeBetaFast: c.Float("rope.scaling.beta_fast", 64),
fullRopeBetaSlow: c.Float("rope.scaling.beta_slow", 1),
swaRopeDim: int(c.Uint("rope.swa.dimension_count", c.Uint("attention.key_length"))),
swaRopeBase: c.Float("rope.swa.freq_base", 10000),
swaRopeScale: swaRopeScale,
numExperts: int(c.Uint("expert_count")),
numExpertsUsed: int(c.Uint("expert_used_count")),
normTopKProb: c.Bool("expert_weights_norm", true),
routedScalingFactor: c.Float("expert_weights_scale", 1),
decoderSparseStep: int(c.Uint("decoder_sparse_step", 1)),
denseLayers: denseLayers,
}
}
func lagunaAttentionFactor(ropeType string, scaleFactor, attentionFactor float32) float32 {
if attentionFactor != 0 {
return attentionFactor
}
if ropeType == "yarn" && scaleFactor > 1 {
return float32(0.1*math.Log(float64(scaleFactor)) + 1)
}
return 1
}
func slidingWindowPattern(c fs.Config, numLayers int) []bool {
pattern := c.Bools("attention.sliding_window_pattern")
if len(pattern) == numLayers {
return pattern
}
layerTypes := configUints(c, "attention.layer_types")
if len(layerTypes) == numLayers {
pattern = make([]bool, numLayers)
for i, layerType := range layerTypes {
pattern[i] = layerType == 1
}
return pattern
}
return make([]bool, numLayers)
}
func configUints(c fs.Config, key string) []uint32 {
keyExists := c.Value(c.Architecture()+"."+key) != nil || c.Value(key) != nil
if cc, ok := c.(interface {
Uints(string, ...[]uint32) []uint32
}); ok {
if values := cc.Uints(key); len(values) > 0 && (keyExists || !(len(values) == 1 && values[0] == 0)) {
return values
}
}
ints := c.Ints(key)
if len(ints) > 0 && (keyExists || !(len(ints) == 1 && ints[0] == 0)) {
values := make([]uint32, len(ints))
for i, v := range ints {
values[i] = uint32(v)
}
return values
}
if scalar := c.Uint(key); scalar != 0 {
return []uint32{scalar}
}
return nil
}
func expandIntArray(values []uint32, n int, fallback uint32) []int {
if len(values) == 0 {
values = []uint32{fallback}
}
defaultValue := values[0]
if len(values) == 1 {
defaultValue = values[0]
}
out := make([]int, n)
for i := range out {
if i < len(values) {
out[i] = int(values[i])
} else {
out[i] = int(defaultValue)
}
}
return out
}
func (m *Model) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) {
return m.Options.applyRotaryPositionEmbeddings(ctx, layer, key, shift), nil
}
func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
positions := ctx.Input().FromInts(batch.Positions, len(batch.Positions))
hiddenStates := m.TokenEmbedding.Forward(ctx, batch.Inputs)
for i, layer := range m.Layers {
if m.Cache != nil {
m.Cache.SetLayer(i)
if wrapper, ok := m.Cache.(*kvcache.WrapperCache); ok {
cacheType := cacheTypeCausal
if m.Options.layerIsSliding(i) {
cacheType = cacheTypeSWA
}
wrapper.SetLayerType(cacheType)
}
}
var outputs ml.Tensor
if i == len(m.Layers)-1 {
outputs = batch.Outputs
}
hiddenStates = layer.Forward(ctx, i, hiddenStates, positions, outputs, m.Cache, m.Options)
}
hiddenStates = m.OutputNorm.Forward(ctx, hiddenStates, m.eps)
return m.Output.Forward(ctx, hiddenStates), nil
}
func init() {
model.Register("laguna", New)
}
var _ model.Model = (*Model)(nil)
+237
View File
@@ -0,0 +1,237 @@
package laguna
import (
"iter"
"math"
"testing"
)
type testConfig map[string]any
func (c testConfig) Architecture() string { return "laguna" }
func (c testConfig) key(key string) string {
switch {
case len(key) >= len("tokenizer.") && key[:len("tokenizer.")] == "tokenizer.":
return key
case len(key) >= len("general.") && key[:len("general.")] == "general.":
return key
default:
return "laguna." + key
}
}
func (c testConfig) String(key string, defaultValue ...string) string {
if v, ok := c[c.key(key)].(string); ok {
return v
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return ""
}
func (c testConfig) Uint(key string, defaultValue ...uint32) uint32 {
switch v := c[c.key(key)].(type) {
case uint32:
return v
case int:
return uint32(v)
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return 0
}
func (c testConfig) Float(key string, defaultValue ...float32) float32 {
if v, ok := c[c.key(key)].(float32); ok {
return v
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return 0
}
func (c testConfig) Bool(key string, defaultValue ...bool) bool {
if v, ok := c[c.key(key)].(bool); ok {
return v
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return false
}
func (c testConfig) Strings(key string, defaultValue ...[]string) []string {
if v, ok := c[c.key(key)].([]string); ok {
return v
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return nil
}
func (c testConfig) Ints(key string, defaultValue ...[]int32) []int32 {
if v, ok := c[c.key(key)].([]int32); ok {
return v
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return nil
}
func (c testConfig) Uints(key string, defaultValue ...[]uint32) []uint32 {
if v, ok := c[c.key(key)].([]uint32); ok {
return v
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return nil
}
func (c testConfig) Floats(key string, defaultValue ...[]float32) []float32 {
if v, ok := c[c.key(key)].([]float32); ok {
return v
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return nil
}
func (c testConfig) Bools(key string, defaultValue ...[]bool) []bool {
if v, ok := c[c.key(key)].([]bool); ok {
return v
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return nil
}
func (c testConfig) Len() int { return len(c) }
func (c testConfig) Keys() iter.Seq[string] {
return func(yield func(string) bool) {
for key := range c {
if !yield(key) {
return
}
}
}
}
func (c testConfig) Value(key string) any { return c[key] }
func TestNewOptionsLayerConfig(t *testing.T) {
cfg := testConfig{
"laguna.block_count": uint32(4),
"laguna.embedding_length": uint32(128),
"laguna.attention.key_length": uint32(16),
"laguna.attention.head_count": []uint32{8, 16, 16, 16},
"laguna.attention.head_count_kv": uint32(4),
"laguna.attention.layer_norm_rms_epsilon": float32(1e-6),
"laguna.attention.sliding_window": uint32(512),
"laguna.attention.sliding_window_pattern": []bool{false, true, true, true},
"laguna.rope.dimension_count": uint32(8),
"laguna.rope.freq_base": float32(500000),
"laguna.rope.scaling.factor": float32(32),
"laguna.rope.scaling.original_context_length": uint32(4096),
"laguna.rope.swa.dimension_count": uint32(16),
"laguna.rope.swa.freq_base": float32(10000),
"laguna.expert_count": uint32(32),
"laguna.expert_used_count": uint32(4),
"laguna.expert_weights_norm": true,
"laguna.expert_weights_scale": float32(2.5),
"laguna.decoder_sparse_step": uint32(1),
"laguna.dense_layers": []uint32{0},
}
opts := newOptions(cfg, 4)
if got := opts.numHeadsForLayer(0); got != 8 {
t.Fatalf("layer 0 heads = %d, want 8", got)
}
if got := opts.numHeadsForLayer(1); got != 16 {
t.Fatalf("layer 1 heads = %d, want 16", got)
}
if opts.layerIsSliding(0) {
t.Fatal("layer 0 should be full attention")
}
if !opts.layerIsSliding(1) {
t.Fatal("layer 1 should be sliding attention")
}
if opts.layerUsesMoE(0) {
t.Fatal("layer 0 should be dense")
}
if !opts.layerUsesMoE(1) {
t.Fatal("layer 1 should use MoE")
}
if opts.fullRopeDim != 8 || opts.swaRopeDim != 16 {
t.Fatalf("rope dims = full %d swa %d, want 8/16", opts.fullRopeDim, opts.swaRopeDim)
}
}
func TestNewOptionsYarnAttentionFactorFallback(t *testing.T) {
cfg := testConfig{
"laguna.block_count": uint32(1),
"laguna.embedding_length": uint32(128),
"laguna.attention.key_length": uint32(16),
"laguna.attention.head_count": uint32(8),
"laguna.attention.head_count_kv": uint32(4),
"laguna.rope.scaling.type": "yarn",
"laguna.rope.scaling.factor": float32(32),
}
opts := newOptions(cfg, 1)
want := float32(0.1*math.Log(32) + 1)
if got := opts.fullRopeAttentionFactor; math.Abs(float64(got-want)) > 1e-6 {
t.Fatalf("fullRopeAttentionFactor = %v, want %v", got, want)
}
}
func TestNewRejectsUnsupportedLagunaVariants(t *testing.T) {
tests := []struct {
name string
cfg testConfig
}{
{
name: "attention sinks",
cfg: testConfig{
"laguna.attention.sink_enabled": true,
},
},
{
name: "non per-head gate",
cfg: testConfig{
"laguna.attention.gating_type": uint32(0),
},
},
{
name: "missing qk norm",
cfg: testConfig{
"laguna.attention.gating_type": uint32(1),
},
},
{
name: "non sigmoid experts",
cfg: testConfig{
"laguna.attention.gating_type": uint32(1),
"laguna.attention.qk_norm": true,
"laguna.expert_gating_func": uint32(1),
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if _, err := New(tt.cfg); err == nil {
t.Fatal("expected unsupported variant error")
}
})
}
}
+1
View File
@@ -11,6 +11,7 @@ import (
_ "github.com/ollama/ollama/model/models/glm4moelite"
_ "github.com/ollama/ollama/model/models/glmocr"
_ "github.com/ollama/ollama/model/models/gptoss"
_ "github.com/ollama/ollama/model/models/laguna"
_ "github.com/ollama/ollama/model/models/lfm2"
_ "github.com/ollama/ollama/model/models/llama"
_ "github.com/ollama/ollama/model/models/llama4"
+355
View File
@@ -0,0 +1,355 @@
package nemotronh
import (
"errors"
"image"
"math"
"slices"
"github.com/ollama/ollama/fs"
"github.com/ollama/ollama/model/imageproc"
)
type ImageProcessor struct {
imageSize int
patchSize int
numChannels int
maxTiles int
minNumPatches int
maxNumPatches int
useThumbnail bool
projectorScale int
imageMean [3]float32
imageStd [3]float32
}
type processedVisionTile struct {
data []float32
size image.Point
}
func newImageProcessor(c fs.Config) ImageProcessor {
mean := c.Floats("vision.image_mean")
std := c.Floats("vision.image_std")
processor := ImageProcessor{
imageSize: int(c.Uint("vision.image_size", 512)),
patchSize: int(c.Uint("vision.patch_size", 16)),
numChannels: int(c.Uint("vision.num_channels", 3)),
maxTiles: int(c.Uint("vision.max_tiles", 12)),
minNumPatches: int(c.Uint("vision.min_num_patches")),
maxNumPatches: int(c.Uint("vision.max_num_patches")),
useThumbnail: c.Bool("vision.use_thumbnail", true),
projectorScale: int(c.Uint("vision.projector.scale_factor", 2)),
imageMean: imageproc.ClipDefaultMean,
imageStd: imageproc.ClipDefaultSTD,
}
if len(mean) >= 3 {
processor.imageMean = [3]float32{mean[0], mean[1], mean[2]}
}
if len(std) >= 3 {
processor.imageStd = [3]float32{std[0], std[1], std[2]}
}
if processor.imageSize <= 0 {
processor.imageSize = 512
}
if processor.patchSize <= 0 {
processor.patchSize = 16
}
if processor.numChannels <= 0 {
processor.numChannels = 3
}
if processor.maxTiles <= 0 {
processor.maxTiles = 12
}
if processor.projectorScale <= 0 {
processor.projectorScale = 2
}
return processor
}
func (p ImageProcessor) ProcessImage(img image.Image) ([]processedVisionTile, error) {
img = imageproc.Composite(img)
if p.useDynamicResolution() {
return p.processDynamicImage(img)
}
return p.processTiledImage(img), nil
}
func (p ImageProcessor) useDynamicResolution() bool {
return p.minNumPatches > 0 || p.maxNumPatches > 0
}
func (p ImageProcessor) processTiledImage(img image.Image) []processedVisionTile {
bounds := img.Bounds()
origWidth := bounds.Dx()
origHeight := bounds.Dy()
targetRatios := nemotronTargetRatios(p.maxTiles)
gridWidth, gridHeight := findClosestAspectRatio(float64(origWidth)/float64(origHeight), targetRatios, origWidth, origHeight, p.imageSize)
targetWidth := p.imageSize * gridWidth
targetHeight := p.imageSize * gridHeight
resized := resizeImageBicubicCHW(img, targetWidth, targetHeight)
tiles := make([]processedVisionTile, 0, gridWidth*gridHeight+1)
for row := range gridHeight {
for col := range gridWidth {
tile := cropCHWRegion(
resized,
targetWidth,
targetHeight,
p.numChannels,
col*p.imageSize,
row*p.imageSize,
p.imageSize,
p.imageSize,
)
tiles = append(tiles, processedVisionTile{
data: normalizeVisionCHW(tile, p.imageMean, p.imageStd),
size: image.Point{X: p.imageSize, Y: p.imageSize},
})
}
}
if p.useThumbnail && len(tiles) > 1 {
thumbnail := resizeImageBicubicCHW(img, p.imageSize, p.imageSize)
tiles = append(tiles, processedVisionTile{
data: normalizeVisionCHW(thumbnail, p.imageMean, p.imageStd),
size: image.Point{X: p.imageSize, Y: p.imageSize},
})
}
return tiles
}
func (p ImageProcessor) processDynamicImage(img image.Image) ([]processedVisionTile, error) {
bounds := img.Bounds()
origWidth := bounds.Dx()
origHeight := bounds.Dy()
patchesWidth, patchesHeight := p.dynamicPatchGrid(origWidth, origHeight)
if patchesWidth <= 0 || patchesHeight <= 0 {
return nil, errors.New("nemotron_h_omni: invalid dynamic image patch grid")
}
targetWidth := patchesWidth * p.patchSize
targetHeight := patchesHeight * p.patchSize
resized := resizeImageBicubicCHW(img, targetWidth, targetHeight)
return []processedVisionTile{{
data: normalizeVisionCHW(resized, p.imageMean, p.imageStd),
size: image.Point{X: targetWidth, Y: targetHeight},
}}, nil
}
func (p ImageProcessor) dynamicPatchGrid(origWidth, origHeight int) (int, int) {
patchesHeight := max(1, int(math.Round(float64(origHeight)/float64(p.patchSize)+0.5)))
patchesWidth := max(1, int(math.Round(float64(origWidth)/float64(p.patchSize)+0.5)))
patches := patchesHeight * patchesWidth
currentNumPatchesAvailable := p.maxNumPatches
if currentNumPatchesAvailable <= 0 {
currentNumPatchesAvailable = max(patches, p.minNumPatches)
}
factor := math.Min(math.Sqrt(float64(currentNumPatchesAvailable)/float64(patches)), 1.0)
targetPatchesHeight := max(1, int(math.Floor(factor*float64(patchesHeight))))
targetPatchesWidth := max(1, int(math.Floor(factor*float64(patchesWidth))))
if currentNumPatchesAvailable > p.minNumPatches && targetPatchesHeight*targetPatchesWidth < p.minNumPatches {
upFactor := math.Sqrt(float64(p.minNumPatches) / float64(targetPatchesHeight*targetPatchesWidth))
targetPatchesHeight = int(math.Ceil(upFactor * float64(targetPatchesHeight)))
targetPatchesWidth = int(math.Ceil(upFactor * float64(targetPatchesWidth)))
}
targetPatchesHeight = roundPatchGridForPixelShuffle(targetPatchesHeight, targetPatchesWidth, currentNumPatchesAvailable, p.projectorScale)
targetPatchesWidth = roundPatchGridForPixelShuffle(targetPatchesWidth, targetPatchesHeight, currentNumPatchesAvailable, p.projectorScale)
return targetPatchesWidth, targetPatchesHeight
}
func roundPatchGridForPixelShuffle(v, other, maxPatches, divisor int) int {
if divisor <= 1 {
return v
}
rem := v % divisor
if rem == 0 {
return v
}
inc := divisor - rem
if (v+inc)*other <= maxPatches {
return v + inc
}
return max(divisor, v-rem)
}
type nemotronImageRatio struct {
width int
height int
}
func nemotronTargetRatios(maxTiles int) []nemotronImageRatio {
targetRatios := make([]nemotronImageRatio, 0, maxTiles*maxTiles)
for n := 1; n <= maxTiles; n++ {
for w := 1; w <= n; w++ {
for h := 1; h <= n; h++ {
if w*h > maxTiles {
continue
}
targetRatios = append(targetRatios, nemotronImageRatio{width: w, height: h})
}
}
}
unique := targetRatios[:0]
for _, ratio := range targetRatios {
if slices.Contains(unique, ratio) {
continue
}
unique = append(unique, ratio)
}
slices.SortFunc(unique, func(a, b nemotronImageRatio) int {
return a.width*a.height - b.width*b.height
})
return unique
}
func findClosestAspectRatio(aspectRatio float64, targetRatios []nemotronImageRatio, width, height, imageSize int) (int, int) {
bestRatio := nemotronImageRatio{width: 1, height: 1}
bestRatioDiff := math.MaxFloat64
area := width * height
for _, ratio := range targetRatios {
targetAspectRatio := float64(ratio.width) / float64(ratio.height)
ratioDiff := math.Abs(aspectRatio - targetAspectRatio)
if ratioDiff < bestRatioDiff {
bestRatioDiff = ratioDiff
bestRatio = ratio
continue
}
if ratioDiff == bestRatioDiff && area > int(0.5*float64(imageSize*imageSize*ratio.width*ratio.height)) {
bestRatio = ratio
}
}
return bestRatio.width, bestRatio.height
}
func resizeImageBicubicCHW(img image.Image, outW, outH int) []float32 {
bounds := img.Bounds()
inW := bounds.Dx()
inH := bounds.Dy()
src := make([]float32, 3*inW*inH)
for y := range inH {
for x := range inW {
r, g, b, _ := img.At(bounds.Min.X+x, bounds.Min.Y+y).RGBA()
src[y*inW+x] = float32(r>>8) / 255.0
src[inW*inH+y*inW+x] = float32(g>>8) / 255.0
src[2*inW*inH+y*inW+x] = float32(b>>8) / 255.0
}
}
dst := make([]float32, 3*outW*outH)
scaleX := float64(inW) / float64(outW)
scaleY := float64(inH) / float64(outH)
for oy := range outH {
srcY := scaleY*(float64(oy)+0.5) - 0.5
yBase := int(math.Floor(srcY))
yFrac := clampUnit(srcY - float64(yBase))
wy := torchBicubicWeights(yFrac)
for ox := range outW {
srcX := scaleX*(float64(ox)+0.5) - 0.5
xBase := int(math.Floor(srcX))
xFrac := clampUnit(srcX - float64(xBase))
wx := torchBicubicWeights(xFrac)
for c := range 3 {
var sum float64
channelOffset := c * inW * inH
for ky := range 4 {
iy := clampIndex(yBase-1+ky, 0, inH-1)
for kx := range 4 {
ix := clampIndex(xBase-1+kx, 0, inW-1)
sum += float64(src[channelOffset+iy*inW+ix]) * wy[ky] * wx[kx]
}
}
dst[c*outW*outH+oy*outW+ox] = float32(sum)
}
}
}
return dst
}
func cropCHWRegion(values []float32, width, height, channels, left, top, cropW, cropH int) []float32 {
out := make([]float32, channels*cropW*cropH)
channelSize := width * height
cropSize := cropW * cropH
for c := range channels {
srcBase := c * channelSize
dstBase := c * cropSize
for y := range cropH {
copy(out[dstBase+y*cropW:dstBase+(y+1)*cropW], values[srcBase+(top+y)*width+left:srcBase+(top+y)*width+left+cropW])
}
}
return out
}
func normalizeVisionCHW(values []float32, mean, std [3]float32) []float32 {
out := make([]float32, len(values))
channelSize := len(values) / 3
for c := range 3 {
base := c * channelSize
for i := range channelSize {
out[base+i] = (values[base+i] - mean[c]) / std[c]
}
}
return out
}
func torchBicubicWeights(t float64) [4]float64 {
const a = -0.75
return [4]float64{
bicubicConvolution2(t+1.0, a),
bicubicConvolution1(t, a),
bicubicConvolution1(1.0-t, a),
bicubicConvolution2(2.0-t, a),
}
}
func bicubicConvolution1(x, a float64) float64 {
return ((a+2)*x-(a+3))*x*x + 1
}
func bicubicConvolution2(x, a float64) float64 {
return ((a*x-5*a)*x+8*a)*x - 4*a
}
func clampUnit(v float64) float64 {
if v < 0 {
return 0
}
if v > 1 {
return 1
}
return v
}
func clampIndex(v, lo, hi int) int {
if v < lo {
return lo
}
if v > hi {
return hi
}
return v
}
+20 -5
View File
@@ -117,9 +117,7 @@ func Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) {
return key, nil
}
func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
hiddenStates := m.TokenEmbedding.Forward(ctx, batch.Inputs)
func (m *Model) forwardHiddenStates(ctx ml.Context, batch input.Batch, hiddenStates ml.Tensor) (ml.Tensor, error) {
cache := m.Cache.(*HybridCache)
for i, layer := range m.Layers {
@@ -137,11 +135,24 @@ func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
}
}
hiddenStates = m.OutputNorm.Forward(ctx, hiddenStates, m.eps)
return m.OutputNorm.Forward(ctx, hiddenStates, m.eps), nil
}
func (m *Model) forwardLogits(ctx ml.Context, batch input.Batch, hiddenStates ml.Tensor) (ml.Tensor, error) {
hiddenStates, err := m.forwardHiddenStates(ctx, batch, hiddenStates)
if err != nil {
return nil, err
}
return m.Output.Forward(ctx, hiddenStates), nil
}
func New(c fs.Config) (model.Model, error) {
func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
hiddenStates := m.TokenEmbedding.Forward(ctx, batch.Inputs)
return m.forwardLogits(ctx, batch, hiddenStates)
}
func newTextModel(c fs.Config) (*Model, error) {
numLayers := int(c.Uint("block_count"))
layers := make([]Layer, numLayers)
@@ -306,6 +317,10 @@ func New(c fs.Config) (model.Model, error) {
return &m, nil
}
func New(c fs.Config) (model.Model, error) {
return newTextModel(c)
}
func init() {
model.Register("nemotron_h", New)
model.Register("nemotron_h_moe", New)
+511
View File
@@ -0,0 +1,511 @@
package nemotronh
import (
"math"
"sync"
"github.com/ollama/ollama/fs"
"github.com/ollama/ollama/ml"
"github.com/ollama/ollama/ml/nn"
)
type AudioOptions struct {
hiddenSize int
numHeads int
headDim int
intermediateSize int
convKernelSize int
melBins int
sampleRate int
subsamplingKernel int
subsamplingStride int
scaleInput bool
eps float32
}
type AudioFeatureExtractor struct {
FB ml.Tensor `gguf:"fb"`
Window ml.Tensor `gguf:"window"`
mu sync.Mutex
fb []float32
window []float32
fbShape [2]int
}
func (f *AudioFeatureExtractor) windowAndFilters(melBins, freqBins, sampleRate int) ([]float32, []float32) {
if f == nil {
return defaultParakeetWindow(), buildSlaneyMelFilterBank(freqBins, melBins, sampleRate)
}
f.mu.Lock()
defer f.mu.Unlock()
if f.window == nil {
if f.Window != nil {
if values := f.Window.BackendGet(); len(values) == parakeetWinLength {
f.window = values
}
}
if f.window == nil {
f.window = defaultParakeetWindow()
}
}
if f.fb == nil {
if f.FB != nil {
if values := f.FB.BackendGet(); len(values) == melBins*freqBins {
f.fb = values
f.fbShape = [2]int{melBins, freqBins}
}
}
if f.fb == nil {
f.fb = buildSlaneyMelFilterBank(freqBins, melBins, sampleRate)
f.fbShape = [2]int{melBins, freqBins}
}
}
return f.window, f.fb
}
type AudioSubsampling struct {
Conv0 *nn.Conv2D `gguf:"conv0"`
DW1 *AudioDepthwiseConv2D `gguf:"dw1"`
PW1 *nn.Conv2D `gguf:"pw1"`
DW2 *AudioDepthwiseConv2D `gguf:"dw2"`
PW2 *nn.Conv2D `gguf:"pw2"`
Linear *nn.Linear `gguf:"linear"`
}
type AudioDepthwiseConv2D struct {
Weight ml.Tensor `gguf:"weight"`
Bias ml.Tensor `gguf:"bias"`
}
type AudioFeedForward struct {
Up *nn.Linear `gguf:"up"`
Down *nn.Linear `gguf:"down"`
}
type AudioSelfAttention struct {
Query *nn.Linear `gguf:"attn_q"`
Key *nn.Linear `gguf:"attn_k"`
Value *nn.Linear `gguf:"attn_v"`
Output *nn.Linear `gguf:"attn_out"`
RelativeKey *nn.Linear `gguf:"attn_rel_k"`
BiasU ml.Tensor `gguf:"attn_bias_u"`
BiasV ml.Tensor `gguf:"attn_bias_v"`
}
type AudioConvolutionModule struct {
Pointwise1 *nn.Linear `gguf:"conv_pw1"`
Depthwise ml.Tensor `gguf:"conv_dw.weight"`
BatchNorm *AudioBatchNorm1D `gguf:"conv_bn"`
Pointwise2 *nn.Linear `gguf:"conv_pw2"`
}
type AudioBatchNorm1D struct {
Weight ml.Tensor `gguf:"weight"`
Bias ml.Tensor `gguf:"bias"`
RunningMean ml.Tensor `gguf:"running_mean"`
RunningVar ml.Tensor `gguf:"running_var"`
}
type AudioLayer struct {
FFN1Norm *nn.LayerNorm `gguf:"ffn1_norm"`
FFN1Up *nn.Linear `gguf:"ffn1_up"`
FFN1Down *nn.Linear `gguf:"ffn1_down"`
AttentionNorm *nn.LayerNorm `gguf:"attn_norm"`
Attention *AudioSelfAttention
ConvNorm *nn.LayerNorm `gguf:"conv_norm"`
Conv *AudioConvolutionModule
FFN2Norm *nn.LayerNorm `gguf:"ffn2_norm"`
FFN2Up *nn.Linear `gguf:"ffn2_up"`
FFN2Down *nn.Linear `gguf:"ffn2_down"`
OutputNorm *nn.LayerNorm `gguf:"out_norm"`
}
type AudioModel struct {
FeatureExtractor *AudioFeatureExtractor `gguf:"feature_extractor"`
Subsampling *AudioSubsampling `gguf:"subsampling"`
Layers []AudioLayer `gguf:"blk"`
*AudioOptions
}
type AudioProjector struct {
Norm *nn.RMSNorm `gguf:"norm"`
Linear1 *nn.Linear `gguf:"1"`
Linear2 *nn.Linear `gguf:"2"`
}
func (p *AudioProjector) Forward(ctx ml.Context, x ml.Tensor, eps float32) ml.Tensor {
x = p.Norm.Forward(ctx, x, eps)
x = audioF32(ctx, p.Linear1.Forward(ctx, x))
x = x.RELU(ctx)
x = x.Mul(ctx, x)
return audioF32(ctx, p.Linear2.Forward(ctx, x))
}
func (m *AudioModel) ForwardAudio(ctx ml.Context, melFeatures ml.Tensor, validFrames int, projector *AudioProjector) ml.Tensor {
x := melFeatures.Reshape(ctx, melFeatures.Dim(0), melFeatures.Dim(1), 1, 1)
validLen := validFrames
x = forwardAudioConv2D(ctx, m.Subsampling.Conv0, x, m.subsamplingStride, m.subsamplingStride, audioConvPadding(m.subsamplingKernel), audioConvPadding(m.subsamplingKernel), 1, 1)
x = x.RELU(ctx)
validLen = convOutputLength(validLen, m.subsamplingKernel, m.subsamplingStride, audioConvPadding(m.subsamplingKernel))
x = applyAudioTimeMask(ctx, x, validLen)
x = forwardAudioDepthwiseConv2D(ctx, m.Subsampling.DW1, x, m.subsamplingStride, m.subsamplingStride, audioConvPadding(m.subsamplingKernel), audioConvPadding(m.subsamplingKernel), 1, 1)
x = forwardAudioConv2D(ctx, m.Subsampling.PW1, x, 1, 1, 0, 0, 1, 1)
x = x.RELU(ctx)
validLen = convOutputLength(validLen, m.subsamplingKernel, m.subsamplingStride, audioConvPadding(m.subsamplingKernel))
x = applyAudioTimeMask(ctx, x, validLen)
x = forwardAudioDepthwiseConv2D(ctx, m.Subsampling.DW2, x, m.subsamplingStride, m.subsamplingStride, audioConvPadding(m.subsamplingKernel), audioConvPadding(m.subsamplingKernel), 1, 1)
x = forwardAudioConv2D(ctx, m.Subsampling.PW2, x, 1, 1, 0, 0, 1, 1)
x = x.RELU(ctx)
validLen = convOutputLength(validLen, m.subsamplingKernel, m.subsamplingStride, audioConvPadding(m.subsamplingKernel))
x = applyAudioTimeMask(ctx, x, validLen)
x = flattenAudioSubsamplingOutput(ctx, x)
x = m.Subsampling.Linear.Forward(ctx, x)
if m.scaleInput {
x = x.Scale(ctx, math.Sqrt(float64(m.hiddenSize)))
}
if validLen > 0 && validLen < x.Dim(1) {
x = x.Slice(ctx, 1, 0, validLen, 1).Contiguous(ctx)
}
for i := range m.Layers {
x = m.Layers[i].Forward(ctx, x, validLen, m.AudioOptions)
}
if projector != nil {
x = projector.Forward(ctx, x, m.eps)
}
return x
}
func flattenAudioSubsamplingOutput(ctx ml.Context, x ml.Tensor) ml.Tensor {
fOut, tOut, cOut := x.Dim(0), x.Dim(1), x.Dim(2)
// PyTorch flattens the subsampling output after [B, C, T, F] ->
// [B, T, C, F], so F must remain the fastest dimension inside each
// channel block before the linear projection.
x = x.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
return x.Reshape(ctx, cOut*fOut, tOut)
}
func (l *AudioLayer) Forward(ctx ml.Context, x ml.Tensor, validLen int, opts *AudioOptions) ml.Tensor {
residual := x
x = audioFeedForward(ctx, l.FFN1Up, l.FFN1Down, l.FFN1Norm.Forward(ctx, x, opts.eps)).Scale(ctx, 0.5)
x = residual.Add(ctx, x)
residual = x
x = l.Attention.Forward(ctx, l.AttentionNorm.Forward(ctx, x, opts.eps), validLen, opts)
x = residual.Add(ctx, x)
residual = x
x = l.Conv.Forward(ctx, l.ConvNorm.Forward(ctx, x, opts.eps), opts)
x = residual.Add(ctx, x)
residual = x
x = audioFeedForward(ctx, l.FFN2Up, l.FFN2Down, l.FFN2Norm.Forward(ctx, x, opts.eps)).Scale(ctx, 0.5)
x = residual.Add(ctx, x)
return l.OutputNorm.Forward(ctx, x, opts.eps)
}
func audioFeedForward(ctx ml.Context, up, down *nn.Linear, x ml.Tensor) ml.Tensor {
x = audioF32(ctx, up.Forward(ctx, x))
x = x.SILU(ctx)
return audioF32(ctx, down.Forward(ctx, x))
}
func (a *AudioSelfAttention) Forward(ctx ml.Context, x ml.Tensor, validLen int, opts *AudioOptions) ml.Tensor {
seqLen := x.Dim(1)
headDim := opts.headDim
numHeads := opts.numHeads
q := audioF32(ctx, a.Query.Forward(ctx, x)).Reshape(ctx, headDim, numHeads, seqLen)
k := audioF32(ctx, a.Key.Forward(ctx, x)).Reshape(ctx, headDim, numHeads, seqLen)
v := audioF32(ctx, a.Value.Forward(ctx, x)).Reshape(ctx, headDim, numHeads, seqLen)
qU := q
if a.BiasU != nil {
qU = qU.Add(ctx, audioF32(ctx, a.BiasU).Reshape(ctx, headDim, numHeads, 1))
}
qV := q
if a.BiasV != nil {
qV = qV.Add(ctx, audioF32(ctx, a.BiasV).Reshape(ctx, headDim, numHeads, 1))
}
qP := qU.Permute(ctx, 0, 2, 1, 3)
kP := k.Permute(ctx, 0, 2, 1, 3)
logits := kP.MulmatFullPrec(ctx, qP)
positionEmbeddings := parakeetPositionEmbeddings(ctx, seqLen, opts.hiddenSize)
relKey := audioF32(ctx, a.RelativeKey.Forward(ctx, positionEmbeddings)).Reshape(ctx, headDim, numHeads, 2*seqLen-1)
pP := relKey.Permute(ctx, 0, 2, 1, 3)
qVP := qV.Permute(ctx, 0, 2, 1, 3)
relLogits := pP.MulmatFullPrec(ctx, qVP)
relLogits = relativeShiftParakeet(ctx, relLogits, seqLen, numHeads)
logits = logits.Add(ctx, relLogits)
logits = logits.Scale(ctx, math.Pow(float64(headDim), -0.5))
if validLen > 0 && validLen < seqLen {
logits = logits.Add(ctx, audioAttentionMask(ctx, seqLen, validLen))
}
logits = logits.Softmax(ctx)
vP := v.Permute(ctx, 0, 2, 1, 3)
vPT := vP.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx)
out := vPT.Mulmat(ctx, logits)
out = out.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
out = out.Reshape(ctx, opts.hiddenSize, seqLen)
return audioF32(ctx, a.Output.Forward(ctx, out))
}
func (c *AudioConvolutionModule) Forward(ctx ml.Context, x ml.Tensor, opts *AudioOptions) ml.Tensor {
x = audioF32(ctx, c.Pointwise1.Forward(ctx, x))
hidden := x.Dim(0) / 2
value := x.Slice(ctx, 0, 0, hidden, 1).Contiguous(ctx)
gate := x.Slice(ctx, 0, hidden, 2*hidden, 1).Contiguous(ctx).Sigmoid(ctx)
x = value.Mul(ctx, gate)
x = audioDepthwiseConv1DSame(ctx, x, c.Depthwise, audioConvPadding(opts.convKernelSize))
x = c.BatchNorm.Forward(ctx, x, opts.eps)
x = x.SILU(ctx)
return audioF32(ctx, c.Pointwise2.Forward(ctx, x))
}
func audioF32(ctx ml.Context, x ml.Tensor) ml.Tensor {
if x.DType() == ml.DTypeF32 {
return x
}
// Metal binary kernels used by the audio graph require F32 operands here.
// This likely slows audio and should be revisited once the precision vs.
// speed tradeoff is validated against BF16-native elementwise paths.
return x.Cast(ctx, ml.DTypeF32)
}
func (b *AudioBatchNorm1D) Forward(ctx ml.Context, x ml.Tensor, eps float32) ml.Tensor {
if b == nil || b.RunningMean == nil || b.RunningVar == nil {
return x
}
hidden := x.Dim(0)
epsValues := make([]float32, hidden)
for i := range epsValues {
epsValues[i] = eps
}
variance := b.RunningVar.Add(ctx, ctx.Input().FromFloats(epsValues, hidden))
x = x.Sub(ctx, b.RunningMean)
x = x.Div(ctx, variance.Sqrt(ctx))
if b.Weight != nil {
x = x.Mul(ctx, b.Weight)
}
if b.Bias != nil {
x = x.Add(ctx, b.Bias)
}
return x
}
func forwardAudioConv2D(ctx ml.Context, conv *nn.Conv2D, x ml.Tensor, s0, s1, p0, p1, d0, d1 int) ml.Tensor {
weight := conv.Weight.Contiguous(ctx)
x = weight.Conv2D(ctx, x, s0, s1, p0, p1, d0, d1)
if conv.Bias != nil {
x = x.Add(ctx, conv.Bias.Reshape(ctx, 1, 1, -1))
}
return x
}
func forwardAudioDepthwiseConv2D(ctx ml.Context, conv *AudioDepthwiseConv2D, x ml.Tensor, s0, s1, p0, p1, d0, d1 int) ml.Tensor {
x = audioDepthwiseConv2D(ctx, x, conv.Weight, s0, s1, p0, p1, d0, d1)
if conv.Bias != nil {
x = x.Add(ctx, conv.Bias.Reshape(ctx, 1, 1, -1))
}
return x
}
func applyAudioTimeMask(ctx ml.Context, x ml.Tensor, validLen int) ml.Tensor {
if validLen <= 0 || validLen >= x.Dim(1) {
return x
}
mask := make([]float32, x.Dim(1))
for i := range validLen {
mask[i] = 1
}
return x.Mul(ctx, ctx.Input().FromFloats(mask, 1, x.Dim(1), 1, 1))
}
func audioDepthwiseConv1DSame(ctx ml.Context, x, kernel ml.Tensor, padding int) ml.Tensor {
kernelSize := kernel.Dim(0)
seqLen := x.Dim(1)
kernelT := kernel.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx)
var out ml.Tensor
for k := range kernelSize {
offset := k - padding
shifted := x
switch {
case offset > 0:
shifted = x.Slice(ctx, 1, offset, seqLen, 1).Contiguous(ctx)
shifted = shifted.PadExt(ctx, 0, 0, 0, offset, 0, 0, 0, 0)
case offset < 0:
shift := -offset
shifted = x.Slice(ctx, 1, 0, seqLen-shift, 1).Contiguous(ctx)
shifted = shifted.PadExt(ctx, 0, 0, shift, 0, 0, 0, 0, 0)
}
wk := kernelT.Slice(ctx, 1, k, k+1, 1).Contiguous(ctx)
term := shifted.Mul(ctx, wk)
if out == nil {
out = term
} else {
out = out.Add(ctx, term)
}
}
return out
}
func audioDepthwiseConv2D(ctx ml.Context, x, kernel ml.Tensor, s0, s1, p0, p1, d0, d1 int) ml.Tensor {
if d0 != 1 || d1 != 1 {
panic("audio depthwise conv2d only supports dilation 1")
}
kernel = kernel.Contiguous(ctx)
kernelW, kernelH := kernel.Dim(0), kernel.Dim(1)
outW := convOutputLength(x.Dim(0), kernelW, s0, p0)
outH := convOutputLength(x.Dim(1), kernelH, s1, p1)
padded := x.PadExt(ctx, p0, p0, p1, p1, 0, 0, 0, 0)
var out ml.Tensor
for ky := range kernelH {
for kx := range kernelW {
patch := padded.Slice(ctx, 0, kx, kx+s0*(outW-1)+1, s0).Contiguous(ctx)
patch = patch.Slice(ctx, 1, ky, ky+s1*(outH-1)+1, s1).Contiguous(ctx)
wk := kernel.Slice(ctx, 0, kx, kx+1, 1).Slice(ctx, 1, ky, ky+1, 1).Contiguous(ctx)
if wk.Dim(2) == 1 {
wk = wk.Permute(ctx, 0, 1, 3, 2).Contiguous(ctx)
} else {
wk = wk.Reshape(ctx, 1, 1, wk.Dim(2), wk.Dim(3))
}
term := patch.Mul(ctx, wk)
if out == nil {
out = term
} else {
out = out.Add(ctx, term)
}
}
}
return out
}
func convOutputLength(inputLength, kernel, stride, padding int) int {
if inputLength <= 0 {
return 0
}
return (inputLength+2*padding-kernel)/stride + 1
}
func audioConvPadding(kernel int) int {
return (kernel - 1) / 2
}
func parakeetPositionEmbeddings(ctx ml.Context, seqLen, hiddenSize int) ml.Tensor {
half := hiddenSize / 2
values := make([]float32, hiddenSize*(2*seqLen-1))
for posIdx, pos := 0, seqLen-1; posIdx < 2*seqLen-1; posIdx, pos = posIdx+1, pos-1 {
for i := range half {
invFreq := math.Pow(10000, -float64(2*i)/float64(hiddenSize))
angle := float64(pos) * invFreq
values[posIdx*hiddenSize+2*i] = float32(math.Sin(angle))
values[posIdx*hiddenSize+2*i+1] = float32(math.Cos(angle))
}
}
return ctx.Input().FromFloats(values, hiddenSize, 2*seqLen-1)
}
func relativeShiftParakeet(ctx ml.Context, x ml.Tensor, seqLen, numHeads int) ml.Tensor {
positionLen := 2*seqLen - 1
x = x.PadExt(ctx, 1, 0, 0, 0, 0, 0, 0, 0)
x = x.Reshape(ctx, seqLen, positionLen+1, numHeads)
x = x.Slice(ctx, 1, 1, positionLen+1, 1).Contiguous(ctx)
x = x.Reshape(ctx, positionLen, seqLen, numHeads)
return x.Slice(ctx, 0, 0, seqLen, 1).Contiguous(ctx)
}
func audioAttentionMask(ctx ml.Context, seqLen, validLen int) ml.Tensor {
values := make([]float32, seqLen*seqLen)
for q := range seqLen {
for k := range seqLen {
if q >= validLen || k >= validLen {
values[q*seqLen+k] = -1e9
}
}
}
return ctx.Input().FromFloats(values, seqLen, seqLen, 1)
}
func newAudioModel(c fs.Config) *AudioModel {
numLayers := int(c.Uint("audio.block_count", 0))
if numLayers == 0 {
return nil
}
return &AudioModel{
Layers: make([]AudioLayer, numLayers),
AudioOptions: newAudioOptions(c),
}
}
func newAudioProjector(c fs.Config) *AudioProjector {
if c.Uint("audio.block_count", 0) == 0 {
return nil
}
return &AudioProjector{}
}
func newAudioOptions(c fs.Config) *AudioOptions {
hiddenSize := int(c.Uint("audio.embedding_length", 1024))
numHeads := int(c.Uint("audio.attention.head_count", 8))
headDim := hiddenSize / max(1, numHeads)
return &AudioOptions{
hiddenSize: hiddenSize,
numHeads: numHeads,
headDim: headDim,
intermediateSize: int(c.Uint("audio.feed_forward_length", uint32(hiddenSize*4))),
convKernelSize: int(c.Uint("audio.conv_kernel_size", 9)),
melBins: int(c.Uint("audio.num_mel_bins", 128)),
sampleRate: int(c.Uint("audio.sample_rate", 16000)),
subsamplingKernel: int(c.Uint("audio.subsampling_conv_kernel_size", 3)),
subsamplingStride: int(c.Uint("audio.subsampling_conv_stride", 2)),
scaleInput: c.Bool("audio.scale_input", false),
eps: c.Float("audio.attention.layer_norm_epsilon", 1e-5),
}
}
func defaultAudioOptions() *AudioOptions {
return &AudioOptions{
hiddenSize: 1024,
numHeads: 8,
headDim: 128,
intermediateSize: 4096,
convKernelSize: 9,
melBins: 128,
sampleRate: 16000,
subsamplingKernel: 3,
subsamplingStride: 2,
eps: 1e-5,
}
}
+239
View File
@@ -0,0 +1,239 @@
package nemotronh
import (
"bytes"
"errors"
"image"
"slices"
"github.com/ollama/ollama/fs"
"github.com/ollama/ollama/ml"
"github.com/ollama/ollama/model"
"github.com/ollama/ollama/model/input"
)
type OmniModel struct {
*Model
*VisionModel `gguf:"v"`
*AudioModel `gguf:"a"`
*MultiModalProjector `gguf:"mm"`
*AudioProjector `gguf:"mm.a"`
ImageProcessor
imageTokenID int32
imageStartToken int32
imageEndToken int32
audioTokenID int32
}
var _ model.MultimodalProcessor = (*OmniModel)(nil)
func NewOmni(c fs.Config) (model.Model, error) {
textModel, err := newTextModel(c)
if err != nil {
return nil, err
}
imageTokenID := int32(c.Uint("vision.image_token_id", 18))
imageStartToken := int32(c.Uint("vision.image_start_token_id", 19))
imageEndToken := int32(c.Uint("vision.image_end_token_id", 20))
audioTokenID := int32(c.Uint("audio.sound_token_id", 27))
return &OmniModel{
Model: textModel,
VisionModel: newVisionModel(c),
AudioModel: newAudioModel(c),
MultiModalProjector: newMultiModalProjector(c),
AudioProjector: newAudioProjector(c),
ImageProcessor: newImageProcessor(c),
imageTokenID: imageTokenID,
imageStartToken: imageStartToken,
imageEndToken: imageEndToken,
audioTokenID: audioTokenID,
}, nil
}
func (m *OmniModel) EncodeMultimodal(ctx ml.Context, multimodalData []byte) ([]input.Multimodal, error) {
if isAudioData(multimodalData) {
return m.encodeAudioMultimodal(ctx, multimodalData)
}
if m.VisionModel == nil || m.MultiModalProjector == nil || len(m.VisionModel.Layers) == 0 {
return nil, model.ErrNoVisionModel
}
img, _, err := image.Decode(bytes.NewReader(multimodalData))
if err != nil {
return nil, err
}
tiles, err := m.ImageProcessor.ProcessImage(img)
if err != nil {
return nil, err
}
mm := make([]input.Multimodal, 0, len(tiles))
for _, tile := range tiles {
patches := visionPatchGrid{
Width: tile.size.X / m.ImageProcessor.patchSize,
Height: tile.size.Y / m.ImageProcessor.patchSize,
}
if patches.Width == 0 || patches.Height == 0 {
return nil, errors.New("nemotron_h_omni: invalid resized image dimensions")
}
patchInput := packVisionPatchesCHW(tile.data, tile.size.X, tile.size.Y, m.ImageProcessor.numChannels, m.ImageProcessor.patchSize)
visionOutputs := m.VisionModel.ForwardPacked(ctx, patchInput, patches)
projected := m.MultiModalProjector.Forward(ctx, visionOutputs, patches)
mm = append(mm, input.Multimodal{Tensor: projected})
}
return mm, nil
}
type audioTag struct{}
func (m *OmniModel) encodeAudioMultimodal(ctx ml.Context, data []byte) ([]input.Multimodal, error) {
if m.AudioModel == nil || m.AudioProjector == nil || len(m.AudioModel.Layers) == 0 {
return nil, model.ErrNoVisionModel
}
samples, err := decodeWAV(data, m.AudioModel.sampleRate)
if err != nil {
return nil, err
}
melData, frames, validFrames, err := computeParakeetMelSpectrogram(samples, m.AudioModel.FeatureExtractor, m.AudioModel.AudioOptions)
if err != nil {
return nil, err
}
melTensor := ctx.Input().FromFloats(melData, m.AudioModel.melBins, frames)
audioOutputs := m.AudioModel.ForwardAudio(ctx, melTensor, validFrames, m.AudioProjector)
return []input.Multimodal{{Tensor: audioOutputs, Data: audioTag{}}}, nil
}
func (m *OmniModel) PostLoad() error {
return nil
}
func (m *OmniModel) PostTokenize(inputs []*input.Input) ([]*input.Input, error) {
var result []*input.Input
imageToken := m.imageTokenID
if imageToken == 0 {
imageToken = 18
}
for _, inp := range inputs {
if len(inp.Multimodal) == 0 {
result = append(result, inp)
continue
}
totalTokens := 0
for _, mm := range inp.Multimodal {
if mm.Tensor == nil {
continue
}
totalTokens += mm.Tensor.Dim(1)
}
if totalTokens <= 0 {
return nil, errors.New("nemotron_h_omni: multimodal input has no tokens")
}
if _, ok := inp.Multimodal[0].Data.(audioTag); ok {
audioToken := m.audioTokenID
if audioToken == 0 {
audioToken = 27
}
for i, mm := range inp.Multimodal {
tokenCount := 0
if mm.Tensor != nil {
tokenCount = mm.Tensor.Dim(1)
}
if tokenCount <= 0 {
return nil, errors.New("nemotron_h_omni: multimodal input has no tokens")
}
first := &input.Input{Token: audioToken, SameBatch: tokenCount - 1}
if i == 0 {
first.MultimodalHash = inp.MultimodalHash
}
first.Multimodal = []input.Multimodal{mm}
result = append(result, first)
if tokenCount > 1 {
result = append(result, slices.Repeat([]*input.Input{{Token: audioToken}}, tokenCount-1)...)
}
}
continue
}
if m.imageStartToken > 0 {
result = append(result, &input.Input{
Token: m.imageStartToken,
SameBatch: totalTokens + btoi(m.imageEndToken > 0),
})
}
for _, mm := range inp.Multimodal {
tokenCount := 0
if mm.Tensor != nil {
tokenCount = mm.Tensor.Dim(1)
}
if tokenCount <= 0 {
return nil, errors.New("nemotron_h_omni: multimodal input has no tokens")
}
result = append(result, &input.Input{
Token: imageToken,
Multimodal: []input.Multimodal{mm},
MultimodalHash: inp.MultimodalHash,
})
if tokenCount > 1 {
result = append(result, slices.Repeat([]*input.Input{{Token: imageToken}}, tokenCount-1)...)
}
}
if m.imageEndToken > 0 {
result = append(result, &input.Input{Token: m.imageEndToken})
}
}
return result, nil
}
func btoi(v bool) int {
if v {
return 1
}
return 0
}
func (m *OmniModel) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
hiddenStates := m.TokenEmbedding.Forward(ctx, batch.Inputs)
if len(batch.Multimodal) > 0 {
hiddenStates = hiddenStates.Duplicate(ctx)
}
for _, mm := range batch.Multimodal {
offset := mm.Index
for _, multimodal := range mm.Multimodal {
if multimodal.Tensor == nil {
continue
}
tensor := multimodal.Tensor
ctx.Forward(tensor.Copy(ctx, hiddenStates.View(ctx, offset*hiddenStates.Stride(1), tensor.Dim(0)*tensor.Dim(1))))
offset += tensor.Dim(1)
}
}
return m.forwardLogits(ctx, batch, hiddenStates)
}
func init() {
model.Register("nemotron_h_omni", NewOmni)
}
+606
View File
@@ -0,0 +1,606 @@
package nemotronh
import (
"bytes"
"encoding/base64"
"encoding/binary"
"image"
"image/color"
"math"
"os"
"path/filepath"
"slices"
"strings"
"testing"
fsggml "github.com/ollama/ollama/fs/ggml"
"github.com/ollama/ollama/ml"
backendggml "github.com/ollama/ollama/ml/backend/ggml"
"github.com/ollama/ollama/ml/nn"
"github.com/ollama/ollama/model/input"
)
type fakeTensor struct {
*backendggml.Tensor
dims []int
}
func (t *fakeTensor) Dim(i int) int {
return t.dims[i]
}
func setupTestContext(t *testing.T) ml.Context {
t.Helper()
f, err := os.CreateTemp(t.TempDir(), "*.gguf")
if err != nil {
t.Fatal(err)
}
defer f.Close()
if err := fsggml.WriteGGUF(f, fsggml.KV{"general.architecture": "test"}, nil); err != nil {
t.Fatal(err)
}
b, err := ml.NewBackend(f.Name(), ml.BackendParams{AllocMemory: true})
if err != nil {
t.Fatal(err)
}
ctx := b.NewContext().Input()
t.Cleanup(func() {
ctx.Close()
b.Close()
})
return ctx
}
func TestPostTokenizeImageSpans(t *testing.T) {
m := &OmniModel{
imageTokenID: 18,
imageStartToken: 19,
imageEndToken: 20,
}
makeChunk := func() input.Multimodal {
return input.Multimodal{Tensor: &fakeTensor{dims: []int{2688, 256, 1, 1}}}
}
in := []*input.Input{
{Token: 7},
{
Multimodal: []input.Multimodal{makeChunk(), makeChunk()},
MultimodalHash: 99,
},
{Token: 8},
}
out, err := m.PostTokenize(in)
if err != nil {
t.Fatalf("PostTokenize() error = %v", err)
}
if len(out) != 516 {
t.Fatalf("len(out) = %d, want 516", len(out))
}
if out[0].Token != 7 {
t.Fatalf("out[0].Token = %d, want 7", out[0].Token)
}
if out[1].Token != 19 {
t.Fatalf("out[1].Token = %d, want 19", out[1].Token)
}
if out[1].SameBatch != 513 {
t.Fatalf("out[1].SameBatch = %d, want 513", out[1].SameBatch)
}
if out[2].Token != 18 || len(out[2].Multimodal) != 1 || out[2].MultimodalHash != 99 || out[2].SameBatch != 0 {
t.Fatalf("unexpected first image token: %+v", *out[2])
}
if out[258].Token != 18 || len(out[258].Multimodal) != 1 || out[258].MultimodalHash != 99 || out[258].SameBatch != 0 {
t.Fatalf("unexpected second image token: %+v", *out[258])
}
if out[514].Token != 20 {
t.Fatalf("out[514].Token = %d, want 20", out[514].Token)
}
if out[515].Token != 8 {
t.Fatalf("out[515].Token = %d, want 8", out[515].Token)
}
}
func TestProjectorPixelShuffleMatchesReferenceV2Order(t *testing.T) {
ctx := setupTestContext(t)
hidden := 2
width := 4
height := 2
values := make([]float32, 0, hidden*width*height)
for y := range height {
for x := range width {
for c := range hidden {
values = append(values, float32(100*y+10*x+c))
}
}
}
got := pixelShuffleVisionOutputs(ctx, ctx.FromFloats(values, hidden, width*height), visionPatchGrid{
Width: width,
Height: height,
}, 2)
ctx.Forward(got).Compute(got)
want := []float32{
0, 1, 10, 11, 100, 101, 110, 111,
20, 21, 30, 31, 120, 121, 130, 131,
}
if got.Shape()[0] != 8 || got.Shape()[1] != 2 {
t.Fatalf("shape = %v, want [8 2 1]", got.Shape())
}
gotValues := got.BackendGet()
if len(gotValues) != len(want) {
t.Fatalf("len(got) = %d, want %d", len(gotValues), len(want))
}
for i := range want {
if gotValues[i] != want[i] {
t.Fatalf("got[%d] = %v, want %v", i, gotValues[i], want[i])
}
}
}
func TestPostTokenizeAudioSpans(t *testing.T) {
m := &OmniModel{
audioTokenID: 27,
}
in := []*input.Input{
{Token: 7},
{
Multimodal: []input.Multimodal{{
Tensor: &fakeTensor{dims: []int{2688, 13, 1, 1}},
Data: audioTag{},
}},
MultimodalHash: 99,
},
{Token: 8},
}
out, err := m.PostTokenize(in)
if err != nil {
t.Fatalf("PostTokenize() error = %v", err)
}
if len(out) != 15 {
t.Fatalf("len(out) = %d, want 15", len(out))
}
if out[0].Token != 7 || out[14].Token != 8 {
t.Fatalf("unexpected surrounding tokens: first=%d last=%d", out[0].Token, out[14].Token)
}
for i := 1; i <= 13; i++ {
if out[i].Token != 27 {
t.Fatalf("out[%d].Token = %d, want 27", i, out[i].Token)
}
}
if len(out[1].Multimodal) != 1 || out[1].MultimodalHash != 99 {
t.Fatalf("first audio token did not carry multimodal payload: %+v", *out[1])
}
if out[1].SameBatch != 12 {
t.Fatalf("first audio token SameBatch = %d, want 12", out[1].SameBatch)
}
if len(out[2].Multimodal) != 0 {
t.Fatalf("only the first audio token should carry multimodal payload: %+v", *out[2])
}
}
func TestParakeetAudioPreprocessShapes(t *testing.T) {
data := sineWAV(t, 16000, 440, 1.0)
samples, err := decodeWAV(data, 16000)
if err != nil {
t.Fatal(err)
}
if got, want := len(samples), 16000; got != want {
t.Fatalf("sample count = %d, want %d", got, want)
}
mel, frames, validFrames, err := computeParakeetMelSpectrogram(samples, nil, defaultAudioOptions())
if err != nil {
t.Fatal(err)
}
if frames != 101 {
t.Fatalf("frames = %d, want 101", frames)
}
if validFrames != 100 {
t.Fatalf("validFrames = %d, want 100", validFrames)
}
if len(mel) != 101*128 {
t.Fatalf("len(mel) = %d, want %d", len(mel), 101*128)
}
lastFrame := mel[100*128 : 101*128]
if !slices.Equal(lastFrame, make([]float32, 128)) {
t.Fatal("expected masked final frame to be zero")
}
}
func TestParakeetAudioPreprocessMatchesIntegrationWAVReference(t *testing.T) {
data := integrationAudioWAV(t)
samples, err := decodeWAV(data, 16000)
if err != nil {
t.Fatal(err)
}
if got, want := len(samples), 42083; got != want {
t.Fatalf("sample count = %d, want %d", got, want)
}
mel, frames, validFrames, err := computeParakeetMelSpectrogram(samples, nil, defaultAudioOptions())
if err != nil {
t.Fatal(err)
}
if frames != 264 {
t.Fatalf("frames = %d, want 264", frames)
}
if validFrames != 263 {
t.Fatalf("validFrames = %d, want 263", validFrames)
}
if len(mel) != 264*128 {
t.Fatalf("len(mel) = %d, want %d", len(mel), 264*128)
}
lastFrame := mel[263*128 : 264*128]
if !slices.Equal(lastFrame, make([]float32, 128)) {
t.Fatal("expected masked final frame to be zero")
}
// Reference values come from the ParakeetExtractor path used by vLLM:
// pre-emphasis, torch.stft(center=True, pad_mode="constant"), Slaney mel
// filters, log guard 2^-24, and per-mel normalization over valid frames.
checks := map[[2]int]float32{
{0, 0}: -1.0855197,
{0, 50}: -0.93212974,
{1, 10}: -0.9735168,
{2, 100}: -0.6533053,
{50, 0}: 2.2483668,
{50, 127}: -0.3828735,
{100, 50}: 2.9742377,
{262, 0}: -0.9521758,
{262, 127}: -0.4602786,
{263, 50}: 0,
}
for pos, want := range checks {
got := mel[pos[0]*128+pos[1]]
if math.Abs(float64(got-want)) > 1e-4 {
t.Errorf("mel[%d,%d] = %v, want %v", pos[0], pos[1], got, want)
}
}
}
func integrationAudioWAV(t *testing.T) []byte {
t.Helper()
path := filepath.Join("..", "..", "..", "integration", "audio_test_data_test.go")
b, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
const marker = "const audioEncodingPrompt = `"
s := string(b)
start := strings.Index(s, marker)
if start < 0 {
t.Fatal("audioEncodingPrompt marker not found")
}
start += len(marker)
end := strings.Index(s[start:], "`")
if end < 0 {
t.Fatal("audioEncodingPrompt terminator not found")
}
data, err := base64.StdEncoding.DecodeString(strings.TrimSpace(s[start : start+end]))
if err != nil {
t.Fatal(err)
}
return data
}
func TestRelativeShiftParakeetMatchesReference(t *testing.T) {
ctx := setupTestContext(t)
seqLen := 3
positionLen := 2*seqLen - 1
values := make([]float32, seqLen*positionLen)
for q := range seqLen {
for p := range positionLen {
values[q*positionLen+p] = float32(q*10 + p)
}
}
x := ctx.FromFloats(values, positionLen, seqLen, 1)
got := relativeShiftParakeet(ctx, x, seqLen, 1)
ctx.Forward(got).Compute(got)
want := []float32{
2, 3, 4,
11, 12, 13,
20, 21, 22,
}
if !slices.Equal(got.BackendGet(), want) {
t.Fatalf("relative shift mismatch:\n got %v\nwant %v", got.BackendGet(), want)
}
}
func TestAudioDepthwiseConv2DMatchesReference(t *testing.T) {
ctx := setupTestContext(t)
freq, frames, channels := 4, 5, 2
xValues := make([]float32, freq*frames*channels)
for i := range xValues {
xValues[i] = float32(i)/10 - 1
}
kernelValues := make([]float32, 3*3*channels)
for i := range kernelValues {
kernelValues[i] = float32(i)/7 - 1
}
x := ctx.FromFloats(xValues, freq, frames, channels, 1)
kernel := ctx.FromFloats(kernelValues, 3, 3, 1, channels)
bias := ctx.FromFloats([]float32{0.25, -0.5}, channels)
got := audioDepthwiseConv2D(ctx, x, kernel, 2, 2, 1, 1, 1, 1).Add(ctx, bias.Reshape(ctx, 1, 1, -1))
ctx.Forward(got).Compute(got)
want := []float32{
0.86428565, 1.3357141,
1.2785715, 1.3642857,
-0.5928571, -1.7499999,
5.4000001, 8.8142853,
10.514286, 16.042856,
6.6857138, 9.8428574,
}
assertCloseSlice(t, got.BackendGet(), want, 1e-5)
}
func TestFlattenAudioSubsamplingOutputMatchesReference(t *testing.T) {
ctx := setupTestContext(t)
const (
freq = 2
frames = 3
channels = 2
)
values := make([]float32, freq*frames*channels)
for c := range channels {
for t := range frames {
for f := range freq {
values[f+freq*(t+frames*c)] = float32(100*c + 10*t + f)
}
}
}
got := flattenAudioSubsamplingOutput(ctx, ctx.FromFloats(values, freq, frames, channels, 1))
ctx.Forward(got).Compute(got)
want := []float32{
0, 1, 100, 101,
10, 11, 110, 111,
20, 21, 120, 121,
}
assertCloseSlice(t, got.BackendGet(), want, 0)
}
func TestAudioDepthwiseConv1DMatchesReference(t *testing.T) {
ctx := setupTestContext(t)
xValues := make([]float32, 2*5)
for i := range xValues {
xValues[i] = float32(i)/5 - 0.7
}
kernelValues := make([]float32, 3*2)
for i := range kernelValues {
kernelValues[i] = float32(i)/3 - 0.5
}
x := ctx.FromFloats(xValues, 2, 5)
kernel := ctx.FromFloats(kernelValues, 3, 2)
got := audioDepthwiseConv1DSame(ctx, x, kernel, 1)
ctx.Forward(got).Compute(got)
want := []float32{
0.066666655, -0.5333333,
0.41666666, 0.016666688,
0.21666668, 1.0166667,
0.01666667, 2.0166664,
-0.40000004, 1.2666667,
}
assertCloseSlice(t, got.BackendGet(), want, 1e-5)
}
func TestAudioSelfAttentionMatchesReference(t *testing.T) {
ctx := setupTestContext(t)
const (
hiddenSize = 4
numHeads = 2
headDim = 2
seqLen = 3
)
xValues := make([]float32, hiddenSize*seqLen)
for i := range xValues {
xValues[i] = float32(i)/10 - 0.5
}
identity := make([]float32, hiddenSize*hiddenSize)
for i := range hiddenSize {
identity[i*hiddenSize+i] = 1
}
linear := func() *nn.Linear {
return &nn.Linear{Weight: ctx.FromFloats(identity, hiddenSize, hiddenSize)}
}
attn := &AudioSelfAttention{
Query: linear(),
Key: linear(),
Value: linear(),
Output: linear(),
RelativeKey: linear(),
BiasU: ctx.FromFloats([]float32{0.1, -0.2, 0.3, -0.4}, headDim, numHeads),
BiasV: ctx.FromFloats([]float32{-0.05, 0.07, 0.11, -0.13}, headDim, numHeads),
}
got := attn.Forward(ctx, ctx.FromFloats(xValues, hiddenSize, seqLen), seqLen, &AudioOptions{
hiddenSize: hiddenSize,
numHeads: numHeads,
headDim: headDim,
})
ctx.Forward(got).Compute(got)
want := []float32{
-0.08471569, 0.015284289, 0.05532019, 0.1553202,
-0.09135241, 0.008647568, 0.11468154, 0.21468155,
-0.019152153, 0.08084783, 0.1733382, 0.2733382,
}
assertCloseSlice(t, got.BackendGet(), want, 1e-5)
}
func assertCloseSlice(t *testing.T, got, want []float32, tolerance float64) {
t.Helper()
if len(got) != len(want) {
t.Fatalf("len(got) = %d, want %d", len(got), len(want))
}
for i := range want {
if math.Abs(float64(got[i]-want[i])) > tolerance {
t.Fatalf("got[%d] = %v, want %v\nall got: %v", i, got[i], want[i], got)
}
}
}
func TestPackPatchesCHW(t *testing.T) {
values := []float32{
0, 1, 2, 3,
4, 5, 6, 7,
8, 9, 10, 11,
12, 13, 14, 15,
100, 101, 102, 103,
104, 105, 106, 107,
108, 109, 110, 111,
112, 113, 114, 115,
}
got := packVisionPatchesCHW(values, 4, 4, 2, 2)
want := []float32{
0, 1, 4, 5, 100, 101, 104, 105,
2, 3, 6, 7, 102, 103, 106, 107,
8, 9, 12, 13, 108, 109, 112, 113,
10, 11, 14, 15, 110, 111, 114, 115,
}
if len(got) != len(want) {
t.Fatalf("len(got) = %d, want %d", len(got), len(want))
}
for i := range want {
if got[i] != want[i] {
t.Fatalf("got[%d] = %v, want %v", i, got[i], want[i])
}
}
}
func TestResizePositionEmbeddingMatchesReferenceInterpolation(t *testing.T) {
values := []float32{
0, 10,
20, 30,
}
got := resizePositionEmbedding(values, 1, 2, 2, 3, 3)
want := []float32{
0, 5, 10,
10, 15, 20,
20, 25, 30,
}
if len(got) != len(want) {
t.Fatalf("len(got) = %d, want %d", len(got), len(want))
}
for i := range want {
if got[i] != want[i] {
t.Fatalf("got[%d] = %v, want %v", i, got[i], want[i])
}
}
}
func TestDynamicImageProcessorMatchesReferencePatchBudget(t *testing.T) {
p := ImageProcessor{
imageSize: 512,
patchSize: 16,
numChannels: 3,
minNumPatches: 1024,
maxNumPatches: 13312,
projectorScale: 2,
imageMean: [3]float32{0.48145466, 0.4578275, 0.40821073},
imageStd: [3]float32{0.26862954, 0.26130258, 0.27577711},
}
img := image.NewRGBA(image.Rect(0, 0, 400, 250))
bounds := img.Bounds()
width, height := bounds.Dx(), bounds.Dy()
for y := range height {
for x := range width {
img.SetRGBA(x, y, color.RGBA{R: uint8(x), G: uint8(y), B: 128, A: 255})
}
}
tiles, err := p.ProcessImage(img)
if err != nil {
t.Fatalf("ProcessImage() error = %v", err)
}
if got, want := len(tiles), 1; got != want {
t.Fatalf("len(tiles) = %d, want %d", got, want)
}
if got, want := tiles[0].size, (image.Point{X: 672, Y: 416}); got != want {
t.Fatalf("tile size = %v, want %v", got, want)
}
if got, want := len(tiles[0].data), 3*672*416; got != want {
t.Fatalf("tile data len = %d, want %d", got, want)
}
}
func sineWAV(t *testing.T, sampleRate int, frequency float64, seconds float64) []byte {
t.Helper()
samples := int(float64(sampleRate) * seconds)
var pcm bytes.Buffer
for i := range samples {
v := int16(math.Sin(2*math.Pi*frequency*float64(i)/float64(sampleRate)) * 32767)
if err := binary.Write(&pcm, binary.LittleEndian, v); err != nil {
t.Fatal(err)
}
}
var out bytes.Buffer
out.WriteString("RIFF")
if err := binary.Write(&out, binary.LittleEndian, uint32(36+pcm.Len())); err != nil {
t.Fatal(err)
}
out.WriteString("WAVE")
out.WriteString("fmt ")
if err := binary.Write(&out, binary.LittleEndian, uint32(16)); err != nil {
t.Fatal(err)
}
if err := binary.Write(&out, binary.LittleEndian, uint16(1)); err != nil {
t.Fatal(err)
}
if err := binary.Write(&out, binary.LittleEndian, uint16(1)); err != nil {
t.Fatal(err)
}
if err := binary.Write(&out, binary.LittleEndian, uint32(sampleRate)); err != nil {
t.Fatal(err)
}
if err := binary.Write(&out, binary.LittleEndian, uint32(sampleRate*2)); err != nil {
t.Fatal(err)
}
if err := binary.Write(&out, binary.LittleEndian, uint16(2)); err != nil {
t.Fatal(err)
}
if err := binary.Write(&out, binary.LittleEndian, uint16(16)); err != nil {
t.Fatal(err)
}
out.WriteString("data")
if err := binary.Write(&out, binary.LittleEndian, uint32(pcm.Len())); err != nil {
t.Fatal(err)
}
out.Write(pcm.Bytes())
return out.Bytes()
}
+348
View File
@@ -0,0 +1,348 @@
package nemotronh
import (
"math"
"sync"
"github.com/ollama/ollama/fs"
"github.com/ollama/ollama/ml"
"github.com/ollama/ollama/ml/nn"
)
const nemotronVisionBatchSize = 1
type visionPatchGrid struct {
Width int
Height int
}
type VisionPatchEmbedding struct {
*nn.Linear
}
func packVisionPatchesCHW(values []float32, width, height, channels, patchSize int) []float32 {
patchesX, patchesY := width/patchSize, height/patchSize
patchDim := channels * patchSize * patchSize
plane := width * height
patches := make([]float32, patchDim*patchesX*patchesY)
offset := 0
for py := range patchesY {
for px := range patchesX {
for c := range channels {
channelBase := c * plane
for yy := range patchSize {
rowBase := (py*patchSize + yy) * width
for xx := range patchSize {
patches[offset] = values[channelBase+rowBase+px*patchSize+xx]
offset++
}
}
}
}
}
return patches
}
func (p *VisionPatchEmbedding) ForwardPacked(ctx ml.Context, patches []float32, patchDim, numPatches int) ml.Tensor {
hiddenState := ctx.Input().FromFloats(patches, patchDim, numPatches)
hiddenState = hiddenState.Duplicate(ctx)
return p.Linear.Forward(ctx, hiddenState)
}
func (p *VisionPatchEmbedding) Forward(ctx ml.Context, pixelValues ml.Tensor, patchSize int) ml.Tensor {
// Match the RADIO patch generator's exact flattening order: patches are laid
// out token-major with each token packed as channel, then patch-row, then
// patch-col. This is more explicit than the prior IM2Col path and likely
// slower, but it avoids backend-specific packing differences that caused the
// converted patch embedder to diverge badly from the reference model.
width, height, channels := pixelValues.Dim(0), pixelValues.Dim(1), pixelValues.Dim(2)
patchesX, patchesY := width/patchSize, height/patchSize
patchDim := channels * patchSize * patchSize
values := pixelValues.BackendGet()
return p.ForwardPacked(ctx, packVisionPatchesCHW(values, width, height, channels, patchSize), patchDim, patchesX*patchesY)
}
type VisionSelfAttention struct {
Query *nn.Linear `gguf:"attn_q"`
Key *nn.Linear `gguf:"attn_k"`
Value *nn.Linear `gguf:"attn_v"`
Output *nn.Linear `gguf:"attn_out"`
}
func (sa *VisionSelfAttention) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *VisionOptions) ml.Tensor {
headDim := opts.hiddenSize / opts.numHeads
query := sa.Query.Forward(ctx, hiddenState)
key := sa.Key.Forward(ctx, hiddenState)
value := sa.Value.Forward(ctx, hiddenState)
query = query.Reshape(ctx, headDim, opts.numHeads, query.Dim(1), nemotronVisionBatchSize)
key = key.Reshape(ctx, headDim, opts.numHeads, key.Dim(1), nemotronVisionBatchSize)
value = value.Reshape(ctx, headDim, opts.numHeads, value.Dim(1), nemotronVisionBatchSize)
attention := nn.Attention(ctx, query, key, value, 1.0/math.Sqrt(float64(headDim)), nil)
attention = attention.Reshape(ctx, opts.hiddenSize, attention.Dim(2), nemotronVisionBatchSize)
return sa.Output.Forward(ctx, attention)
}
type VisionMLP struct {
Up *nn.Linear `gguf:"ffn_up"`
Down *nn.Linear `gguf:"ffn_down"`
}
func (mlp *VisionMLP) Forward(ctx ml.Context, hiddenState ml.Tensor) ml.Tensor {
return mlp.Down.Forward(ctx, mlp.Up.Forward(ctx, hiddenState).GELU(ctx))
}
type VisionEncoderLayer struct {
LayerNorm1 *nn.LayerNorm `gguf:"ln1"`
SelfAttention *VisionSelfAttention
LayerNorm2 *nn.LayerNorm `gguf:"ln2"`
MLP *VisionMLP
}
func (l *VisionEncoderLayer) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *VisionOptions) ml.Tensor {
residual := hiddenState
hiddenState = l.LayerNorm1.Forward(ctx, hiddenState, opts.eps)
hiddenState = l.SelfAttention.Forward(ctx, hiddenState, opts)
hiddenState = hiddenState.Add(ctx, residual)
residual = hiddenState
hiddenState = l.LayerNorm2.Forward(ctx, hiddenState, opts.eps)
hiddenState = l.MLP.Forward(ctx, hiddenState)
return hiddenState.Add(ctx, residual)
}
type VisionOptions struct {
hiddenSize int
numHeads int
imageSize int
patchSize int
eps float32
}
type VisionModel struct {
PatchEmbedding *VisionPatchEmbedding `gguf:"patch_embd"`
PositionEmbedding ml.Tensor `gguf:"position_embd"`
ClassEmbedding ml.Tensor `gguf:"cls_embd"`
Layers []VisionEncoderLayer `gguf:"blk"`
*VisionOptions
resizedPositionEmbeddingsMu sync.Mutex
resizedPositionEmbeddings map[visionPatchGrid][]float32
}
func (m *VisionModel) Forward(ctx ml.Context, pixelValues ml.Tensor, patches visionPatchGrid) ml.Tensor {
numPatches := patches.Width * patches.Height
hiddenState := m.PatchEmbedding.Forward(ctx, pixelValues, m.patchSize)
return m.forwardPatchEmbeddings(ctx, hiddenState, patches, numPatches)
}
func (m *VisionModel) ForwardPacked(ctx ml.Context, patchValues []float32, patches visionPatchGrid) ml.Tensor {
numPatches := patches.Width * patches.Height
patchDim := 0
if numPatches > 0 {
patchDim = len(patchValues) / numPatches
}
hiddenState := m.PatchEmbedding.ForwardPacked(ctx, patchValues, patchDim, numPatches)
return m.forwardPatchEmbeddings(ctx, hiddenState, patches, numPatches)
}
func (m *VisionModel) forwardPatchEmbeddings(ctx ml.Context, hiddenState ml.Tensor, patches visionPatchGrid, numPatches int) ml.Tensor {
if m.PositionEmbedding != nil {
positionEmbeddings := m.positionEmbeddings(ctx, hiddenState, patches, numPatches)
hiddenState = hiddenState.Add(ctx, positionEmbeddings)
}
if m.ClassEmbedding != nil {
numPrefixTokens := m.ClassEmbedding.Dim(1)
classEmbeddings := m.ClassEmbedding.Cast(ctx, hiddenState.DType())
classEmbeddings = classEmbeddings.Reshape(ctx, classEmbeddings.Dim(0), numPrefixTokens, 1)
hiddenState = classEmbeddings.Concat(ctx, hiddenState, 1)
}
for _, layer := range m.Layers {
hiddenState = layer.Forward(ctx, hiddenState, m.VisionOptions)
}
if m.ClassEmbedding != nil {
hiddenState = hiddenState.Slice(ctx, 1, m.ClassEmbedding.Dim(1), hiddenState.Dim(1), 1)
}
return hiddenState.Reshape(ctx, hiddenState.Dim(0), hiddenState.Dim(1))
}
func (m *VisionModel) positionEmbeddings(ctx ml.Context, hiddenState ml.Tensor, patches visionPatchGrid, numPatches int) ml.Tensor {
posTokens := m.PositionEmbedding.Dim(1)
source := int(math.Sqrt(float64(posTokens)))
positionEmbeddings := m.PositionEmbedding.Cast(ctx, hiddenState.DType())
if !(source > 0 && source*source == posTokens && (source != patches.Width || source != patches.Height)) {
if positionEmbeddings.Dim(1) > numPatches {
positionEmbeddings = positionEmbeddings.Slice(ctx, 1, 0, numPatches, 1)
}
return positionEmbeddings
}
if cached, ok := m.cachePositionEmbeddings(ctx, hiddenState.Dim(0), patches); ok {
return ctx.Input().FromFloats(cached, hiddenState.Dim(0), numPatches)
}
// Runner fit/reserve builds worst-case multimodal graphs before weights are
// loaded, so the align-corners CPU cache path cannot materialize source
// values there. Fall back to a graph-only bilinear resize for reservation;
// the loaded inference path above still uses the cached align-corners data.
positionEmbeddings = positionEmbeddings.Reshape(ctx, -1, source, source)
positionEmbeddings = positionEmbeddings.Permute(ctx, 2, 0, 1, 3).Contiguous(ctx)
positionEmbeddings = positionEmbeddings.Interpolate(ctx, [4]int{
patches.Width,
patches.Height,
hiddenState.Dim(0),
1,
}, ml.SamplingModeBilinear)
positionEmbeddings = positionEmbeddings.Permute(ctx, 1, 2, 0, 3)
return positionEmbeddings.Contiguous(ctx, -1, patches.Width*patches.Height)
}
func (m *VisionModel) cachePositionEmbeddings(ctx ml.Context, hidden int, patches visionPatchGrid) ([]float32, bool) {
m.resizedPositionEmbeddingsMu.Lock()
cached := m.resizedPositionEmbeddings[patches]
m.resizedPositionEmbeddingsMu.Unlock()
if cached != nil {
return cached, true
}
if len(m.PositionEmbedding.Bytes()) == 0 {
return nil, false
}
posTokens := m.PositionEmbedding.Dim(1)
source := int(math.Sqrt(float64(posTokens)))
positionEmbeddingsF32 := m.PositionEmbedding.Cast(ctx, ml.DTypeF32)
ctx.Forward(positionEmbeddingsF32).Compute(positionEmbeddingsF32)
// RADIO eval-time CPE uses bilinear interpolation with align_corners=false.
// Cache a CPU-resized token-major embedding here for correctness first. This
// is likely slower than a native graph path and should be revisited if this
// precision vs speed tradeoff is not worthwhile.
cached = resizePositionEmbedding(positionEmbeddingsF32.Floats(), hidden, source, source, patches.Width, patches.Height)
m.resizedPositionEmbeddingsMu.Lock()
if m.resizedPositionEmbeddings == nil {
m.resizedPositionEmbeddings = make(map[visionPatchGrid][]float32)
}
if existing := m.resizedPositionEmbeddings[patches]; existing != nil {
cached = existing
} else {
m.resizedPositionEmbeddings[patches] = cached
}
m.resizedPositionEmbeddingsMu.Unlock()
return cached, true
}
func resizePositionEmbedding(values []float32, hidden, sourceWidth, sourceHeight, targetWidth, targetHeight int) []float32 {
out := make([]float32, hidden*targetWidth*targetHeight)
scaleX := float64(sourceWidth) / float64(targetWidth)
scaleY := float64(sourceHeight) / float64(targetHeight)
for oy := range targetHeight {
srcY := scaleY*(float64(oy)+0.5) - 0.5
y0 := int(math.Floor(srcY))
y1 := min(y0+1, sourceHeight-1)
wy := float32(srcY - float64(y0))
y0 = max(y0, 0)
for ox := range targetWidth {
srcX := scaleX*(float64(ox)+0.5) - 0.5
x0 := int(math.Floor(srcX))
x1 := min(x0+1, sourceWidth-1)
wx := float32(srcX - float64(x0))
x0 = max(x0, 0)
t00 := (y0*sourceWidth + x0) * hidden
t01 := (y0*sourceWidth + x1) * hidden
t10 := (y1*sourceWidth + x0) * hidden
t11 := (y1*sourceWidth + x1) * hidden
dst := (oy*targetWidth + ox) * hidden
for h := range hidden {
v00 := values[t00+h]
v01 := values[t01+h]
v10 := values[t10+h]
v11 := values[t11+h]
top := v00 + (v01-v00)*wx
bot := v10 + (v11-v10)*wx
out[dst+h] = top + (bot-top)*wy
}
}
}
return out
}
func newVisionModel(c fs.Config) *VisionModel {
return &VisionModel{
Layers: make([]VisionEncoderLayer, c.Uint("vision.block_count", 32)),
VisionOptions: &VisionOptions{
hiddenSize: int(c.Uint("vision.embedding_length", 1280)),
numHeads: int(c.Uint("vision.attention.head_count", 16)),
imageSize: int(c.Uint("vision.image_size", 512)),
patchSize: int(c.Uint("vision.patch_size", 16)),
eps: c.Float("vision.attention.layer_norm_epsilon", 1e-6),
},
}
}
type MultiModalProjector struct {
Norm *nn.RMSNorm `gguf:"norm"`
Linear1 *nn.Linear `gguf:"1"`
Linear2 *nn.Linear `gguf:"2"`
scaleFactor int
}
func (p *MultiModalProjector) Forward(ctx ml.Context, visionOutputs ml.Tensor, patches visionPatchGrid) ml.Tensor {
scaleFactor := max(p.scaleFactor, 1)
// The reference projector first pixel-shuffles the vision grid with
// downsample_ratio=0.5 before applying the RMSNorm/MLP. Preserve that exact
// v2 packing order here rather than flattening 2x2 neighborhoods via IM2Col.
merged := pixelShuffleVisionOutputs(ctx, visionOutputs, patches, scaleFactor)
merged = p.Norm.Forward(ctx, merged, 1e-5)
merged = p.Linear1.Forward(ctx, merged)
merged = merged.RELU(ctx)
merged = merged.Mul(ctx, merged)
return p.Linear2.Forward(ctx, merged)
}
func pixelShuffleVisionOutputs(ctx ml.Context, visionOutputs ml.Tensor, patches visionPatchGrid, scaleFactor int) ml.Tensor {
hiddenSize := visionOutputs.Dim(0)
scaleFactor = max(scaleFactor, 1)
merged := visionOutputs.Reshape(ctx, hiddenSize, patches.Width, patches.Height, 1)
width := patches.Width / scaleFactor
height := patches.Height / scaleFactor
channels := hiddenSize * scaleFactor
merged = merged.Reshape(ctx, channels, width, patches.Height, 1)
merged = merged.Reshape(ctx, channels, width, scaleFactor, height)
merged = merged.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
return merged.Reshape(ctx, channels*scaleFactor, width*height, 1)
}
func newMultiModalProjector(c fs.Config) *MultiModalProjector {
return &MultiModalProjector{
scaleFactor: int(c.Uint("vision.projector.scale_factor", 2)),
}
}
+328
View File
@@ -0,0 +1,328 @@
package nemotronh
import (
"encoding/binary"
"fmt"
"math"
"math/cmplx"
)
const (
parakeetHopLength = 160
parakeetNFFT = 512
parakeetWinLength = 400
parakeetPreemphasis = 0.97
parakeetLogZeroGuardValue = 1.0 / (1 << 24)
parakeetNormalizeEps = 1e-5
)
func isAudioData(data []byte) bool {
return len(data) >= 12 && string(data[:4]) == "RIFF" && string(data[8:12]) == "WAVE"
}
func decodeWAV(data []byte, targetSampleRate int) ([]float32, error) {
if len(data) < 12 {
return nil, fmt.Errorf("WAV file too short")
}
if !isAudioData(data) {
return nil, fmt.Errorf("not a WAV file")
}
var audioFormat uint16
var numChannels, sampleRate, bitsPerSample int
var audioData []byte
foundFmt := false
offset := 12
for offset+8 <= len(data) {
chunkID := string(data[offset : offset+4])
chunkSize := int(binary.LittleEndian.Uint32(data[offset+4 : offset+8]))
chunkEnd := min(offset+8+chunkSize, len(data))
chunkData := data[offset+8 : chunkEnd]
switch chunkID {
case "fmt ":
if len(chunkData) < 16 {
return nil, fmt.Errorf("fmt chunk too short")
}
audioFormat = binary.LittleEndian.Uint16(chunkData[0:2])
numChannels = int(binary.LittleEndian.Uint16(chunkData[2:4]))
sampleRate = int(binary.LittleEndian.Uint32(chunkData[4:8]))
bitsPerSample = int(binary.LittleEndian.Uint16(chunkData[14:16]))
if audioFormat == 0xfffe && len(chunkData) >= 26 {
audioFormat = binary.LittleEndian.Uint16(chunkData[24:26])
}
foundFmt = true
case "data":
audioData = chunkData
}
offset += 8 + chunkSize
if chunkSize%2 != 0 {
offset++
}
}
if !foundFmt {
return nil, fmt.Errorf("no fmt chunk found in WAV file")
}
if audioFormat != 1 && audioFormat != 3 {
return nil, fmt.Errorf("unsupported WAV format: %d (need PCM=1 or float=3)", audioFormat)
}
if audioData == nil {
return nil, fmt.Errorf("no data chunk found in WAV file")
}
if numChannels <= 0 {
return nil, fmt.Errorf("invalid WAV channel count: %d", numChannels)
}
samples := decodeWAVSamples(audioData, audioFormat, bitsPerSample, numChannels)
if sampleRate != targetSampleRate {
samples = resampleLinear(samples, sampleRate, targetSampleRate)
}
return samples, nil
}
func decodeWAVSamples(data []byte, format uint16, bits, channels int) []float32 {
bytesPerSample := bits / 8
if bytesPerSample <= 0 || channels <= 0 {
return nil
}
totalSamples := len(data) / (bytesPerSample * channels)
mono := make([]float32, totalSamples)
for i := range totalSamples {
var sum float64
for ch := range channels {
off := (i*channels + ch) * bytesPerSample
if off+bytesPerSample > len(data) {
break
}
switch {
case format == 1 && bits == 16:
v := int16(binary.LittleEndian.Uint16(data[off : off+2]))
sum += float64(v) / 32768.0
case format == 1 && bits == 32:
v := int32(binary.LittleEndian.Uint32(data[off : off+4]))
sum += float64(v) / 2147483648.0
case format == 1 && bits == 24:
v := int32(data[off]) | int32(data[off+1])<<8 | int32(data[off+2])<<16
if v&0x800000 != 0 {
v |= ^0xffffff
}
sum += float64(v) / 8388608.0
case format == 3 && bits == 32:
sum += float64(math.Float32frombits(binary.LittleEndian.Uint32(data[off : off+4])))
case format == 1 && bits == 8:
sum += (float64(data[off]) - 128.0) / 128.0
}
}
mono[i] = float32(sum / float64(channels))
}
return mono
}
func resampleLinear(samples []float32, fromRate, toRate int) []float32 {
if fromRate <= 0 || toRate <= 0 || len(samples) == 0 {
return samples
}
n := int(float64(len(samples)) / float64(fromRate) * float64(toRate))
if n <= 1 {
return slicesCloneOne(samples)
}
out := make([]float32, n)
for i := range n {
pos := float64(i) * float64(len(samples)-1) / float64(n-1)
idx := int(pos)
frac := float32(pos - float64(idx))
if idx+1 < len(samples) {
out[i] = samples[idx]*(1-frac) + samples[idx+1]*frac
} else {
out[i] = samples[idx]
}
}
return out
}
func slicesCloneOne(samples []float32) []float32 {
if len(samples) == 0 {
return nil
}
return []float32{samples[0]}
}
func computeParakeetMelSpectrogram(samples []float32, extractor *AudioFeatureExtractor, opts *AudioOptions) ([]float32, int, int, error) {
if len(samples) == 0 {
return nil, 0, 0, fmt.Errorf("audio too short to encode")
}
if opts == nil {
opts = defaultAudioOptions()
}
melBins := opts.melBins
freqBins := parakeetNFFT/2 + 1
window, melFilters := extractor.windowAndFilters(melBins, freqBins, opts.sampleRate)
if len(window) != parakeetWinLength {
return nil, 0, 0, fmt.Errorf("invalid Parakeet window length: %d", len(window))
}
if len(melFilters) != melBins*freqBins {
return nil, 0, 0, fmt.Errorf("invalid Parakeet mel filter shape: %d", len(melFilters))
}
emphasized := make([]float32, len(samples))
emphasized[0] = samples[0]
for i := 1; i < len(samples); i++ {
emphasized[i] = samples[i] - parakeetPreemphasis*samples[i-1]
}
frames := len(samples)/parakeetHopLength + 1
validFrames := max(1, len(samples)/parakeetHopLength)
if validFrames > frames {
validFrames = frames
}
result := make([]float32, frames*melBins)
fftInput := make([]complex128, parakeetNFFT)
winOffset := (parakeetNFFT - parakeetWinLength) / 2
centerPad := parakeetNFFT / 2
for frame := range frames {
for i := range parakeetNFFT {
fftInput[i] = 0
}
for i := range parakeetWinLength {
src := frame*parakeetHopLength + i + winOffset - centerPad
if src >= 0 && src < len(emphasized) {
fftInput[i+winOffset] = complex(float64(emphasized[src])*float64(window[i]), 0)
}
}
fft(fftInput)
for mel := range melBins {
var v float64
filterOffset := mel * freqBins
for freq := range freqBins {
mag := cmplx.Abs(fftInput[freq])
v += float64(melFilters[filterOffset+freq]) * mag * mag
}
result[frame*melBins+mel] = float32(math.Log(v + parakeetLogZeroGuardValue))
}
}
for mel := range melBins {
var sum float64
for frame := range validFrames {
sum += float64(result[frame*melBins+mel])
}
mean := sum / float64(validFrames)
var variance float64
for frame := range validFrames {
d := float64(result[frame*melBins+mel]) - mean
variance += d * d
}
denom := max(1, validFrames-1)
std := math.Sqrt(variance / float64(denom))
for frame := range frames {
idx := frame*melBins + mel
if frame >= validFrames {
result[idx] = 0
continue
}
result[idx] = float32((float64(result[idx]) - mean) / (std + parakeetNormalizeEps))
}
}
return result, frames, validFrames, nil
}
func defaultParakeetWindow() []float32 {
window := make([]float32, parakeetWinLength)
for i := range window {
window[i] = float32(0.5 - 0.5*math.Cos(2*math.Pi*float64(i)/float64(parakeetWinLength-1)))
}
return window
}
func buildSlaneyMelFilterBank(numFreqBins, numMels int, sampleRate int) []float32 {
hzToMel := func(f float64) float64 {
if f < 1000 {
return 3 * f / 200
}
return 15 + math.Log(f/1000)*27/math.Log(6.4)
}
melToHz := func(m float64) float64 {
if m < 15 {
return 200 * m / 3
}
return 1000 * math.Exp(math.Log(6.4)*(m-15)/27)
}
minMel := hzToMel(0)
maxMel := hzToMel(float64(sampleRate) / 2)
mels := make([]float64, numMels+2)
freqs := make([]float64, numMels+2)
for i := range mels {
mels[i] = minMel + (maxMel-minMel)*float64(i)/float64(numMels+1)
freqs[i] = melToHz(mels[i])
}
fftFreqs := make([]float64, numFreqBins)
for i := range fftFreqs {
fftFreqs[i] = float64(i) * float64(sampleRate) / float64(parakeetNFFT)
}
filters := make([]float32, numMels*numFreqBins)
for mel := range numMels {
left, center, right := freqs[mel], freqs[mel+1], freqs[mel+2]
enorm := 2.0 / (right - left)
for freq, fftFreq := range fftFreqs {
var lower, upper float64
if center > left {
lower = (fftFreq - left) / (center - left)
}
if right > center {
upper = (right - fftFreq) / (right - center)
}
v := math.Max(0, math.Min(lower, upper))
filters[mel*numFreqBins+freq] = float32(v * enorm)
}
}
return filters
}
func fft(x []complex128) {
n := len(x)
if n <= 1 {
return
}
j := 0
for i := 1; i < n; i++ {
bit := n >> 1
for j&bit != 0 {
j ^= bit
bit >>= 1
}
j ^= bit
if i < j {
x[i], x[j] = x[j], x[i]
}
}
for size := 2; size <= n; size <<= 1 {
halfSize := size / 2
w := complex(math.Cos(2*math.Pi/float64(size)), -math.Sin(2*math.Pi/float64(size)))
for start := 0; start < n; start += size {
wn := complex(1, 0)
for k := range halfSize {
t := wn * x[start+k+halfSize]
x[start+k+halfSize] = x[start+k] - t
x[start+k] = x[start+k] + t
wn *= w
}
}
}
}
+498
View File
@@ -0,0 +1,498 @@
package parsers
import (
"encoding/json"
"fmt"
"regexp"
"strings"
"unicode"
"github.com/ollama/ollama/api"
)
const (
lagunaThinkingOpenTag = "<think>"
lagunaThinkingCloseTag = "</think>"
lagunaToolCallOpenTag = "<tool_call>"
lagunaToolCallCloseTag = "</tool_call>"
lagunaUserOpenTag = "<user>"
lagunaUserCloseTag = "</user>"
)
type lagunaParserState int
const (
lagunaParserStateThinking lagunaParserState = iota
lagunaParserStateContent
lagunaParserStateTool
)
type LagunaParser struct {
state lagunaParserState
buffer strings.Builder
tools []api.Tool
callIndex int
thinkingEnabled bool
thinkingSuppressed bool
allowLeadingThinkOpen bool
}
func (p *LagunaParser) HasToolSupport() bool {
return true
}
func (p *LagunaParser) HasThinkingSupport() bool {
return true
}
func (p *LagunaParser) Init(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue) []api.Tool {
p.tools = tools
p.callIndex = 0
p.buffer.Reset()
p.thinkingEnabled = thinkValue == nil || thinkValue.Bool()
p.thinkingSuppressed = thinkValue != nil && !thinkValue.Bool()
p.state = lagunaParserStateContent
p.allowLeadingThinkOpen = false
return tools
}
func (p *LagunaParser) Add(s string, done bool) (content string, thinking string, calls []api.ToolCall, err error) {
p.buffer.WriteString(s)
var contentSB, thinkingSB strings.Builder
for {
progress := false
switch p.state {
case lagunaParserStateThinking:
progress, thinking = p.consumeThinking(done)
if p.thinkingEnabled {
thinkingSB.WriteString(thinking)
}
case lagunaParserStateContent:
var parsedCalls []api.ToolCall
progress, content, parsedCalls, err = p.consumeContent(done)
if err != nil {
return "", "", nil, err
}
contentSB.WriteString(content)
calls = append(calls, parsedCalls...)
case lagunaParserStateTool:
var call api.ToolCall
progress, call, err = p.consumeTool(done)
if err != nil {
return "", "", nil, err
}
if progress {
calls = append(calls, call)
}
}
if !progress {
break
}
}
return contentSB.String(), thinkingSB.String(), calls, nil
}
func (p *LagunaParser) consumeThinking(done bool) (bool, string) {
acc := p.buffer.String()
if p.allowLeadingThinkOpen {
trimmed := strings.TrimLeftFunc(acc, unicode.IsSpace)
if strings.HasPrefix(trimmed, lagunaThinkingOpenTag) {
p.buffer.Reset()
p.buffer.WriteString(strings.TrimLeftFunc(strings.TrimPrefix(trimmed, lagunaThinkingOpenTag), unicode.IsSpace))
p.allowLeadingThinkOpen = false
return true, ""
}
if strings.HasPrefix(lagunaThinkingOpenTag, trimmed) && !done {
return false, ""
}
p.allowLeadingThinkOpen = false
}
if idx := strings.Index(acc, lagunaThinkingCloseTag); idx != -1 {
thinking := acc[:idx]
after := strings.TrimLeftFunc(acc[idx+len(lagunaThinkingCloseTag):], unicode.IsSpace)
p.buffer.Reset()
p.buffer.WriteString(after)
p.state = lagunaParserStateContent
return true, thinking
}
if idx := strings.Index(acc, lagunaToolCallOpenTag); idx != -1 {
thinking := strings.TrimRightFunc(acc[:idx], unicode.IsSpace)
after := acc[idx+len(lagunaToolCallOpenTag):]
p.buffer.Reset()
p.buffer.WriteString(after)
p.state = lagunaParserStateTool
return true, thinking
}
if done {
p.buffer.Reset()
p.state = lagunaParserStateContent
return acc != "", acc
}
overlapLen := max(overlap(acc, lagunaThinkingCloseTag), overlap(acc, lagunaToolCallOpenTag))
trailingLen := trailingWhitespaceLen(acc)
keep := max(overlapLen, trailingLen)
if keep > 0 && keep < len(acc) {
emit := acc[:len(acc)-keep]
p.buffer.Reset()
p.buffer.WriteString(acc[len(acc)-keep:])
return emit != "", emit
}
return false, ""
}
func (p *LagunaParser) consumeContent(done bool) (bool, string, []api.ToolCall, error) {
acc := p.buffer.String()
if p.thinkingEnabled || p.thinkingSuppressed {
if idx := strings.Index(acc, lagunaThinkingOpenTag); idx != -1 {
content := acc[:idx]
after := strings.TrimLeftFunc(acc[idx+len(lagunaThinkingOpenTag):], unicode.IsSpace)
p.buffer.Reset()
p.buffer.WriteString(after)
p.state = lagunaParserStateThinking
p.allowLeadingThinkOpen = false
return true, content, nil, nil
}
if !done {
overlapLen := overlap(acc, lagunaThinkingOpenTag)
if overlapLen > 0 && overlapLen < len(acc) {
content := acc[:len(acc)-overlapLen]
p.buffer.Reset()
p.buffer.WriteString(acc[len(acc)-overlapLen:])
return content != "", content, nil, nil
}
}
}
if p.thinkingEnabled {
trimmed := strings.TrimLeftFunc(acc, unicode.IsSpace)
if strings.HasPrefix(trimmed, lagunaThinkingCloseTag) {
p.buffer.Reset()
p.buffer.WriteString(strings.TrimLeftFunc(strings.TrimPrefix(trimmed, lagunaThinkingCloseTag), unicode.IsSpace))
return true, "", nil, nil
}
if strings.HasPrefix(lagunaThinkingCloseTag, trimmed) && !done {
return false, "", nil, nil
}
}
if p.thinkingSuppressed {
trimmed := strings.TrimLeftFunc(acc, unicode.IsSpace)
if strings.HasPrefix(trimmed, lagunaThinkingCloseTag) {
p.buffer.Reset()
p.buffer.WriteString(strings.TrimLeftFunc(strings.TrimPrefix(trimmed, lagunaThinkingCloseTag), unicode.IsSpace))
return true, "", nil, nil
}
if strings.HasPrefix(lagunaThinkingCloseTag, trimmed) && !done {
return false, "", nil, nil
}
}
if idx := strings.Index(acc, lagunaToolCallOpenTag); idx != -1 {
content := strings.TrimRightFunc(acc[:idx], unicode.IsSpace)
after := acc[idx+len(lagunaToolCallOpenTag):]
p.buffer.Reset()
p.buffer.WriteString(after)
p.state = lagunaParserStateTool
return true, content, nil, nil
}
if idx := strings.Index(acc, lagunaUserOpenTag); idx != -1 && len(p.tools) > 0 {
before := strings.TrimRightFunc(acc[:idx], unicode.IsSpace)
afterOpen := acc[idx+len(lagunaUserOpenTag):]
if closeIdx := strings.Index(afterOpen, lagunaUserCloseTag); closeIdx != -1 {
raw := afterOpen[:closeIdx]
if call, ok := p.parseToolAlias(raw); ok {
after := strings.TrimLeftFunc(afterOpen[closeIdx+len(lagunaUserCloseTag):], unicode.IsSpace)
p.buffer.Reset()
p.buffer.WriteString(after)
return true, before, []api.ToolCall{call}, nil
}
} else if !done {
if idx > 0 {
p.buffer.Reset()
p.buffer.WriteString(acc[idx:])
return true, before, nil, nil
}
return false, "", nil, nil
}
}
if len(p.tools) > 0 {
if progress, content, call, ok, err := p.consumeStandaloneJSONTool(done); ok || err != nil {
if err != nil {
return false, "", nil, err
}
if progress {
return true, content, []api.ToolCall{call}, nil
}
return false, "", nil, nil
}
}
if done {
p.buffer.Reset()
return acc != "", acc, nil, nil
}
overlapLen := max(overlap(acc, lagunaToolCallOpenTag), overlap(acc, lagunaUserOpenTag))
if p.thinkingEnabled || p.thinkingSuppressed {
overlapLen = max(overlapLen, overlap(acc, lagunaThinkingOpenTag))
}
if p.thinkingSuppressed {
overlapLen = max(overlapLen, overlap(acc, lagunaThinkingCloseTag))
}
trailingLen := trailingWhitespaceLen(acc)
keep := max(overlapLen, trailingLen)
if keep > 0 && keep < len(acc) {
emit := acc[:len(acc)-keep]
p.buffer.Reset()
p.buffer.WriteString(acc[len(acc)-keep:])
return emit != "", emit, nil, nil
}
if keep == 0 && acc != "" {
p.buffer.Reset()
return true, acc, nil, nil
}
return false, "", nil, nil
}
func (p *LagunaParser) consumeStandaloneJSONTool(done bool) (progress bool, content string, call api.ToolCall, ok bool, err error) {
acc := p.buffer.String()
jsonIdx := strings.Index(acc, "{")
if jsonIdx == -1 {
return false, "", api.ToolCall{}, false, nil
}
before := strings.TrimRightFunc(acc[:jsonIdx], unicode.IsSpace)
raw := strings.TrimLeftFunc(acc[jsonIdx:], unicode.IsSpace)
if !lagunaLooksLikeJSONToolCall(raw, done) {
return false, "", api.ToolCall{}, false, nil
}
if !done && !json.Valid([]byte(strings.TrimSpace(raw))) {
if before != "" {
p.buffer.Reset()
p.buffer.WriteString(acc[jsonIdx:])
return true, before, api.ToolCall{}, true, nil
}
return false, "", api.ToolCall{}, true, nil
}
call, err = parseLagunaToolCall(raw, p.tools)
if err != nil {
return false, "", api.ToolCall{}, true, err
}
call.Function.Index = p.callIndex
p.callIndex++
p.buffer.Reset()
p.state = lagunaParserStateContent
return true, before, call, true, nil
}
func lagunaLooksLikeJSONToolCall(raw string, done bool) bool {
trimmed := strings.TrimLeftFunc(raw, unicode.IsSpace)
if !strings.HasPrefix(trimmed, "{") {
return false
}
if strings.Contains(trimmed, `"name"`) || strings.Contains(trimmed, `"arguments"`) {
return true
}
if done {
return false
}
return strings.HasPrefix(trimmed, `{"`) || strings.HasPrefix(trimmed, "{\n") || strings.HasPrefix(trimmed, "{\r\n")
}
func (p *LagunaParser) parseToolAlias(raw string) (api.ToolCall, bool) {
raw = cleanLagunaToolCallRaw(raw)
name, ok := lagunaToolCallName(raw)
if !ok {
return api.ToolCall{}, false
}
if _, ok := lagunaResolveToolName(name, p.tools); !ok {
return api.ToolCall{}, false
}
call, err := parseLagunaToolCall(raw, p.tools)
if err != nil {
return api.ToolCall{}, false
}
call.Function.Index = p.callIndex
p.callIndex++
return call, true
}
func lagunaResolveToolName(name string, tools []api.Tool) (string, bool) {
for i := range tools {
if tools[i].Function.Name == name {
return name, true
}
}
aliases := map[string]string{
"read_file": "read",
"write_file": "write",
"edit_file": "edit",
"web_fetch": "webfetch",
}
if alias, ok := aliases[name]; ok {
for i := range tools {
if tools[i].Function.Name == alias {
return alias, true
}
}
}
return name, false
}
func cleanLagunaToolCallRaw(raw string) string {
raw = strings.TrimSpace(raw)
for strings.HasPrefix(raw, lagunaToolCallOpenTag) {
raw = strings.TrimSpace(strings.TrimPrefix(raw, lagunaToolCallOpenTag))
}
if idx := strings.Index(raw, lagunaToolCallCloseTag); idx != -1 {
raw = strings.TrimSpace(raw[:idx])
}
if idx := strings.Index(raw, lagunaToolCallOpenTag); idx != -1 {
before := strings.TrimSpace(raw[:idx])
if before != "" {
return before
}
raw = strings.TrimSpace(raw[idx+len(lagunaToolCallOpenTag):])
}
return raw
}
func lagunaToolCallName(raw string) (string, bool) {
raw = cleanLagunaToolCallRaw(raw)
if strings.HasPrefix(raw, "{") {
var parsed struct {
Name string `json:"name"`
}
if err := json.Unmarshal([]byte(raw), &parsed); err != nil {
return "", false
}
name := strings.TrimSpace(parsed.Name)
return name, name != ""
}
nameEnd := strings.Index(raw, "<arg_key>")
if nameEnd < 0 {
nameEnd = strings.Index(raw, "{")
}
if nameEnd < 0 {
nameEnd = strings.IndexAny(raw, "\r\n")
}
if nameEnd < 0 {
nameEnd = len(raw)
}
name := strings.TrimSpace(raw[:nameEnd])
return name, name != ""
}
func (p *LagunaParser) consumeTool(done bool) (bool, api.ToolCall, error) {
acc := p.buffer.String()
if idx := strings.Index(acc, lagunaToolCallCloseTag); idx != -1 {
raw := acc[:idx]
after := strings.TrimLeftFunc(acc[idx+len(lagunaToolCallCloseTag):], unicode.IsSpace)
p.buffer.Reset()
p.buffer.WriteString(after)
p.state = lagunaParserStateContent
call, err := parseLagunaToolCall(raw, p.tools)
if err != nil {
return false, api.ToolCall{}, err
}
call.Function.Index = p.callIndex
p.callIndex++
return true, call, nil
}
if done && strings.TrimSpace(acc) != "" {
p.buffer.Reset()
p.state = lagunaParserStateContent
call, err := parseLagunaToolCall(acc, p.tools)
if err != nil {
return false, api.ToolCall{}, err
}
call.Function.Index = p.callIndex
p.callIndex++
return true, call, nil
}
return false, api.ToolCall{}, nil
}
var lagunaArgRE = regexp.MustCompile(`(?s)<arg_key>(.*?)</arg_key>\s*<arg_value>(.*?)</arg_value>`)
func parseLagunaToolCall(raw string, tools []api.Tool) (api.ToolCall, error) {
raw = cleanLagunaToolCallRaw(raw)
if strings.HasPrefix(raw, "{") {
var parsed struct {
Name string `json:"name"`
Arguments api.ToolCallFunctionArguments `json:"arguments"`
}
if err := json.Unmarshal([]byte(raw), &parsed); err != nil {
return api.ToolCall{}, fmt.Errorf("failed to parse Laguna JSON tool call: %w", err)
}
if parsed.Name == "" {
return api.ToolCall{}, fmt.Errorf("empty Laguna tool call name")
}
if name, ok := lagunaResolveToolName(parsed.Name, tools); ok {
parsed.Name = name
}
return api.ToolCall{
Function: api.ToolCallFunction{
Name: parsed.Name,
Arguments: parsed.Arguments,
},
}, nil
}
nameEnd := strings.Index(raw, "<arg_key>")
name := raw
argsText := ""
if nameEnd >= 0 {
name = raw[:nameEnd]
argsText = raw[nameEnd:]
} else if jsonStart := strings.Index(raw, "{"); jsonStart >= 0 {
name = raw[:jsonStart]
argsText = raw[jsonStart:]
}
name = strings.TrimSpace(name)
if resolved, ok := lagunaResolveToolName(name, tools); ok {
name = resolved
}
var matchedTool *api.Tool
for i := range tools {
if tools[i].Function.Name == name {
matchedTool = &tools[i]
break
}
}
call := api.ToolCall{
Function: api.ToolCallFunction{
Name: name,
Arguments: api.NewToolCallFunctionArguments(),
},
}
if strings.HasPrefix(strings.TrimSpace(argsText), "{") {
if err := json.Unmarshal([]byte(strings.TrimSpace(argsText)), &call.Function.Arguments); err != nil {
return api.ToolCall{}, fmt.Errorf("failed to parse Laguna JSON tool call arguments: %w", err)
}
return call, nil
}
for _, match := range lagunaArgRE.FindAllStringSubmatch(argsText, -1) {
key := strings.TrimSpace(match[1])
value := match[2]
var paramType api.PropertyType
if matchedTool != nil && matchedTool.Function.Parameters.Properties != nil {
if prop, ok := matchedTool.Function.Parameters.Properties.Get(key); ok {
if len(prop.AnyOf) > 0 {
for _, anyOfProp := range prop.AnyOf {
paramType = append(paramType, anyOfProp.Type...)
}
} else {
paramType = prop.Type
}
}
}
call.Function.Arguments.Set(key, parseValue(value, paramType))
}
return call, nil
}
+484
View File
@@ -0,0 +1,484 @@
package parsers
import (
"testing"
"github.com/ollama/ollama/api"
)
func lagunaTestTools() []api.Tool {
props := api.NewToolPropertiesMap()
props.Set("location", api.ToolProperty{Type: api.PropertyType{"string"}})
props.Set("days", api.ToolProperty{Type: api.PropertyType{"integer"}})
return []api.Tool{{
Function: api.ToolFunction{
Name: "get_weather",
Parameters: api.ToolFunctionParameters{
Properties: props,
},
},
}}
}
func TestLagunaParserToolCall(t *testing.T) {
parser := ParserForName("laguna")
if parser == nil {
t.Fatal("expected laguna parser")
}
if !parser.HasToolSupport() || !parser.HasThinkingSupport() {
t.Fatal("laguna parser should advertise tools and thinking")
}
parser.Init(lagunaTestTools(), nil, nil)
content, thinking, calls, err := parser.Add("<tool_call>get_weather\n<arg_key>location</arg_key>\n<arg_value>Paris</arg_value>\n<arg_key>days</arg_key>\n<arg_value>3</arg_value>\n</tool_call>", true)
if err != nil {
t.Fatal(err)
}
if content != "" || thinking != "" {
t.Fatalf("content=%q thinking=%q, want empty", content, thinking)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "get_weather" {
t.Fatalf("name=%q, want get_weather", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("location"); got != "Paris" {
t.Fatalf("location=%v, want Paris", got)
}
if got, _ := calls[0].Function.Arguments.Get("days"); got != 3 {
t.Fatalf("days=%v, want 3", got)
}
}
func TestLagunaParserJSONToolCall(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, nil)
_, _, calls, err := parser.Add("<tool_call>\n{\"name\":\"get_weather\",\"arguments\":{\"location\":\"Paris\",\"days\":3}}\n</tool_call>", true)
if err != nil {
t.Fatal(err)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "get_weather" {
t.Fatalf("name=%q, want get_weather", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("location"); got != "Paris" {
t.Fatalf("location=%v, want Paris", got)
}
if got, _ := calls[0].Function.Arguments.Get("days"); got != float64(3) {
t.Fatalf("days=%v, want 3", got)
}
}
func TestLagunaParserStandaloneJSONToolCall(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, nil)
content, thinking, calls, err := parser.Add("{\"name\":\"get_weather\",\"arguments\":{\"location\":\"Paris\",\"days\":3}}", true)
if err != nil {
t.Fatal(err)
}
if content != "" || thinking != "" {
t.Fatalf("content=%q thinking=%q", content, thinking)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "get_weather" {
t.Fatalf("name=%q, want get_weather", calls[0].Function.Name)
}
}
func TestLagunaParserStandaloneJSONToolCallAfterLeadingContent(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, nil)
content, thinking, calls, err := parser.Add("Let me call the weather tool.\n{\"name\":\"get_weather\",\"arguments\":{\"location\":\"Paris\"}}", true)
if err != nil {
t.Fatal(err)
}
if content != "Let me call the weather tool." || thinking != "" {
t.Fatalf("content=%q thinking=%q", content, thinking)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "get_weather" {
t.Fatalf("name=%q, want get_weather", calls[0].Function.Name)
}
}
func TestLagunaParserStreamingStandaloneJSONToolCall(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, nil)
content, thinking, calls, err := parser.Add("{\"name\":\"get_weather\",\"arguments\":{\"location\":\"San Francisco,", false)
if err != nil {
t.Fatal(err)
}
if content != "" || thinking != "" || len(calls) != 0 {
t.Fatalf("first chunk content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
content, thinking, calls, err = parser.Add(" CA\"}}", true)
if err != nil {
t.Fatal(err)
}
if content != "" || thinking != "" || len(calls) != 1 {
t.Fatalf("second chunk content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
if calls[0].Function.Name != "get_weather" {
t.Fatalf("name=%q, want get_weather", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("location"); got != "San Francisco, CA" {
t.Fatalf("location=%v, want San Francisco, CA", got)
}
}
func TestLagunaParserNameLineJSONToolCall(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, nil)
_, _, calls, err := parser.Add("<tool_call>get_weather\n{\"location\":\"San Francisco\"}</tool_call>", true)
if err != nil {
t.Fatal(err)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "get_weather" {
t.Fatalf("name=%q, want get_weather", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("location"); got != "San Francisco" {
t.Fatalf("location=%v, want San Francisco", got)
}
}
func TestLagunaParserNormalizesCommonToolAliases(t *testing.T) {
props := api.NewToolPropertiesMap()
props.Set("path", api.ToolProperty{Type: api.PropertyType{"string"}})
tools := []api.Tool{{
Function: api.ToolFunction{
Name: "read",
Parameters: api.ToolFunctionParameters{
Properties: props,
},
},
}}
parser := ParserForName("laguna")
parser.Init(tools, nil, nil)
_, _, calls, err := parser.Add("<tool_call>\n{\"name\":\"read_file\",\"arguments\":{\"path\":\"./go.mod\"}}\n</tool_call>", true)
if err != nil {
t.Fatal(err)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "read" {
t.Fatalf("name=%q, want read", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("path"); got != "./go.mod" {
t.Fatalf("path=%v, want ./go.mod", got)
}
}
func TestLagunaParserIgnoresDuplicatedNestedToolCall(t *testing.T) {
props := api.NewToolPropertiesMap()
props.Set("name", api.ToolProperty{Type: api.PropertyType{"string"}})
tools := []api.Tool{{
Function: api.ToolFunction{
Name: "skill",
Parameters: api.ToolFunctionParameters{
Properties: props,
},
},
}}
parser := ParserForName("laguna")
parser.Init(tools, nil, nil)
_, _, calls, err := parser.Add("<tool_call>skill\n{\"name\":\"git-diff-review\"}\n<tool_call>skill\n{\"name\":\"git-diff-review\"}</tool_call>", true)
if err != nil {
t.Fatal(err)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "skill" {
t.Fatalf("name=%q, want skill", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("name"); got != "git-diff-review" {
t.Fatalf("name arg=%v, want git-diff-review", got)
}
}
func TestLagunaParserThinkingThenTool(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, &api.ThinkValue{Value: true})
content, thinking, calls, err := parser.Add("<think>Need current weather.</think>\n<tool_call>get_weather\n<arg_key>location</arg_key>\n<arg_value>SF</arg_value>\n</tool_call>", true)
if err != nil {
t.Fatal(err)
}
if content != "" {
t.Fatalf("content=%q, want empty", content)
}
if thinking != "Need current weather." {
t.Fatalf("thinking=%q, want reasoning", thinking)
}
if len(calls) != 1 || calls[0].Function.Name != "get_weather" {
t.Fatalf("unexpected calls: %#v", calls)
}
}
func TestLagunaParserUserTaggedToolAlias(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, nil)
content, thinking, calls, err := parser.Add("<user>get_weather\n<arg_key>location</arg_key>\n<arg_value>San Francisco, CA</arg_value>\n</user>", true)
if err != nil {
t.Fatal(err)
}
if content != "" || thinking != "" {
t.Fatalf("content=%q thinking=%q, want empty", content, thinking)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "get_weather" {
t.Fatalf("name=%q, want get_weather", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("location"); got != "San Francisco, CA" {
t.Fatalf("location=%v, want San Francisco, CA", got)
}
}
func TestLagunaParserUserTaggedToolAliasWithLeadingContent(t *testing.T) {
parser := ParserForName("laguna")
props := api.NewToolPropertiesMap()
props.Set("path", api.ToolProperty{Type: api.PropertyType{"string"}})
tools := []api.Tool{{
Function: api.ToolFunction{
Name: "read",
Parameters: api.ToolFunctionParameters{
Properties: props,
},
},
}}
parser.Init(tools, nil, nil)
content, thinking, calls, err := parser.Add("I'll read the file for you.\n<user>read\n<arg_key>path</arg_key>\n<arg_value>/Users/test/code/myproject/go.mod</arg_value>\n</user>", true)
if err != nil {
t.Fatal(err)
}
if content != "I'll read the file for you." || thinking != "" {
t.Fatalf("content=%q thinking=%q", content, thinking)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "read" {
t.Fatalf("name=%q, want read", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("path"); got != "/Users/test/code/myproject/go.mod" {
t.Fatalf("path=%v, want /Users/test/code/myproject/go.mod", got)
}
}
func TestLagunaParserUserTaggedJSONToolCallWithLeadingContent(t *testing.T) {
parser := ParserForName("laguna")
props := api.NewToolPropertiesMap()
props.Set("command", api.ToolProperty{Type: api.PropertyType{"string"}})
tools := []api.Tool{{
Function: api.ToolFunction{
Name: "bash",
Parameters: api.ToolFunctionParameters{
Properties: props,
},
},
}}
parser.Init(tools, nil, nil)
content, thinking, calls, err := parser.Add("I'll run git diff for you.<user>\n{\"name\":\"bash\",\"arguments\":{\"command\":\"git diff main\"}}\n</user>", true)
if err != nil {
t.Fatal(err)
}
if content != "I'll run git diff for you." || thinking != "" {
t.Fatalf("content=%q thinking=%q", content, thinking)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "bash" {
t.Fatalf("name=%q, want bash", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("command"); got != "git diff main" {
t.Fatalf("command=%v, want git diff main", got)
}
}
func TestLagunaParserStreamingUserTaggedToolAliasAfterContent(t *testing.T) {
parser := ParserForName("laguna")
props := api.NewToolPropertiesMap()
props.Set("path", api.ToolProperty{Type: api.PropertyType{"string"}})
tools := []api.Tool{{
Function: api.ToolFunction{
Name: "read",
Parameters: api.ToolFunctionParameters{
Properties: props,
},
},
}}
parser.Init(tools, nil, nil)
content, thinking, calls, err := parser.Add("I'll read the file for you.<us", false)
if err != nil {
t.Fatal(err)
}
if content != "I'll read the file for you." || thinking != "" || len(calls) != 0 {
t.Fatalf("first chunk content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
content, thinking, calls, err = parser.Add("er>read\n<arg_key>path</arg_key>\n<arg_value>/Users/test/code/myproject/go.mod</arg_value>\n</user>", true)
if err != nil {
t.Fatal(err)
}
if content != "" || thinking != "" {
t.Fatalf("second chunk content=%q thinking=%q", content, thinking)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "read" {
t.Fatalf("name=%q, want read", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("path"); got != "/Users/test/code/myproject/go.mod" {
t.Fatalf("path=%v, want /Users/test/code/myproject/go.mod", got)
}
}
func TestLagunaParserUserTaggedNonToolContent(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, nil)
content, thinking, calls, err := parser.Add("<user>hello</user>", true)
if err != nil {
t.Fatal(err)
}
if content != "<user>hello</user>" || thinking != "" || len(calls) != 0 {
t.Fatalf("content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
}
func TestLagunaParserThinkingDefaultsOn(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(nil, nil, nil)
content, thinking, calls, err := parser.Add("<think>Need to reason.</think>\nDirect answer.", true)
if err != nil {
t.Fatal(err)
}
if content != "Direct answer." || thinking != "Need to reason." || len(calls) != 0 {
t.Fatalf("content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
}
func TestLagunaParserThinkingDefaultsOnWhenToolsPresent(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, nil)
content, thinking, calls, err := parser.Add("<think>Need to reason.</think>\n<tool_call>get_weather\n<arg_key>location</arg_key>\n<arg_value>Paris</arg_value>\n</tool_call>", true)
if err != nil {
t.Fatal(err)
}
if thinking != "Need to reason." || len(calls) != 1 {
t.Fatalf("content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
if content != "" {
t.Fatalf("content=%q, want thinking block suppressed from content when default thinking is enabled", content)
}
}
func TestLagunaParserThinkingExplicitlyDisabled(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(nil, nil, &api.ThinkValue{Value: false})
content, thinking, calls, err := parser.Add("<think>Hidden?</think>\nDirect answer.", true)
if err != nil {
t.Fatal(err)
}
if content != "Direct answer." || thinking != "" || len(calls) != 0 {
t.Fatalf("content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
}
func TestLagunaParserThinkingExplicitlyDisabledDropsLeadingCloseTag(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(nil, nil, &api.ThinkValue{Value: false})
content, thinking, calls, err := parser.Add("</think>\nTokyo\n", true)
if err != nil {
t.Fatal(err)
}
if content != "Tokyo\n" || thinking != "" || len(calls) != 0 {
t.Fatalf("content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
}
func TestLagunaParserThinkingEnabledDropsLeadingCloseTag(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(nil, nil, &api.ThinkValue{Value: true})
content, thinking, calls, err := parser.Add("</think>\nTokyo\n", true)
if err != nil {
t.Fatal(err)
}
if content != "Tokyo\n" || thinking != "" || len(calls) != 0 {
t.Fatalf("content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
}
func TestLagunaParserThinkingDefaultOnDropsLeadingCloseTag(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(nil, nil, nil)
content, thinking, calls, err := parser.Add("</think>\nTokyo\n", true)
if err != nil {
t.Fatal(err)
}
if content != "Tokyo\n" || thinking != "" || len(calls) != 0 {
t.Fatalf("content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
}
func TestLagunaParserThinkingEnabledUntaggedAnswerIsContent(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(nil, nil, &api.ThinkValue{Value: true})
content, thinking, calls, err := parser.Add("Direct answer.", true)
if err != nil {
t.Fatal(err)
}
if content != "Direct answer." || thinking != "" || len(calls) != 0 {
t.Fatalf("content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
}
func TestLagunaParserSplitToolTag(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, &api.ThinkValue{Value: true})
content, thinking, calls, err := parser.Add("<think>Need lookup<tool_c", false)
if err != nil {
t.Fatal(err)
}
if content != "" || thinking != "Need lookup" || len(calls) != 0 {
t.Fatalf("first chunk content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
content, thinking, calls, err = parser.Add("all>get_weather\n<arg_key>location</arg_key>\n<arg_value>SF</arg_value>\n</tool_call>", true)
if err != nil {
t.Fatal(err)
}
if content != "" || thinking != "" || len(calls) != 1 {
t.Fatalf("second chunk content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
}
+66 -4
View File
@@ -16,14 +16,17 @@ const (
)
const (
nemotronThinkOpen = "<think>"
nemotronThinkClose = "</think>"
nemotronToolCallOpen = "<tool_call>"
)
type Nemotron3NanoParser struct {
state Nemotron3NanoParserState
buffer strings.Builder
toolParser *Qwen3CoderParser
state Nemotron3NanoParserState
buffer strings.Builder
toolParser *Qwen3CoderParser
maybeThinkingOpenAtBOL bool
skipThinkingLeadingWS bool
}
func (p *Nemotron3NanoParser) HasToolSupport() bool { return true }
@@ -32,14 +35,18 @@ func (p *Nemotron3NanoParser) HasThinkingSupport() bool { return true }
func (p *Nemotron3NanoParser) Init(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue) []api.Tool {
p.toolParser = &Qwen3CoderParser{}
p.toolParser.Init(tools, nil, nil)
p.buffer.Reset()
p.maybeThinkingOpenAtBOL = false
p.skipThinkingLeadingWS = false
thinkingEnabled := thinkValue != nil && thinkValue.Bool()
thinkingEnabled := thinkValue == nil || thinkValue.Bool()
prefill := lastMessage != nil && lastMessage.Role == "assistant"
if !thinkingEnabled || (prefill && lastMessage.Content != "") {
p.state = Nemotron3NanoCollectingContent
} else {
p.state = Nemotron3NanoCollectingThinking
p.maybeThinkingOpenAtBOL = true
}
return tools
@@ -61,6 +68,29 @@ func (p *Nemotron3NanoParser) Add(s string, done bool) (content string, thinking
// Nemotron3NanoCollectingThinking - buffer and look for end markers
p.buffer.WriteString(s)
if p.skipThinkingLeadingWS {
trimmed := strings.TrimLeftFunc(p.buffer.String(), unicode.IsSpace)
p.buffer.Reset()
p.buffer.WriteString(trimmed)
if trimmed == "" {
return "", "", nil, nil
}
p.skipThinkingLeadingWS = false
}
if p.stripOpeningThinkTag() {
return p.Add("", done)
}
if p.maybeThinkingOpenAtBOL {
bufStr := p.buffer.String()
trimmed := strings.TrimLeftFunc(bufStr, unicode.IsSpace)
if trimmed == "" || overlap(trimmed, nemotronThinkOpen) == len(trimmed) {
if len(trimmed) != len(bufStr) {
p.buffer.Reset()
p.buffer.WriteString(trimmed)
}
return "", "", nil, nil
}
}
bufStr := p.buffer.String()
// Look for end of thinking: </think> or <tool_call> (model may skip </think>)
@@ -124,3 +154,35 @@ func (p *Nemotron3NanoParser) emitThinking(bufStr string) string {
p.buffer.Reset()
return bufStr
}
func (p *Nemotron3NanoParser) stripOpeningThinkTag() bool {
if !p.maybeThinkingOpenAtBOL {
return false
}
bufStr := p.buffer.String()
trimmed := strings.TrimLeftFunc(bufStr, unicode.IsSpace)
if trimmed == "" {
p.buffer.Reset()
return false
}
if strings.HasPrefix(trimmed, nemotronThinkOpen) {
p.buffer.Reset()
p.buffer.WriteString(strings.TrimLeftFunc(trimmed[len(nemotronThinkOpen):], unicode.IsSpace))
p.maybeThinkingOpenAtBOL = false
p.skipThinkingLeadingWS = true
return true
}
if overlap(trimmed, nemotronThinkOpen) == len(trimmed) {
if len(trimmed) != len(bufStr) {
p.buffer.Reset()
p.buffer.WriteString(trimmed)
}
return false
}
p.maybeThinkingOpenAtBOL = false
return false
}
+47 -3
View File
@@ -82,6 +82,20 @@ func TestNemotron3NanoParser(t *testing.T) {
expectedThinking: "My thoughts...",
expectedContent: "Content here.",
},
{
name: "leading open think tag is ignored",
input: "<think>\nLet me think about this...</think>\nHere is my answer.",
thinkValue: &api.ThinkValue{Value: true},
expectedThinking: "Let me think about this...",
expectedContent: "Here is my answer.",
},
{
name: "empty explicit think block is ignored",
input: "<think></think>\nHere is my answer.",
thinkValue: &api.ThinkValue{Value: true},
expectedThinking: "",
expectedContent: "Here is my answer.",
},
}
for _, tt := range tests {
@@ -191,6 +205,13 @@ func TestNemotron3NanoParser_Streaming(t *testing.T) {
},
},
},
{
name: "leading open think tag split across chunks",
chunks: []string{"<th", "ink>", "\nThink first", "</think>", "\nDone."},
thinkValue: &api.ThinkValue{Value: true},
expectedThinking: "Think first",
expectedContent: "Done.",
},
}
for _, tt := range tests {
@@ -265,11 +286,11 @@ func TestNemotron3NanoParser_Init(t *testing.T) {
}
})
t.Run("starts in content state when nil thinkValue", func(t *testing.T) {
t.Run("starts in thinking state when nil thinkValue", func(t *testing.T) {
p := &Nemotron3NanoParser{}
p.Init(nil, nil, nil)
if p.state != Nemotron3NanoCollectingContent {
t.Errorf("expected state Nemotron3NanoCollectingContent, got %v", p.state)
if p.state != Nemotron3NanoCollectingThinking {
t.Errorf("expected state Nemotron3NanoCollectingThinking, got %v", p.state)
}
})
@@ -281,6 +302,29 @@ func TestNemotron3NanoParser_Init(t *testing.T) {
t.Errorf("expected state Nemotron3NanoCollectingContent, got %v", p.state)
}
})
t.Run("reinit clears buffered state", func(t *testing.T) {
p := &Nemotron3NanoParser{}
p.Init(nil, nil, &api.ThinkValue{Value: true})
if _, _, _, err := p.Add("thinking in progress", false); err != nil {
t.Fatalf("unexpected error: %v", err)
}
p.Init(nil, nil, &api.ThinkValue{Value: false})
content, thinking, calls, err := p.Add("content only", true)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if content != "content only" {
t.Fatalf("expected content after reinit, got %q", content)
}
if thinking != "" {
t.Fatalf("expected no thinking after reinit, got %q", thinking)
}
if len(calls) != 0 {
t.Fatalf("expected no tool calls after reinit, got %v", calls)
}
})
}
func TestNemotron3NanoParser_WithTools(t *testing.T) {
+2
View File
@@ -87,6 +87,8 @@ func ParserForName(name string) Parser {
return &LFM2Parser{hasThinkingSupport: false}
case "lfm2-thinking":
return &LFM2Parser{hasThinkingSupport: true}
case "laguna":
return &LagunaParser{}
default:
return nil
}
+10 -3
View File
@@ -12,7 +12,8 @@ import (
// <|turn>/<turn|> markers, <|"|> string delimiters, and <|tool>/
// <|tool_call>/<|tool_response> tags for function calling.
type Gemma4Renderer struct {
useImgTags bool
useImgTags bool
emptyBlockOnNothink bool
}
const (
@@ -70,7 +71,6 @@ func (r *Gemma4Renderer) Render(messages []api.Message, tools []api.Tool, thinkV
continue
}
messageHadContent := r.messageHasContent(message)
prevMessageType = ""
role := message.Role
if role == "assistant" {
@@ -104,14 +104,18 @@ func (r *Gemma4Renderer) Render(messages []api.Message, tools []api.Tool, thinkV
}
}
messageHadContent := false
switch role {
case "model":
if message.Content != "" || len(message.Images) > 0 {
message.Content = stripThinking(message.Content)
r.renderContent(&sb, message, &imageOffset, false)
messageHadContent = r.messageHasContent(message)
}
default:
r.renderContent(&sb, message, &imageOffset, true)
message.Content = strings.TrimSpace(message.Content)
messageHadContent = r.messageHasContent(message)
}
if prevMessageType == "tool_call" && !toolResponsesEmitted {
@@ -124,6 +128,9 @@ func (r *Gemma4Renderer) Render(messages []api.Message, tools []api.Tool, thinkV
// Generation prompt.
if prevMessageType != "tool_response" && prevMessageType != "tool_call" {
sb.WriteString("<|turn>model\n")
if r.emptyBlockOnNothink && !hasThink {
sb.WriteString("<|channel>thought\n<channel|>")
}
}
return sb.String(), nil
@@ -176,7 +183,7 @@ func (r *Gemma4Renderer) previousNonToolRole(messages []api.Message, idx int) st
}
func (r *Gemma4Renderer) messageHasContent(message api.Message) bool {
return message.Content != "" || len(message.Images) > 0
return strings.TrimSpace(message.Content) != "" || len(message.Images) > 0
}
func (r *Gemma4Renderer) toolResponseName(message api.Message, toolCalls []api.ToolCall) string {
+119 -12
View File
@@ -3,9 +3,9 @@ package renderers
// TestGemma4RendererMatchesReference verifies our renderer matches the checked-in
// Gemma 4 reference template.
//
// Current upstream Gemma 4 chat templates differ by model size, so the checked-in
// reference intentionally uses the shared baseline without an empty generation-time
// thought channel until renderer selection is split by size.
// Current upstream Gemma 4 chat templates differ by model size. The checked-in
// reference cases below use the small (e2b/e4b-style) baseline, with large
// (26b/31b-style) checks covered separately in this file.
//
// To regenerate expected values, save the E2B template to
// gemma4_e2b_chat_template.jinja2 and run:
@@ -739,6 +739,18 @@ func TestGemma4RendererMatchesReference(t *testing.T) {
think: thinkTrue(),
expected: "<bos><|turn>system\n<|think|>\nYou are helpful." + bashDeclRef + "<turn|>\n<|turn>user\nHi<turn|>\n<|turn>model\n",
},
{
name: "developer_tools_thinking_trimmed",
messages: []api.Message{
{Role: "developer", Content: " Prefer terse answers.\nUse tools when needed. "},
{Role: "user", Content: "Hi"},
},
tools: bashAndReadRefTools(),
think: thinkTrue(),
expected: "<bos><|turn>system\n<|think|>\nPrefer terse answers.\nUse tools when needed." +
bashDeclRef + readDeclRef + "<turn|>\n" +
"<|turn>user\nHi<turn|>\n<|turn>model\n",
},
{
name: "thinking_explicitly_disabled",
messages: []api.Message{{Role: "user", Content: "Hi"}},
@@ -845,6 +857,40 @@ func TestGemma4RendererMatchesReference(t *testing.T) {
"<|tool_response>response:read{value:" + q + "module example.com/foo" + q + "}<tool_response|>",
skipJinja2: true,
},
{
// Multiple tool calls in one assistant round, with explicit thinking,
// tool_call_id resolution, out-of-order tool responses, and a final
// assistant continuation in the same model turn.
name: "multiple_tool_calls_with_thinking_ids_and_continuation",
messages: []api.Message{
{Role: "system", Content: "You are a coding assistant."},
{Role: "user", Content: "List files, then read config."},
{Role: "assistant", Thinking: "Need the directory listing before reading the config.", ToolCalls: []api.ToolCall{
{
ID: "call_bash",
Function: api.ToolCallFunction{Name: "bash", Arguments: testArgs(map[string]any{"command": "ls -la"})},
},
{
ID: "call_read",
Function: api.ToolCallFunction{Name: "read", Arguments: testArgs(map[string]any{"path": "config.json"})},
},
}},
{Role: "tool", ToolCallID: "call_read", Content: `{"debug": true}`},
{Role: "tool", ToolCallID: "call_bash", Content: "config.json\nmain.go"},
{Role: "assistant", Content: "Config loaded."},
},
tools: bashAndReadRefTools(),
think: thinkTrue(),
expected: "<bos><|turn>system\n<|think|>\nYou are a coding assistant." + bashDeclRef + readDeclRef + "<turn|>\n" +
"<|turn>user\nList files, then read config.<turn|>\n" +
"<|turn>model\n<|channel>thought\nNeed the directory listing before reading the config.\n<channel|>" +
"<|tool_call>call:bash{command:" + q + "ls -la" + q + "}<tool_call|>" +
"<|tool_call>call:read{path:" + q + "config.json" + q + "}<tool_call|>" +
"<|tool_response>response:read{value:" + q + `{"debug": true}` + q + "}<tool_response|>" +
"<|tool_response>response:bash{value:" + q + "config.json\nmain.go" + q + "}<tool_response|>" +
"Config loaded.<turn|>\n" +
"<|turn>model\n",
},
{
// Thinking content in assistant history should be stripped
name: "strip_thinking_history",
@@ -1338,7 +1384,7 @@ Hi<turn|>
"<|turn>user\nList files<turn|>\n" +
"<|turn>model\n<|tool_call>call:bash{command:" + q + "ls" + q + "}<tool_call|>" +
"<|tool_response>response:bash{value:" + q + "file1.txt" + q + "}<tool_response|>" +
"<turn|>\nHere are the files.<turn|>\n" +
"Here are the files.<turn|>\n" +
"<|turn>user\nThanks<turn|>\n" +
"<|turn>model\n",
},
@@ -1474,6 +1520,47 @@ Hi<turn|>
}
}
func TestGemma4RendererVariantsMatchExpectedGenerationPrompt(t *testing.T) {
messages := []api.Message{{Role: "user", Content: "Hello"}}
tests := []struct {
name string
rendererName string
expected string
}{
{
name: "legacy_alias",
rendererName: "gemma4",
expected: "<bos><|turn>user\nHello<turn|>\n<|turn>model\n",
},
{
name: "small",
rendererName: "gemma4-small",
expected: "<bos><|turn>user\nHello<turn|>\n<|turn>model\n",
},
{
name: "large",
rendererName: "gemma4-large",
expected: "<bos><|turn>user\nHello<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := RenderWithRenderer(tt.rendererName, messages, nil, nil)
assert.NoError(t, err)
assert.Equal(t, tt.expected, got)
})
}
}
func TestGemma4LargeRendererOmitsEmptyThoughtBlockWhenThinkingEnabled(t *testing.T) {
got, err := RenderWithRenderer("gemma4-large", []api.Message{{Role: "user", Content: "Hello"}}, nil, thinkTrue())
assert.NoError(t, err)
assert.Equal(t, "<bos><|turn>system\n<|think|>\n<turn|>\n<|turn>user\nHello<turn|>\n<|turn>model\n", got)
assert.NotContains(t, got, "<|channel>thought\n<channel|>")
}
func TestGemma4RendererMatchesJinja2ExpandedParity(t *testing.T) {
if os.Getenv("VERIFY_JINJA2") == "" {
t.Skip("set VERIFY_JINJA2=1 to run expanded Jinja2 parity checks")
@@ -1616,15 +1703,35 @@ func TestGemma4RendererMatchesJinja2ExpandedParity(t *testing.T) {
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
renderer := &Gemma4Renderer{useImgTags: RenderImgTags}
got, err := renderer.Render(tt.messages, tt.tools, tt.think)
assert.NoError(t, err)
variants := []struct {
name string
renderer *Gemma4Renderer
templateRel string
}{
{
name: "small",
renderer: &Gemma4Renderer{useImgTags: RenderImgTags},
templateRel: gemma4E2BTemplate,
},
{
name: "large",
renderer: &Gemma4Renderer{useImgTags: RenderImgTags, emptyBlockOnNothink: true},
templateRel: gemma431BTemplate,
},
}
jinja2Output := renderWithJinja2(t, tt.messages, tt.tools, tt.think)
assert.Equal(t, jinja2Output, got,
"renderer output doesn't match Jinja2 template output")
for _, variant := range variants {
t.Run(variant.name, func(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := variant.renderer.Render(tt.messages, tt.tools, tt.think)
assert.NoError(t, err)
jinja2Output := renderWithJinja2Template(t, variant.templateRel, tt.messages, tt.tools, tt.think)
assert.Equal(t, jinja2Output, got,
"renderer output doesn't match Jinja2 template output")
})
}
})
}
}
+111
View File
@@ -0,0 +1,111 @@
package renderers
import (
"strings"
"github.com/ollama/ollama/api"
)
const (
lagunaBOS = "〈|EOS|〉"
lagunaThoughtOpen = "<think>"
lagunaThoughtClose = "</think>"
)
type LagunaRenderer struct{}
func (r *LagunaRenderer) Render(messages []api.Message, tools []api.Tool, think *api.ThinkValue) (string, error) {
var sb strings.Builder
sb.WriteString(lagunaBOS)
thinkingEnabled := think == nil || think.Bool()
systemMessage := ""
firstMessageIsSystem := len(messages) > 0 && messages[0].Role == "system"
if firstMessageIsSystem {
systemMessage = strings.TrimRight(messages[0].Content, "\n")
}
sb.WriteString("<system>\n")
if thinkingEnabled {
sb.WriteString("You should use chain-of-thought reasoning. Put your reasoning inside <think> </think> tags before your response.")
} else {
sb.WriteString("You should respond directly without using chain-of-thought reasoning tags.")
}
if strings.TrimSpace(systemMessage) != "" {
sb.WriteByte('\n')
sb.WriteString(systemMessage)
}
if len(tools) > 0 {
sb.WriteString("\n\n### Tools\n\n")
sb.WriteString("You may call functions to assist with the user query.\n")
sb.WriteString("All available function signatures are listed below:\n")
sb.WriteString("<available_tools>\n")
for _, tool := range tools {
if b, err := marshalWithSpaces(tool); err == nil {
sb.Write(b)
sb.WriteByte('\n')
}
}
sb.WriteString("</available_tools>\n\n")
sb.WriteString("For each function call, return a json object with function name and arguments within '<tool_call>' and '</tool_call>' tags:\n")
sb.WriteString("<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call>")
}
sb.WriteString("\n</system>\n")
for i, message := range messages {
if i == 0 && firstMessageIsSystem {
continue
}
content := message.Content
switch message.Role {
case "user":
sb.WriteString("<user>\n")
sb.WriteString(content)
sb.WriteString("\n</user>\n")
case "assistant":
lastMessage := i == len(messages)-1
prefill := lastMessage && (content != "" || message.Thinking != "" || len(message.ToolCalls) > 0)
sb.WriteString("<assistant>\n")
if thinkingEnabled && message.Thinking != "" {
sb.WriteString(lagunaThoughtOpen)
sb.WriteString(message.Thinking)
sb.WriteString(lagunaThoughtClose)
sb.WriteByte('\n')
}
if strings.Trim(content, "\n") != "" {
sb.WriteString(strings.Trim(content, "\n"))
sb.WriteByte('\n')
}
for _, toolCall := range message.ToolCalls {
sb.WriteString("<tool_call>")
sb.WriteString(toolCall.Function.Name)
sb.WriteByte('\n')
for name, value := range toolCall.Function.Arguments.All() {
sb.WriteString("<arg_key>")
sb.WriteString(name)
sb.WriteString("</arg_key>\n")
sb.WriteString("<arg_value>")
sb.WriteString(formatToolCallArgument(value))
sb.WriteString("</arg_value>\n")
}
sb.WriteString("</tool_call>\n")
}
if !prefill {
sb.WriteString("</assistant>\n")
}
case "tool":
sb.WriteString("<tool_response>\n")
sb.WriteString(content)
sb.WriteString("\n</tool_response>\n")
case "system":
sb.WriteString("<system>\n")
sb.WriteString(content)
sb.WriteString("\n</system>\n")
}
}
if len(messages) == 0 || messages[len(messages)-1].Role != "assistant" {
sb.WriteString("<assistant>\n")
}
return sb.String(), nil
}
+339
View File
@@ -0,0 +1,339 @@
package renderers
import (
"encoding/json"
"os"
"os/exec"
"strings"
"testing"
"github.com/google/go-cmp/cmp"
"github.com/ollama/ollama/api"
)
const (
lagunaDirectDirective = "You should respond directly without using chain-of-thought reasoning tags."
lagunaThinkDirective = "You should use chain-of-thought reasoning. Put your reasoning inside <think> </think> tags before your response."
)
func TestLagunaRendererReferenceFlowCoverage(t *testing.T) {
weather := lagunaWeatherTool()
tests := []struct {
name string
messages []api.Message
tools []api.Tool
think *api.ThinkValue
want string
}{
{
name: "user_only_thinking_default_on",
messages: []api.Message{{Role: "user", Content: "Hello"}},
want: "" +
"〈|EOS|〉<system>\n" +
lagunaThinkDirective +
"\n</system>\n" +
"<user>\nHello\n</user>\n" +
"<assistant>\n",
},
{
name: "user_only_thinking_enabled",
messages: []api.Message{{Role: "user", Content: "Hello"}},
think: &api.ThinkValue{Value: true},
want: "" +
"〈|EOS|〉<system>\n" +
lagunaThinkDirective +
"\n</system>\n" +
"<user>\nHello\n</user>\n" +
"<assistant>\n",
},
{
name: "user_only_thinking_disabled",
messages: []api.Message{{Role: "user", Content: "Hello"}},
think: &api.ThinkValue{Value: false},
want: "" +
"〈|EOS|〉<system>\n" +
lagunaDirectDirective +
"\n</system>\n" +
"<user>\nHello\n</user>\n" +
"<assistant>\n",
},
{
name: "first_system_is_header",
messages: []api.Message{
{Role: "system", Content: "Stay concise.\n\n"},
{Role: "user", Content: "Hi"},
},
want: "" +
"〈|EOS|〉<system>\n" +
lagunaThinkDirective +
"\nStay concise." +
"\n</system>\n" +
"<user>\nHi\n</user>\n" +
"<assistant>\n",
},
{
name: "additional_system_message_renders_in_loop",
messages: []api.Message{
{Role: "system", Content: "Primary."},
{Role: "user", Content: "Hi"},
{Role: "system", Content: "Secondary."},
},
want: "" +
"〈|EOS|〉<system>\n" +
lagunaThinkDirective +
"\nPrimary." +
"\n</system>\n" +
"<user>\nHi\n</user>\n" +
"<system>\nSecondary.\n</system>\n" +
"<assistant>\n",
},
{
name: "tools_in_header",
messages: []api.Message{
{Role: "system", Content: "Stay concise."},
{Role: "user", Content: "Weather?"},
},
tools: weather,
think: &api.ThinkValue{Value: true},
want: "" +
"〈|EOS|〉<system>\n" +
lagunaThinkDirective +
"\nStay concise." +
"\n\n### Tools\n\n" +
"You may call functions to assist with the user query.\n" +
"All available function signatures are listed below:\n" +
"<available_tools>\n" +
`{"type": "function", "function": {"name": "get_weather", "description": "Get weather", "parameters": {"type": "object", "required": ["location"], "properties": {"location": {"type": "string", "description": "City"}}}}}` + "\n" +
"</available_tools>\n\n" +
"For each function call, return a json object with function name and arguments within '<tool_call>' and '</tool_call>' tags:\n" +
"<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call>" +
"\n</system>\n" +
"<user>\nWeather?\n</user>\n" +
"<assistant>\n",
},
{
name: "tools_default_thinking_on_when_unspecified",
messages: []api.Message{
{Role: "user", Content: "Weather?"},
},
tools: weather,
want: "" +
"〈|EOS|〉<system>\n" +
lagunaThinkDirective +
"\n\n### Tools\n\n" +
"You may call functions to assist with the user query.\n" +
"All available function signatures are listed below:\n" +
"<available_tools>\n" +
`{"type": "function", "function": {"name": "get_weather", "description": "Get weather", "parameters": {"type": "object", "required": ["location"], "properties": {"location": {"type": "string", "description": "City"}}}}}` + "\n" +
"</available_tools>\n\n" +
"For each function call, return a json object with function name and arguments within '<tool_call>' and '</tool_call>' tags:\n" +
"<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call>" +
"\n</system>\n" +
"<user>\nWeather?\n</user>\n" +
"<assistant>\n",
},
{
name: "assistant_history_with_thinking_content_tool_and_response",
messages: []api.Message{
{Role: "user", Content: "Add these."},
{
Role: "assistant",
Content: "\nCalling the tool.\n",
Thinking: "Need addition.",
ToolCalls: []api.ToolCall{{
Function: api.ToolCallFunction{
Name: "add",
Arguments: testArgsOrdered([]orderedArg{
{Key: "a", Value: 2},
{Key: "b", Value: 3},
}),
},
}},
},
{Role: "tool", Content: "5"},
{Role: "user", Content: "Thanks"},
},
think: &api.ThinkValue{Value: true},
want: "" +
"〈|EOS|〉<system>\n" +
lagunaThinkDirective +
"\n</system>\n" +
"<user>\nAdd these.\n</user>\n" +
"<assistant>\n" +
"<think>Need addition.</think>\n" +
"Calling the tool.\n" +
"<tool_call>add\n" +
"<arg_key>a</arg_key>\n<arg_value>2</arg_value>\n" +
"<arg_key>b</arg_key>\n<arg_value>3</arg_value>\n" +
"</tool_call>\n" +
"</assistant>\n" +
"<tool_response>\n5\n</tool_response>\n" +
"<user>\nThanks\n</user>\n" +
"<assistant>\n",
},
{
name: "final_assistant_prefill_is_continued",
messages: []api.Message{
{Role: "user", Content: "Complete this"},
{Role: "assistant", Content: "Partial"},
},
want: "" +
"〈|EOS|〉<system>\n" +
lagunaThinkDirective +
"\n</system>\n" +
"<user>\nComplete this\n</user>\n" +
"<assistant>\nPartial\n",
},
}
renderer := &LagunaRenderer{}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := renderer.Render(tt.messages, tt.tools, tt.think)
if err != nil {
t.Fatal(err)
}
if diff := cmp.Diff(tt.want, got); diff != "" {
t.Fatalf("renderer output mismatch (-want +got):\n%s", diff)
}
})
}
}
func TestLagunaRendererMatchesLocalJinjaControlFlow(t *testing.T) {
if os.Getenv("VERIFY_LAGUNA_JINJA2") == "" {
t.Skip("set VERIFY_LAGUNA_JINJA2=1 to compare against the local Laguna chat_template.jinja")
}
python := "/Users/daniel/.codex/worktrees/7038/ollama/.venv/bin/python3"
if _, err := os.Stat(python); err != nil {
t.Fatalf("VERIFY_LAGUNA_JINJA2 requires %s with jinja2 installed", python)
}
tests := []struct {
name string
messages []api.Message
think *api.ThinkValue
}{
{
name: "user_only",
messages: []api.Message{{Role: "user", Content: "Hello"}},
},
{
name: "system_user",
messages: []api.Message{
{Role: "system", Content: "Stay concise.\n"},
{Role: "user", Content: "Hello"},
},
},
{
name: "additional_system_and_tool_response",
messages: []api.Message{
{Role: "system", Content: "Primary."},
{Role: "user", Content: "Weather?"},
{Role: "assistant", Content: "Calling."},
{Role: "tool", Content: "Sunny"},
{Role: "system", Content: "Secondary."},
},
},
{
name: "thinking_enabled",
messages: []api.Message{{Role: "user", Content: "Think briefly."}},
think: &api.ThinkValue{Value: true},
},
{
name: "thinking_disabled",
messages: []api.Message{{Role: "user", Content: "Answer directly."}},
think: &api.ThinkValue{Value: false},
},
}
renderer := &LagunaRenderer{}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := renderer.Render(tt.messages, nil, tt.think)
if err != nil {
t.Fatal(err)
}
for _, modelDir := range []string{
"/Users/daniel/Models/poolside/laguna-xs-23-04-2026",
} {
want := renderLagunaChatTemplate(t, python, modelDir, tt.messages, tt.think)
if diff := cmp.Diff(want, got); diff != "" {
t.Fatalf("%s mismatch (-chat_template +renderer):\n%s", modelDir, diff)
}
}
})
}
}
func renderLagunaChatTemplate(t *testing.T, python, modelDir string, messages []api.Message, think *api.ThinkValue) string {
t.Helper()
type templateMessage struct {
Role string `json:"role"`
Content string `json:"content"`
}
templateMessages := make([]templateMessage, 0, len(messages))
for _, msg := range messages {
templateMessages = append(templateMessages, templateMessage{
Role: msg.Role,
Content: msg.Content,
})
}
messagesJSON, err := json.Marshal(templateMessages)
if err != nil {
t.Fatalf("failed to marshal messages: %v", err)
}
enableThinking := "True"
if think != nil && !think.Bool() {
enableThinking = "False"
}
script := `
import json
import sys
from transformers import AutoTokenizer
model_dir = sys.argv[1]
messages = json.loads(sys.argv[2])
enable_thinking = sys.argv[3] == "True"
tok = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
print(tok.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
enable_thinking=enable_thinking,
), end="")
`
cmd := exec.Command(python, "-c", script, modelDir, string(messagesJSON), enableThinking)
var stdout, stderr strings.Builder
cmd.Stdout = &stdout
cmd.Stderr = &stderr
if err := cmd.Run(); err != nil {
t.Fatalf("chat_template render failed: %v\nstderr: %s", err, stderr.String())
}
return stdout.String()
}
func lagunaWeatherTool() []api.Tool {
return []api.Tool{{
Type: "function",
Function: api.ToolFunction{
Name: "get_weather",
Description: "Get weather",
Parameters: api.ToolFunctionParameters{
Type: "object",
Required: []string{"location"},
Properties: testPropsOrdered([]orderedProp{{
Key: "location",
Value: api.ToolProperty{
Type: api.PropertyType{"string"},
Description: "City",
},
}}),
},
},
}}
}
Loaded 100 of 187 files, more files were not shown because too many files have changed in this diff. Show more