iop/apps/node/internal/adapters/ollama/ollama_test.go

541 lines
18 KiB
Go

package ollama
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"go.uber.org/zap"
noderuntime "iop/apps/node/internal/runtime"
"iop/packages/go/config"
)
type fakeSink struct {
mu sync.Mutex
events []noderuntime.RuntimeEvent
}
func (s *fakeSink) Emit(_ context.Context, event noderuntime.RuntimeEvent) error {
s.mu.Lock()
defer s.mu.Unlock()
s.events = append(s.events, event)
return nil
}
func (s *fakeSink) all() []noderuntime.RuntimeEvent {
s.mu.Lock()
defer s.mu.Unlock()
return append([]noderuntime.RuntimeEvent(nil), s.events...)
}
func TestOllamaExecuteStreamsChatDeltas(t *testing.T) {
var gotModel string
var gotPrompt string
var gotNumCtx int
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/chat" {
t.Fatalf("unexpected path %s", r.URL.Path)
}
var req ollamaChatRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
t.Fatalf("decode request: %v", err)
}
gotModel = req.Model
if req.Options != nil {
switch v := req.Options["num_ctx"].(type) {
case float64:
gotNumCtx = int(v)
case int:
gotNumCtx = v
}
}
if len(req.Messages) != 1 {
t.Fatalf("expected one message, got %+v", req.Messages)
}
gotPrompt = req.Messages[0].Content
w.Header().Set("Content-Type", "application/x-ndjson")
_, _ = w.Write([]byte(`{"message":{"role":"assistant","content":"hello "},"done":false}` + "\n"))
_, _ = w.Write([]byte(`{"message":{"role":"assistant","content":"world"},"done":false}` + "\n"))
_, _ = w.Write([]byte(`{"done":true,"prompt_eval_count":3,"eval_count":2}` + "\n"))
}))
defer server.Close()
adapter := New(config.OllamaConf{BaseURL: server.URL, ContextSize: 262144}, zap.NewNop())
sink := &fakeSink{}
err := adapter.Execute(context.Background(), noderuntime.ExecutionSpec{
RunID: "run-1",
Target: "llama-test",
Input: map[string]any{"prompt": "say hello"},
}, sink)
if err != nil {
t.Fatalf("Execute failed: %v", err)
}
if gotModel != "llama-test" {
t.Fatalf("model: got %q", gotModel)
}
if gotPrompt != "say hello" {
t.Fatalf("prompt: got %q", gotPrompt)
}
if gotNumCtx != 262144 {
t.Fatalf("num_ctx: got %d", gotNumCtx)
}
events := sink.all()
if len(events) != 4 {
t.Fatalf("expected 4 events, got %+v", events)
}
if events[0].Type != noderuntime.EventTypeStart {
t.Fatalf("expected start event, got %+v", events[0])
}
if events[1].Delta+events[2].Delta != "hello world" {
t.Fatalf("unexpected delta text: %+v", events)
}
if events[3].Type != noderuntime.EventTypeComplete {
t.Fatalf("expected complete event, got %+v", events[3])
}
if events[3].Usage == nil || events[3].Usage.InputTokens != 3 || events[3].Usage.OutputTokens != 2 {
t.Fatalf("unexpected usage: %+v", events[3].Usage)
}
}
func TestOllamaExecutePassesOptionsAndTopLevelFields(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/chat" {
t.Fatalf("unexpected path %s", r.URL.Path)
}
var req ollamaChatRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
t.Fatalf("decode request: %v", err)
}
if req.Options["temperature"].(float64) != 0.2 || req.Options["top_p"].(float64) != 0.9 || req.Options["num_predict"].(float64) != 32 {
t.Fatalf("unexpected options: %+v", req.Options)
}
switch v := req.Options["num_ctx"].(type) {
case float64:
if int(v) != 262144 {
t.Fatalf("profile context_size should override request num_ctx: got %d, want 262144", int(v))
}
case int:
if v != 262144 {
t.Fatalf("profile context_size should override request num_ctx: got %d, want 262144", v)
}
default:
t.Fatalf("unexpected num_ctx type %T in %+v", v, req.Options)
}
if req.KeepAlive != "10m" || req.Think != false || req.Format != "json" {
t.Fatalf("top-level fields not passed: keep_alive=%v think=%v format=%v", req.KeepAlive, req.Think, req.Format)
}
if tools, ok := req.Tools.([]any); !ok || len(tools) != 1 {
t.Fatalf("tools not passed: %+v", req.Tools)
}
w.Header().Set("Content-Type", "application/x-ndjson")
_, _ = w.Write([]byte(`{"message":{"role":"assistant","content":"ok"},"done":false}` + "\n"))
_, _ = w.Write([]byte(`{"done":true}` + "\n"))
}))
defer server.Close()
adapter := New(config.OllamaConf{BaseURL: server.URL, ContextSize: 262144}, zap.NewNop())
sink := &fakeSink{}
err := adapter.Execute(context.Background(), noderuntime.ExecutionSpec{
RunID: "run-options",
Target: "llama-test",
Input: map[string]any{
"prompt": "hi",
"options": map[string]any{
"temperature": 0.2,
"top_p": 0.9,
"num_predict": float64(32),
"num_ctx": float64(4096),
},
"keep_alive": "10m",
"think": false,
"format": "json",
"tools": []any{map[string]any{"type": "function"}},
},
}, sink)
if err != nil {
t.Fatalf("Execute failed: %v", err)
}
}
func TestOllamaExecutePreservesNativeToolCalls(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/chat" {
t.Fatalf("unexpected path %s", r.URL.Path)
}
w.Header().Set("Content-Type", "application/x-ndjson")
_, _ = w.Write([]byte(`{"message":{"role":"assistant","tool_calls":[{"id":"call_1","type":"function","function":{"name":"run_commands","arguments":"{\"commands\":[\"git status\"]}"}}]},"done":false}` + "\n"))
_, _ = w.Write([]byte(`{"done":true}` + "\n"))
}))
defer server.Close()
adapter := New(config.OllamaConf{BaseURL: server.URL}, zap.NewNop())
sink := &fakeSink{}
err := adapter.Execute(context.Background(), noderuntime.ExecutionSpec{
RunID: "run-tools",
Target: "llama-test",
Input: map[string]any{"prompt": "status"},
}, sink)
if err != nil {
t.Fatalf("Execute failed: %v", err)
}
events := sink.all()
complete := events[len(events)-1]
if complete.Metadata["finish_reason"] != "tool_calls" {
t.Fatalf("finish_reason: %+v", complete.Metadata)
}
var toolCalls []map[string]any
if err := json.Unmarshal([]byte(complete.Metadata[runtimeMetadataOpenAIToolCalls]), &toolCalls); err != nil {
t.Fatalf("tool_calls metadata JSON: %v", err)
}
if len(toolCalls) != 1 || toolCalls[0]["id"] != "call_1" {
t.Fatalf("tool_calls: %+v", toolCalls)
}
}
func TestOllamaExecuteEmitsErrorForHTTPFailure(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
http.Error(w, "boom", http.StatusBadGateway)
}))
defer server.Close()
adapter := New(config.OllamaConf{BaseURL: server.URL}, zap.NewNop())
sink := &fakeSink{}
err := adapter.Execute(context.Background(), noderuntime.ExecutionSpec{
RunID: "run-err",
Target: "llama-test",
Input: map[string]any{"prompt": "hi"},
}, sink)
if err == nil {
t.Fatal("expected error")
}
events := sink.all()
if len(events) < 2 || events[1].Type != noderuntime.EventTypeError {
t.Fatalf("expected error event, got %+v", events)
}
if !strings.Contains(events[1].Error, "boom") {
t.Fatalf("expected HTTP body in error, got %q", events[1].Error)
}
}
func TestOllamaCapabilitiesQueryTags(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/tags" {
t.Fatalf("unexpected path %s", r.URL.Path)
}
_, _ = w.Write([]byte(`{"models":[{"name":"llama-a"},{"name":"llama-b"}]}`))
}))
defer server.Close()
adapter := New(config.OllamaConf{
BaseURL: server.URL,
Capacity: 5,
MaxQueue: 10,
QueueTimeoutMS: 15000,
RequestTimeoutMS: 30000,
}, zap.NewNop())
caps, err := adapter.Capabilities(context.Background())
if err != nil {
t.Fatalf("Capabilities failed: %v", err)
}
got := strings.Join(caps.Targets, ",")
if got != "llama-a,llama-b" {
t.Fatalf("targets: got %q", got)
}
if caps.MaxConcurrency != 5 {
t.Fatalf("expected MaxConcurrency 5, got %d", caps.MaxConcurrency)
}
if caps.MaxQueue != 10 {
t.Fatalf("expected MaxQueue 10, got %d", caps.MaxQueue)
}
if caps.QueueTimeoutMS != 15000 {
t.Fatalf("expected QueueTimeoutMS 15000, got %d", caps.QueueTimeoutMS)
}
if caps.RequestTimeoutMS != 30000 {
t.Fatalf("expected RequestTimeoutMS 30000, got %d", caps.RequestTimeoutMS)
}
// Verify default capacity fallback
adapterDefault := New(config.OllamaConf{BaseURL: server.URL}, zap.NewNop())
capsDefault, err := adapterDefault.Capabilities(context.Background())
if err != nil {
t.Fatalf("Capabilities failed: %v", err)
}
if capsDefault.MaxConcurrency != 4 {
t.Fatalf("expected default MaxConcurrency 4, got %d", capsDefault.MaxConcurrency)
}
}
func TestOllamaHandleCommandPassesNativeAPI(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost || r.URL.RequestURI() != "/api/show?verbose=true" {
t.Fatalf("unexpected request %s %s", r.Method, r.URL.RequestURI())
}
var body map[string]any
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Fatalf("decode body: %v", err)
}
if body["model"] != "gemma4:26b" {
t.Fatalf("body not passed: %+v", body)
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusAccepted)
_, _ = w.Write([]byte(`{"license":"ok"}`))
}))
defer server.Close()
adapter := New(config.OllamaConf{BaseURL: server.URL}, zap.NewNop())
resp, err := adapter.HandleCommand(context.Background(), noderuntime.CommandRequest{
RequestID: "ollama-1",
Type: noderuntime.CommandTypeOllamaAPI,
Adapter: "ollama",
Metadata: map[string]string{
"ollama_method": "POST",
"ollama_path": "/api/show?verbose=true",
"ollama_body": `{"model":"gemma4:26b"}`,
},
})
if err != nil {
t.Fatalf("HandleCommand failed: %v", err)
}
if resp.Result["status_code"] != "202" || resp.Result["content_type"] != "application/json" || resp.Result["body"] != `{"license":"ok"}` {
t.Fatalf("unexpected response: %+v", resp.Result)
}
}
func TestOllamaMultiEndpoint(t *testing.T) {
var (
gotModelA string
gotNumCtxA int
gotModelB string
gotNumCtxB int
)
makeServer := func(gotModel *string, gotNumCtx *int, reply string) *httptest.Server {
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/api/tags" {
_, _ = w.Write([]byte(`{"models":[]}`))
return
}
if r.URL.Path != "/api/chat" {
t.Errorf("unexpected path %s", r.URL.Path)
return
}
var req ollamaChatRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
t.Errorf("decode request: %v", err)
return
}
*gotModel = req.Model
if v, ok := req.Options["num_ctx"].(float64); ok {
*gotNumCtx = int(v)
}
w.Header().Set("Content-Type", "application/x-ndjson")
_, _ = w.Write([]byte(`{"message":{"role":"assistant","content":"` + reply + `"},"done":false}` + "\n"))
_, _ = w.Write([]byte(`{"done":true}` + "\n"))
}))
}
serverA := makeServer(&gotModelA, &gotNumCtxA, "a")
defer serverA.Close()
serverB := makeServer(&gotModelB, &gotNumCtxB, "b")
defer serverB.Close()
adapterA := New(config.OllamaConf{BaseURL: serverA.URL, ContextSize: 4096}, zap.NewNop(), "ollama-a")
adapterB := New(config.OllamaConf{BaseURL: serverB.URL, ContextSize: 8192}, zap.NewNop(), "ollama-b")
sinkA := &fakeSink{}
if err := adapterA.Execute(context.Background(), noderuntime.ExecutionSpec{
RunID: "run-a", Target: "model-a", Input: map[string]any{"prompt": "test"},
}, sinkA); err != nil {
t.Fatalf("A execute: %v", err)
}
sinkB := &fakeSink{}
if err := adapterB.Execute(context.Background(), noderuntime.ExecutionSpec{
RunID: "run-b", Target: "model-b", Input: map[string]any{"prompt": "test"},
}, sinkB); err != nil {
t.Fatalf("B execute: %v", err)
}
if gotModelA != "model-a" {
t.Errorf("A model: got %q, want model-a", gotModelA)
}
if gotNumCtxA != 4096 {
t.Errorf("A num_ctx: got %d, want 4096", gotNumCtxA)
}
if gotModelB != "model-b" {
t.Errorf("B model: got %q, want model-b", gotModelB)
}
if gotNumCtxB != 8192 {
t.Errorf("B num_ctx: got %d, want 8192", gotNumCtxB)
}
capsA, _ := adapterA.Capabilities(context.Background())
if capsA.InstanceKey != "ollama-a" {
t.Errorf("A InstanceKey: got %q, want ollama-a", capsA.InstanceKey)
}
capsB, _ := adapterB.Capabilities(context.Background())
if capsB.InstanceKey != "ollama-b" {
t.Errorf("B InstanceKey: got %q, want ollama-b", capsB.InstanceKey)
}
}
func TestOllamaContextSizeForcedOverRequestNumCtx(t *testing.T) {
var gotNumCtx int
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var req ollamaChatRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
t.Fatalf("decode request: %v", err)
}
if v, ok := req.Options["num_ctx"].(float64); ok {
gotNumCtx = int(v)
}
w.Header().Set("Content-Type", "application/x-ndjson")
_, _ = w.Write([]byte(`{"message":{"role":"assistant","content":"ok"},"done":false}` + "\n"))
_, _ = w.Write([]byte(`{"done":true}` + "\n"))
}))
defer server.Close()
// profile context_size=262144 must override request options.num_ctx=8192
adapter := New(config.OllamaConf{BaseURL: server.URL, ContextSize: 262144}, zap.NewNop())
sink := &fakeSink{}
err := adapter.Execute(context.Background(), noderuntime.ExecutionSpec{
RunID: "run-forced",
Target: "llama-test",
Input: map[string]any{
"prompt": "hi",
"options": map[string]any{"num_ctx": float64(8192), "temperature": 0.5},
},
}, sink)
if err != nil {
t.Fatalf("Execute failed: %v", err)
}
if gotNumCtx != 262144 {
t.Fatalf("Edge-owned context_size must override request num_ctx: got %d, want 262144", gotNumCtx)
}
}
func TestOllamaExecuteStreamsThinkingDeltasSeparately(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/chat" {
t.Fatalf("unexpected path %s", r.URL.Path)
}
w.Header().Set("Content-Type", "application/x-ndjson")
_, _ = w.Write([]byte(`{"message":{"role":"assistant","thinking":"thinking..."},"done":false}` + "\n"))
_, _ = w.Write([]byte(`{"message":{"role":"assistant","content":"answer"},"done":false}` + "\n"))
_, _ = w.Write([]byte(`{"done":true,"prompt_eval_count":3,"eval_count":2}` + "\n"))
}))
defer server.Close()
adapter := New(config.OllamaConf{BaseURL: server.URL}, zap.NewNop())
sink := &fakeSink{}
err := adapter.Execute(context.Background(), noderuntime.ExecutionSpec{
RunID: "run-thinking",
Target: "llama-test",
Input: map[string]any{"prompt": "say hello"},
}, sink)
if err != nil {
t.Fatalf("Execute failed: %v", err)
}
events := sink.all()
if len(events) != 4 {
t.Fatalf("expected 4 events, got %+v", events)
}
if events[0].Type != noderuntime.EventTypeStart {
t.Fatalf("expected start event, got %+v", events[0])
}
if events[1].Type != noderuntime.EventTypeReasoningDelta || events[1].Delta != "thinking..." {
t.Fatalf("expected reasoning_delta event, got %+v", events[1])
}
if events[2].Type != noderuntime.EventTypeDelta || events[2].Delta != "answer" {
t.Fatalf("expected delta event, got %+v", events[2])
}
if events[3].Type != noderuntime.EventTypeComplete {
t.Fatalf("expected complete event, got %+v", events[3])
}
if events[3].Usage == nil || events[3].Usage.InputTokens != 3 || events[3].Usage.OutputTokens != 2 {
t.Fatalf("unexpected usage: %+v", events[3].Usage)
}
}
func TestOllamaProbeProviderAvailability(t *testing.T) {
t.Run("200_ok_and_target_hit", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/tags" {
t.Fatalf("unexpected path %s", r.URL.Path)
}
_, _ = w.Write([]byte(`{"models":[{"name":"llama-a"},{"name":"llama-b"}]}`))
}))
defer server.Close()
adapter := New(config.OllamaConf{BaseURL: server.URL}, zap.NewNop())
res, err := adapter.ProbeProvider(context.Background(), "llama-a")
if err != nil {
t.Fatalf("ProbeProvider failed: %v", err)
}
if res.Status != noderuntime.ProviderStatusAvailable {
t.Errorf("expected Status available, got %s", res.Status)
}
if len(res.Targets) != 2 || res.Targets[0] != "llama-a" {
t.Errorf("unexpected Targets: %+v", res.Targets)
}
})
t.Run("200_ok_and_target_miss", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = w.Write([]byte(`{"models":[{"name":"llama-a"}]}`))
}))
defer server.Close()
adapter := New(config.OllamaConf{BaseURL: server.URL}, zap.NewNop())
res, err := adapter.ProbeProvider(context.Background(), "llama-b")
if err != nil {
t.Fatalf("ProbeProvider failed: %v", err)
}
if res.Status != noderuntime.ProviderStatusUnavailable {
t.Errorf("expected Status unavailable, got %s", res.Status)
}
if !strings.Contains(res.Detail, "not found") {
t.Errorf("expected 'not found' in detail, got %s", res.Detail)
}
})
t.Run("500_internal_error", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusInternalServerError)
}))
defer server.Close()
adapter := New(config.OllamaConf{BaseURL: server.URL}, zap.NewNop())
res, err := adapter.ProbeProvider(context.Background(), "llama-a")
if err != nil {
t.Fatalf("ProbeProvider failed: %v", err)
}
if res.Status != noderuntime.ProviderStatusUnavailable {
t.Errorf("expected Status unavailable, got %s", res.Status)
}
})
t.Run("empty_target", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = w.Write([]byte(`{"models":[{"name":"llama-a"}]}`))
}))
defer server.Close()
adapter := New(config.OllamaConf{BaseURL: server.URL}, zap.NewNop())
res, err := adapter.ProbeProvider(context.Background(), "")
if err != nil {
t.Fatalf("ProbeProvider failed: %v", err)
}
if res.Status != noderuntime.ProviderStatusAvailable {
t.Errorf("expected Status available, got %s", res.Status)
}
})
}