package providers import ( "context" "encoding/json" "net/http" "net/http/httptest" "strings" "testing" "time" ) func TestOllamaClient_Chat_Success(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { // Verify request if r.Method != "POST" { t.Errorf("Expected POST, got %s", r.Method) } if r.URL.Path != "/api/chat" { t.Errorf("Expected /api/chat, got %s", r.URL.Path) } // Decode request to verify it var req ollamaRequest json.NewDecoder(r.Body).Decode(&req) if req.Model != "llama2" { t.Errorf("Expected model llama2, got %s", req.Model) } // Return mock response resp := ollamaResponse{ Model: "llama2", CreatedAt: time.Now().Format(time.RFC3339), Message: ollamaMessageResp{ Role: "assistant", Content: "Hello! I'm Llama.", }, Done: true, DoneReason: "stop", PromptEvalCount: 10, EvalCount: 15, } w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(resp) })) defer server.Close() client := NewOllamaClient("llama2", server.URL) ctx := context.Background() resp, err := client.Chat(ctx, ChatRequest{ Messages: []Message{{Role: "user", Content: "Hello"}}, }) if err != nil { t.Fatalf("Unexpected error: %v", err) } if resp.Content != "Hello! I'm Llama." { t.Errorf("Expected content 'Hello! I'm Llama.', got '%s'", resp.Content) } if resp.Model != "llama2" { t.Errorf("Expected model 'llama2', got '%s'", resp.Model) } if resp.InputTokens != 10 { t.Errorf("Expected 10 input tokens, got %d", resp.InputTokens) } if resp.OutputTokens != 15 { t.Errorf("Expected 15 output tokens, got %d", resp.OutputTokens) } } func TestOllamaClient_Chat_WithSystemPrompt(t *testing.T) { var receivedMessages []ollamaMessage server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var req ollamaRequest json.NewDecoder(r.Body).Decode(&req) receivedMessages = req.Messages resp := ollamaResponse{ Model: "llama2", Message: ollamaMessageResp{Role: "assistant", Content: "Response"}, Done: true, } json.NewEncoder(w).Encode(resp) })) defer server.Close() client := NewOllamaClient("llama2", server.URL) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ Messages: []Message{{Role: "user", Content: "Hello"}}, System: "You are a helpful assistant", }) if err != nil { t.Fatalf("Unexpected error: %v", err) } // Verify system message was included if len(receivedMessages) < 2 { t.Fatalf("Expected at least 2 messages, got %d", len(receivedMessages)) } if receivedMessages[0].Role != "system" { t.Errorf("Expected first message to be system, got %s", receivedMessages[0].Role) } if receivedMessages[0].Content != "You are a helpful assistant" { t.Errorf("Expected system content, got %s", receivedMessages[0].Content) } } func TestOllamaClient_Chat_WithOptions(t *testing.T) { var receivedRequest ollamaRequest server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { json.NewDecoder(r.Body).Decode(&receivedRequest) resp := ollamaResponse{ Model: "llama2", Message: ollamaMessageResp{Role: "assistant", Content: "Response"}, Done: true, } json.NewEncoder(w).Encode(resp) })) defer server.Close() client := NewOllamaClient("llama2", server.URL) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ Messages: []Message{{Role: "user", Content: "Hello"}}, MaxTokens: 500, Temperature: 0.7, }) if err != nil { t.Fatalf("Unexpected error: %v", err) } if receivedRequest.Options == nil { t.Fatal("Expected options to be set") } if receivedRequest.Options.NumPredict != 500 { t.Errorf("Expected num_predict 500, got %d", receivedRequest.Options.NumPredict) } if receivedRequest.Options.Temperature != 0.7 { t.Errorf("Expected temperature 0.7, got %f", receivedRequest.Options.Temperature) } } func TestOllamaClient_Chat_APIError(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusInternalServerError) w.Write([]byte(`{"error": "Model not found"}`)) })) defer server.Close() client := NewOllamaClient("nonexistent", server.URL) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ Messages: []Message{{Role: "user", Content: "Hello"}}, }) if err == nil { t.Error("Expected error for API failure") } } func TestOllamaClient_Chat_NetworkError(t *testing.T) { client := NewOllamaClient("llama2", "http://localhost:99999") ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) defer cancel() _, err := client.Chat(ctx, ChatRequest{ Messages: []Message{{Role: "user", Content: "Hello"}}, }) if err == nil { t.Error("Expected error for network failure") } } func TestOllamaClient_Chat_ModelFallback(t *testing.T) { var receivedModel string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var req ollamaRequest json.NewDecoder(r.Body).Decode(&req) receivedModel = req.Model resp := ollamaResponse{ Model: req.Model, Message: ollamaMessageResp{Role: "assistant", Content: "Response"}, Done: true, } json.NewEncoder(w).Encode(resp) })) defer server.Close() // Client with no default model client := NewOllamaClient("", server.URL) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ Messages: []Message{{Role: "user", Content: "Hello"}}, // No model specified in request either }) if err != nil { t.Fatalf("Unexpected error: %v", err) } // Should fallback to llama3 if receivedModel != "llama3" { t.Errorf("Expected fallback to llama3, got %s", receivedModel) } } func TestOllamaClient_Chat_StripModelPrefix(t *testing.T) { var receivedModel string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var req ollamaRequest json.NewDecoder(r.Body).Decode(&req) receivedModel = req.Model resp := ollamaResponse{ Model: req.Model, Message: ollamaMessageResp{Role: "assistant", Content: "Response"}, Done: true, } json.NewEncoder(w).Encode(resp) })) defer server.Close() client := NewOllamaClient("default", server.URL) ctx := context.Background() _, err := client.Chat(ctx, ChatRequest{ Messages: []Message{{Role: "user", Content: "Hello"}}, Model: "ollama:llama2", // With prefix }) if err != nil { t.Fatalf("Unexpected error: %v", err) } // Should strip the prefix if receivedModel != "llama2" { t.Errorf("Expected model 'llama2' (prefix stripped), got %s", receivedModel) } } func TestOllamaClient_TestConnection_Success(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/api/version" { t.Errorf("Expected /api/version, got %s", r.URL.Path) } w.WriteHeader(http.StatusOK) w.Write([]byte(`{"version": "0.1.0"}`)) })) defer server.Close() client := NewOllamaClient("llama2", server.URL) ctx := context.Background() err := client.TestConnection(ctx) if err != nil { t.Errorf("Unexpected error: %v", err) } } func TestOllamaClient_TestConnection_Failure(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusServiceUnavailable) })) defer server.Close() client := NewOllamaClient("llama2", server.URL) ctx := context.Background() err := client.TestConnection(ctx) if err == nil { t.Error("Expected error for failed connection test") } } func TestOllamaClient_ListModels_Success(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/api/tags" { t.Errorf("Expected /api/tags, got %s", r.URL.Path) } resp := struct { Models []struct { Name string `json:"name"` ModifiedAt string `json:"modified_at"` Size int64 `json:"size"` } `json:"models"` }{ Models: []struct { Name string `json:"name"` ModifiedAt string `json:"modified_at"` Size int64 `json:"size"` }{ {Name: "llama2:latest", ModifiedAt: "2024-01-01T00:00:00Z", Size: 1000000}, {Name: "mistral:latest", ModifiedAt: "2024-01-01T00:00:00Z", Size: 2000000}, }, } json.NewEncoder(w).Encode(resp) })) defer server.Close() client := NewOllamaClient("llama2", server.URL) ctx := context.Background() models, err := client.ListModels(ctx) if err != nil { t.Fatalf("Unexpected error: %v", err) } if len(models) != 2 { t.Errorf("Expected 2 models, got %d", len(models)) } if models[0].ID != "llama2:latest" { t.Errorf("Expected first model 'llama2:latest', got '%s'", models[0].ID) } } func TestOllamaClient_ListModels_Failure(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusInternalServerError) w.Write([]byte("Internal error")) })) defer server.Close() client := NewOllamaClient("llama2", server.URL) ctx := context.Background() _, err := client.ListModels(ctx) if err == nil { t.Error("Expected error for failed list models") } } func TestNewOllamaClient_NormalizesBaseURL(t *testing.T) { tests := []struct { in string expected string }{ {"", "http://localhost:11434"}, {"http://example:11434", "http://example:11434"}, {"http://example:11434/", "http://example:11434"}, {"http://example:11434/api", "http://example:11434"}, {"http://example:11434/api/", "http://example:11434"}, } for _, tc := range tests { t.Run(tc.in, func(t *testing.T) { client := NewOllamaClient("llama3", tc.in) if client.baseURL != tc.expected { t.Fatalf("baseURL = %q, want %q", client.baseURL, tc.expected) } }) } } func TestOllamaClient_Chat_ToolCallsResponse(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { resp := ollamaResponse{ Model: "llama3", Message: ollamaMessageResp{ Role: "assistant", Content: "", ToolCalls: []ollamaToolCall{ { ID: "call_1", Function: ollamaFunctionCall{ Name: "get_time", Arguments: map[string]interface{}{"tz": "UTC"}, }, }, }, }, Done: true, DoneReason: "stop", } _ = json.NewEncoder(w).Encode(resp) })) defer server.Close() client := NewOllamaClient("llama3", server.URL) out, err := client.Chat(context.Background(), ChatRequest{ Messages: []Message{{Role: "user", Content: "What time is it?"}}, }) if err != nil { t.Fatalf("Chat: %v", err) } if out.StopReason != "tool_use" { t.Fatalf("StopReason = %q, want tool_use", out.StopReason) } if len(out.ToolCalls) != 1 { t.Fatalf("ToolCalls = %d, want 1", len(out.ToolCalls)) } if out.ToolCalls[0].ID != "call_1" || out.ToolCalls[0].Name != "get_time" { t.Fatalf("unexpected tool call: %+v", out.ToolCalls[0]) } if out.ToolCalls[0].Input["tz"] != "UTC" { t.Fatalf("unexpected tool call input: %+v", out.ToolCalls[0].Input) } } func TestOllamaClient_Chat_ToolCallsAndToolResultsInRequest(t *testing.T) { var got ollamaRequest server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _ = json.NewDecoder(r.Body).Decode(&got) _ = json.NewEncoder(w).Encode(ollamaResponse{ Model: got.Model, Message: ollamaMessageResp{Role: "assistant", Content: "ok"}, Done: true, }) })) defer server.Close() client := NewOllamaClient("llama3", server.URL) _, err := client.Chat(context.Background(), ChatRequest{ System: "system prompt", Messages: []Message{ { Role: "assistant", Content: "calling tool", ToolCalls: []ToolCall{ {Name: "get_time", Input: map[string]any{"tz": "UTC"}}, }, }, { Role: "assistant", ToolResult: &ToolResult{Content: "{\"time\":\"00:00\"}"}, }, }, Tools: []Tool{ { Type: "function", Name: "get_time", Description: "get time", InputSchema: map[string]any{"type": "object"}, }, { Type: "web_search", Name: "search", }, }, }) if err != nil { t.Fatalf("Chat: %v", err) } if got.Messages[0].Role != "system" || got.Messages[0].Content != "system prompt" { t.Fatalf("expected system message first, got: %+v", got.Messages[0]) } var sawAssistantToolCall bool var sawToolResult bool for _, m := range got.Messages { if m.Role == "assistant" && len(m.ToolCalls) == 1 && m.ToolCalls[0].Function.Name == "get_time" { if m.ToolCalls[0].Function.Arguments["tz"] != "UTC" { t.Fatalf("unexpected tool call args: %+v", m.ToolCalls[0].Function.Arguments) } sawAssistantToolCall = true } if m.Role == "tool" && strings.Contains(m.Content, "00:00") { sawToolResult = true } } if !sawAssistantToolCall { t.Fatalf("expected assistant tool call message in request, got: %+v", got.Messages) } if !sawToolResult { t.Fatalf("expected tool result message in request, got: %+v", got.Messages) } if len(got.Tools) != 1 || got.Tools[0].Function.Name != "get_time" { t.Fatalf("expected only function tools to be included, got: %+v", got.Tools) } }