Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9ba5a04914 | ||
|
|
938ca6e274 | ||
|
|
8f39fff70b | ||
|
|
4fe5609563 | ||
|
|
917324bb4d | ||
|
|
c7c2837c96 | ||
|
|
b6447caebc | ||
|
|
bad32c7244 | ||
|
|
ab2e005bf7 | ||
|
|
321cc8a2ba | ||
|
|
87288ced4f | ||
|
|
2bbe2405fe | ||
|
|
bd21678b16 | ||
|
|
088dfd89a8 | ||
|
|
3cab8a7b02 | ||
|
|
03aee88186 | ||
|
|
ec9b4e9e47 | ||
|
|
4656a07e56 | ||
|
|
30f86cb9dd | ||
|
|
ea01af6f76 | ||
|
|
c2ebb4d57c | ||
|
|
590109c835 | ||
|
|
b4442c6d17 | ||
|
|
85ff8e4a21 | ||
|
|
160660e572 | ||
|
|
3b43b9bc4b | ||
|
|
21883571b7 | ||
|
|
ce99f24731 | ||
|
|
04f5f0cdb4 | ||
|
|
fb36a01ffe | ||
|
|
0c65ed33bc | ||
|
|
22d6c817f8 | ||
|
|
ca01373b28 | ||
|
|
24e038d56a | ||
|
|
5d1021603a | ||
|
|
8e05d734b9 | ||
|
|
05e0f21bec | ||
|
|
ff23dd343f | ||
|
|
123b300af6 | ||
|
|
57653b8e42 | ||
|
|
a50ce61c54 | ||
|
|
2bb7ea00d2 | ||
|
|
55fa80d07a | ||
|
|
b9cb535407 | ||
|
|
031baef094 | ||
|
|
7d271e6dc9 | ||
|
|
c88dae2d6b | ||
|
|
9e3618d663 | ||
|
|
e585ecd11f |
No files matched your search
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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" }`,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 |
@@ -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) {
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
});
|
||||
}
|
||||
@@ -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;
|
||||
});
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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");
|
||||
});
|
||||
});
|
||||
@@ -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];
|
||||
}
|
||||
@@ -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())
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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, "'", "''") + "'"
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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":
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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"])
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -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/")
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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, "", " ")
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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") {
|
||||
|
||||
@@ -45,7 +45,7 @@ type menuItem struct {
|
||||
isOthers bool
|
||||
}
|
||||
|
||||
const pinnedIntegrationCount = 3
|
||||
const pinnedIntegrationCount = 4
|
||||
|
||||
var runModelMenuItem = menuItem{
|
||||
title: "Chat with a model",
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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()...)
|
||||
|
||||
@@ -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" ||
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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{
|
||||
|
||||
@@ -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 ""
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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:
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
|
After Width: | Height: | Size: 307 KiB |
|
After Width: | Height: | Size: 297 KiB |
|
After Width: | Height: | Size: 1.4 MiB |
@@ -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:
|
||||
|
||||

|
||||
|
||||
## 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
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
|
||||

|
||||
|
||||
## 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
|
||||
```
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
```
|
||||
|
||||
@@ -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
|
||||
```
|
||||
@@ -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",
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
@@ -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")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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)),
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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) {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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")
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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",
|
||||
},
|
||||
}}),
|
||||
},
|
||||
},
|
||||
}}
|
||||
}
|
||||