gsadmin/pulse
1
0
Fork 0
mirror of https://github.com/rcourtman/Pulse.git synced 2026-09-21 18:53:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: Add configurable AI request timeout for slow hardware. Related to #880

Browse Source 
Adds RequestTimeoutSeconds to AI config (default 300s / 5 min).
Users with low-power hardware running Ollama can increase this
value in Settings to prevent timeouts on slower inference.
This commit is contained in:
rcourtman
2025-12-24 10:52:11 +00:00
parent 2798877aa7
commit 708dc964df
13 changed files with 137 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/ai/providers/anthropic.go
+9 -5
View File
@@ -29,23 +29,27 @@ type AnthropicClient struct {
}
// NewAnthropicClient creates a new Anthropic API client
func NewAnthropicClient(apiKey, model string) *AnthropicClient {
return NewAnthropicClientWithBaseURL(apiKey, model, anthropicAPIURL)
// timeout is optional - pass 0 to use the default 5 minute timeout
func NewAnthropicClient(apiKey, model string, timeout time.Duration) *AnthropicClient {
return NewAnthropicClientWithBaseURL(apiKey, model, anthropicAPIURL, timeout)
}
// NewAnthropicClientWithBaseURL creates a new Anthropic client using a custom messages endpoint.
// This is useful for testing and for deployments that route requests through a proxy.
func NewAnthropicClientWithBaseURL(apiKey, model, baseURL string) *AnthropicClient {
// timeout is optional - pass 0 to use the default 5 minute timeout
func NewAnthropicClientWithBaseURL(apiKey, model, baseURL string, timeout time.Duration) *AnthropicClient {
if baseURL == "" {
baseURL = anthropicAPIURL
}
if timeout <= 0 {
timeout = 300 * time.Second // Default 5 minutes
}
return &AnthropicClient{
apiKey: apiKey,
model: model,
baseURL: baseURL,
client: &http.Client{
// 5 minutes - Opus and other large models can take a very long time
Timeout: 300 * time.Second,
Timeout: timeout,
},
}
}
internal/ai/providers/anthropic_oauth.go
+9 -4
View File
@@ -320,16 +320,21 @@ type AnthropicOAuthClient struct {
}
// NewAnthropicOAuthClient creates a new Anthropic client using OAuth tokens
func NewAnthropicOAuthClient(accessToken, refreshToken string, expiresAt time.Time, model string) *AnthropicOAuthClient {
return NewAnthropicOAuthClientWithBaseURL(accessToken, refreshToken, expiresAt, model, "https://api.anthropic.com/v1/messages?beta=true")
// timeout is optional - pass 0 to use the default 5 minute timeout
func NewAnthropicOAuthClient(accessToken, refreshToken string, expiresAt time.Time, model string, timeout time.Duration) *AnthropicOAuthClient {
return NewAnthropicOAuthClientWithBaseURL(accessToken, refreshToken, expiresAt, model, "https://api.anthropic.com/v1/messages?beta=true", timeout)
}
// NewAnthropicOAuthClientWithBaseURL creates a new Anthropic OAuth client using a custom messages endpoint.
// This is useful for testing and for deployments that route requests through a proxy.
func NewAnthropicOAuthClientWithBaseURL(accessToken, refreshToken string, expiresAt time.Time, model, baseURL string) *AnthropicOAuthClient {
// timeout is optional - pass 0 to use the default 5 minute timeout
func NewAnthropicOAuthClientWithBaseURL(accessToken, refreshToken string, expiresAt time.Time, model, baseURL string, timeout time.Duration) *AnthropicOAuthClient {
if baseURL == "" {
baseURL = "https://api.anthropic.com/v1/messages?beta=true"
}
if timeout <= 0 {
timeout = 300 * time.Second // Default 5 minutes
}
return &AnthropicOAuthClient{
accessToken: accessToken,
refreshToken: refreshToken,
@@ -337,7 +342,7 @@ func NewAnthropicOAuthClientWithBaseURL(accessToken, refreshToken string, expire
model: model,
baseURL: baseURL,
client: &http.Client{
Timeout: 300 * time.Second,
Timeout: timeout,
},
}
}
internal/ai/providers/anthropic_oauth_test.go
+3 -1
View File
@@ -265,7 +265,7 @@ func TestAnthropicOAuthClient_forceRefreshToken_UpdatesAndCallsCallback(t *testi
oauthTokenURL = server.URL
oauthHTTPClient = server.Client()
client := NewAnthropicOAuthClient("access_old", "refresh_old", time.Now().Add(-time.Minute), "claude-3")
client := NewAnthropicOAuthClient("access_old", "refresh_old", time.Now().Add(-time.Minute), "claude-3", 0)
var cbTokens *OAuthTokens
client.SetTokenRefreshCallback(func(tokens *OAuthTokens) { cbTokens = tokens })
@@ -361,6 +361,7 @@ func TestAnthropicOAuthClient_Chat_RefreshesOn401AndRetriesImmediately(t *testin
time.Now().Add(10*time.Minute), // valid token, so refresh is driven by 401
"claude-3",
server.URL+"/v1/messages?beta=true",
0,
)
client.client = server.Client()
@@ -401,6 +402,7 @@ func TestAnthropicOAuthClient_ListModels_UsesConfiguredHost(t *testing.T) {
time.Now().Add(10*time.Minute),
"claude-3",
server.URL+"/v1/messages?beta=true",
0,
)
client.client = server.Client()
internal/ai/providers/anthropic_test.go
+7 -7
View File
@@ -12,14 +12,14 @@ import (
// Anthropic tests focus on request/response correctness and endpoint behavior.
func TestAnthropicClient_Name(t *testing.T) {
client := NewAnthropicClient("test-key", "claude-3-5-sonnet")
client := NewAnthropicClient("test-key", "claude-3-5-sonnet", 0)
if client.Name() != "anthropic" {
t.Errorf("Expected 'anthropic', got '%s'", client.Name())
}
}
func TestAnthropicClient_Chat_ContextCanceled(t *testing.T) {
client := NewAnthropicClient("test-key", "claude-3-5-sonnet")
client := NewAnthropicClient("test-key", "claude-3-5-sonnet", 0)
ctx, cancel := context.WithCancel(context.Background())
cancel() // Cancel immediately
@@ -34,7 +34,7 @@ func TestAnthropicClient_Chat_ContextCanceled(t *testing.T) {
}
func TestAnthropicClient_Chat_Timeout(t *testing.T) {
client := NewAnthropicClient("test-key", "claude-3-5-sonnet")
client := NewAnthropicClient("test-key", "claude-3-5-sonnet", 0)
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Millisecond)
defer cancel()
@@ -108,7 +108,7 @@ func TestAnthropicClient_Chat_Success_TextAndToolCalls(t *testing.T) {
}))
defer server.Close()
client := NewAnthropicClientWithBaseURL("test-key", "claude-3-5-sonnet", server.URL+"/v1/messages")
client := NewAnthropicClientWithBaseURL("test-key", "claude-3-5-sonnet", server.URL+"/v1/messages", 0)
out, err := client.Chat(context.Background(), ChatRequest{
System: "You are helpful",
Messages: []Message{{Role: "user", Content: "Hello"}},
@@ -162,7 +162,7 @@ func TestAnthropicClient_Chat_ToolResultInRequest(t *testing.T) {
}))
defer server.Close()
client := NewAnthropicClientWithBaseURL("test-key", "claude-3-5-sonnet", server.URL+"/v1/messages")
client := NewAnthropicClientWithBaseURL("test-key", "claude-3-5-sonnet", server.URL+"/v1/messages", 0)
_, err := client.Chat(context.Background(), ChatRequest{
Messages: []Message{
{
@@ -227,7 +227,7 @@ func TestAnthropicClient_ListModels_UsesConfiguredHost(t *testing.T) {
}))
defer server.Close()
client := NewAnthropicClientWithBaseURL("test-key", "claude-3-5-sonnet", server.URL+"/v1/messages")
client := NewAnthropicClientWithBaseURL("test-key", "claude-3-5-sonnet", server.URL+"/v1/messages", 0)
models, err := client.ListModels(context.Background())
if err != nil {
t.Fatalf("ListModels: %v", err)
@@ -248,7 +248,7 @@ func TestAnthropicClient_TestConnection_CallsListModels(t *testing.T) {
}))
defer server.Close()
client := NewAnthropicClientWithBaseURL("test-key", "claude-3-5-sonnet", server.URL+"/v1/messages")
client := NewAnthropicClientWithBaseURL("test-key", "claude-3-5-sonnet", server.URL+"/v1/messages", 0)
if err := client.TestConnection(context.Background()); err != nil {
t.Fatalf("TestConnection: %v", err)
}
internal/ai/providers/factory.go
+18 -10
View File
@@ -23,12 +23,15 @@ func NewFromConfig(cfg *config.AIConfig) (Provider, error) {
return providerClient, nil
}
// Get the configured timeout
timeout := cfg.GetRequestTimeout()
// Fall back to legacy single-provider format
switch cfg.Provider {
case config.AIProviderAnthropic:
// If we have an API key (from direct entry or OAuth-created), use regular client
if cfg.APIKey != "" {
return NewAnthropicClient(cfg.APIKey, cfg.GetModel()), nil
return NewAnthropicClient(cfg.APIKey, cfg.GetModel(), timeout), nil
}
// Pro/Max users without org:create_api_key will use OAuth tokens directly
if cfg.IsUsingOAuth() && cfg.OAuthAccessToken != "" {
@@ -37,6 +40,7 @@ func NewFromConfig(cfg *config.AIConfig) (Provider, error) {
cfg.OAuthRefreshToken,
cfg.OAuthExpiresAt,
cfg.GetModel(),
timeout,
)
return client, nil
}
@@ -46,23 +50,23 @@ func NewFromConfig(cfg *config.AIConfig) (Provider, error) {
if cfg.APIKey == "" {
return nil, fmt.Errorf("OpenAI API key is required")
}
return NewOpenAIClient(cfg.APIKey, cfg.GetModel(), cfg.GetBaseURL()), nil
return NewOpenAIClient(cfg.APIKey, cfg.GetModel(), cfg.GetBaseURL(), timeout), nil
case config.AIProviderOllama:
return NewOllamaClient(cfg.GetModel(), cfg.GetBaseURL()), nil
return NewOllamaClient(cfg.GetModel(), cfg.GetBaseURL(), timeout), nil
case config.AIProviderDeepSeek:
if cfg.APIKey == "" {
return nil, fmt.Errorf("DeepSeek API key is required")
}
// DeepSeek uses OpenAI-compatible API
return NewOpenAIClient(cfg.APIKey, cfg.GetModel(), cfg.GetBaseURL()), nil
return NewOpenAIClient(cfg.APIKey, cfg.GetModel(), cfg.GetBaseURL(), timeout), nil
case config.AIProviderGemini:
if cfg.APIKey == "" {
return nil, fmt.Errorf("Gemini API key is required")
}
return NewGeminiClient(cfg.APIKey, cfg.GetModel(), cfg.GetBaseURL()), nil
return NewGeminiClient(cfg.APIKey, cfg.GetModel(), cfg.GetBaseURL(), timeout), nil
default:
return nil, fmt.Errorf("unknown provider: %s", cfg.Provider)
@@ -75,6 +79,9 @@ func NewForProvider(cfg *config.AIConfig, provider, model string) (Provider, err
return nil, fmt.Errorf("AI config is nil")
}
// Get the configured timeout
timeout := cfg.GetRequestTimeout()
switch provider {
case config.AIProviderAnthropic:
// Check for OAuth first
@@ -84,6 +91,7 @@ func NewForProvider(cfg *config.AIConfig, provider, model string) (Provider, err
cfg.OAuthRefreshToken,
cfg.OAuthExpiresAt,
model,
timeout,
), nil
}
// Then check for per-provider API key
@@ -91,7 +99,7 @@ func NewForProvider(cfg *config.AIConfig, provider, model string) (Provider, err
if apiKey == "" {
return nil, fmt.Errorf("Anthropic API key not configured")
}
return NewAnthropicClient(apiKey, model), nil
return NewAnthropicClient(apiKey, model, timeout), nil
case config.AIProviderOpenAI:
apiKey := cfg.GetAPIKeyForProvider(config.AIProviderOpenAI)
@@ -99,7 +107,7 @@ func NewForProvider(cfg *config.AIConfig, provider, model string) (Provider, err
return nil, fmt.Errorf("OpenAI API key not configured")
}
baseURL := cfg.GetBaseURLForProvider(config.AIProviderOpenAI)
return NewOpenAIClient(apiKey, model, baseURL), nil
return NewOpenAIClient(apiKey, model, baseURL, timeout), nil
case config.AIProviderDeepSeek:
apiKey := cfg.GetAPIKeyForProvider(config.AIProviderDeepSeek)
@@ -107,11 +115,11 @@ func NewForProvider(cfg *config.AIConfig, provider, model string) (Provider, err
return nil, fmt.Errorf("DeepSeek API key not configured")
}
baseURL := cfg.GetBaseURLForProvider(config.AIProviderDeepSeek)
return NewOpenAIClient(apiKey, model, baseURL), nil
return NewOpenAIClient(apiKey, model, baseURL, timeout), nil
case config.AIProviderOllama:
baseURL := cfg.GetBaseURLForProvider(config.AIProviderOllama)
return NewOllamaClient(model, baseURL), nil
return NewOllamaClient(model, baseURL, timeout), nil
case config.AIProviderGemini:
apiKey := cfg.GetAPIKeyForProvider(config.AIProviderGemini)
@@ -119,7 +127,7 @@ func NewForProvider(cfg *config.AIConfig, provider, model string) (Provider, err
return nil, fmt.Errorf("Gemini API key not configured")
}
baseURL := cfg.GetBaseURLForProvider(config.AIProviderGemini)
return NewGeminiClient(apiKey, model, baseURL), nil
return NewGeminiClient(apiKey, model, baseURL, timeout), nil
default:
return nil, fmt.Errorf("unknown provider: %s", provider)
internal/ai/providers/gemini.go
+6 -3
View File
@@ -28,7 +28,8 @@ type GeminiClient struct {
}
// NewGeminiClient creates a new Gemini API client
func NewGeminiClient(apiKey, model, baseURL string) *GeminiClient {
// timeout is optional - pass 0 to use the default 5 minute timeout
func NewGeminiClient(apiKey, model, baseURL string, timeout time.Duration) *GeminiClient {
if baseURL == "" {
baseURL = geminiAPIURL
}
@@ -36,13 +37,15 @@ func NewGeminiClient(apiKey, model, baseURL string) *GeminiClient {
if strings.HasPrefix(model, "gemini:") {
model = strings.TrimPrefix(model, "gemini:")
}
if timeout <= 0 {
timeout = 300 * time.Second // Default 5 minutes
}
return &GeminiClient{
apiKey: apiKey,
model: model,
baseURL: baseURL,
client: &http.Client{
// 5 minutes timeout - large models can take a long time
Timeout: 300 * time.Second,
Timeout: timeout,
},
}
}
internal/ai/providers/gemini_test.go
+17 -17
View File
@@ -13,7 +13,7 @@ import (
func TestNewGeminiClient(t *testing.T) {
t.Parallel()
client := NewGeminiClient("test-api-key", "gemini-pro", "")
client := NewGeminiClient("test-api-key", "gemini-pro", "", 0)
if client == nil {
t.Fatal("expected non-nil client")
}
@@ -31,7 +31,7 @@ func TestNewGeminiClient(t *testing.T) {
func TestNewGeminiClient_StripPrefix(t *testing.T) {
t.Parallel()
client := NewGeminiClient("api-key", "gemini:gemini-1.5-pro", "")
client := NewGeminiClient("api-key", "gemini:gemini-1.5-pro", "", 0)
if client.model != "gemini-1.5-pro" {
t.Errorf("expected model with prefix stripped, got %q", client.model)
}
@@ -40,7 +40,7 @@ func TestNewGeminiClient_StripPrefix(t *testing.T) {
func TestNewGeminiClient_CustomBaseURL(t *testing.T) {
t.Parallel()
client := NewGeminiClient("api-key", "gemini-pro", "https://custom.api.example.com")
client := NewGeminiClient("api-key", "gemini-pro", "https://custom.api.example.com", 0)
if client.baseURL != "https://custom.api.example.com" {
t.Errorf("expected custom baseURL, got %q", client.baseURL)
}
@@ -49,7 +49,7 @@ func TestNewGeminiClient_CustomBaseURL(t *testing.T) {
func TestGeminiClient_Name(t *testing.T) {
t.Parallel()
client := NewGeminiClient("key", "model", "")
client := NewGeminiClient("key", "model", "", 0)
if client.Name() != "gemini" {
t.Errorf("expected Name() to return 'gemini', got %q", client.Name())
}
@@ -87,7 +87,7 @@ func TestGeminiClient_Chat_Success(t *testing.T) {
}))
defer server.Close()
client := NewGeminiClient("test-key", "gemini-pro", server.URL)
client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0)
ctx := context.Background()
resp, err := client.Chat(ctx, ChatRequest{
@@ -135,7 +135,7 @@ func TestGeminiClient_Chat_WithSystemPrompt(t *testing.T) {
}))
defer server.Close()
client := NewGeminiClient("test-key", "gemini-pro", server.URL)
client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0)
ctx := context.Background()
_, err := client.Chat(ctx, ChatRequest{
@@ -178,7 +178,7 @@ func TestGeminiClient_Chat_WithMaxTokens(t *testing.T) {
}))
defer server.Close()
client := NewGeminiClient("test-key", "gemini-pro", server.URL)
client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0)
ctx := context.Background()
_, err := client.Chat(ctx, ChatRequest{
@@ -221,7 +221,7 @@ func TestGeminiClient_Chat_ToolCalls(t *testing.T) {
}))
defer server.Close()
client := NewGeminiClient("test-key", "gemini-pro", server.URL)
client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0)
ctx := context.Background()
resp, err := client.Chat(ctx, ChatRequest{
@@ -269,7 +269,7 @@ func TestGeminiClient_Chat_APIError(t *testing.T) {
}))
defer server.Close()
client := NewGeminiClient("invalid-key", "gemini-pro", server.URL)
client := NewGeminiClient("invalid-key", "gemini-pro", server.URL, 0)
ctx := context.Background()
_, err := client.Chat(ctx, ChatRequest{
@@ -293,7 +293,7 @@ func TestGeminiClient_Chat_NoCandidates(t *testing.T) {
}))
defer server.Close()
client := NewGeminiClient("test-key", "gemini-pro", server.URL)
client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0)
ctx := context.Background()
_, err := client.Chat(ctx, ChatRequest{
@@ -325,7 +325,7 @@ func TestGeminiClient_Chat_SafetyBlocked(t *testing.T) {
}))
defer server.Close()
client := NewGeminiClient("test-key", "gemini-pro", server.URL)
client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0)
ctx := context.Background()
_, err := client.Chat(ctx, ChatRequest{
@@ -351,7 +351,7 @@ func TestGeminiClient_Chat_PromptBlocked(t *testing.T) {
}))
defer server.Close()
client := NewGeminiClient("test-key", "gemini-pro", server.URL)
client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0)
ctx := context.Background()
_, err := client.Chat(ctx, ChatRequest{
@@ -413,7 +413,7 @@ func TestGeminiClient_ListModels_Success(t *testing.T) {
}))
defer server.Close()
client := NewGeminiClient("test-key", "gemini-pro", server.URL)
client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0)
ctx := context.Background()
models, err := client.ListModels(ctx)
@@ -435,7 +435,7 @@ func TestGeminiClient_ListModels_Error(t *testing.T) {
}))
defer server.Close()
client := NewGeminiClient("test-key", "gemini-pro", server.URL)
client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0)
ctx := context.Background()
_, err := client.ListModels(ctx)
@@ -454,7 +454,7 @@ func TestGeminiClient_TestConnection(t *testing.T) {
}))
defer server.Close()
client := NewGeminiClient("test-key", "gemini-pro", server.URL)
client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0)
ctx := context.Background()
err := client.TestConnection(ctx)
@@ -465,7 +465,7 @@ func TestGeminiClient_TestConnection(t *testing.T) {
}
func TestGeminiClient_Chat_NetworkError(t *testing.T) {
client := NewGeminiClient("test-key", "gemini-pro", "http://localhost:99999")
client := NewGeminiClient("test-key", "gemini-pro", "http://localhost:99999", 0)
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
@@ -497,7 +497,7 @@ func TestGeminiClient_Chat_RoleConversion(t *testing.T) {
}))
defer server.Close()
client := NewGeminiClient("test-key", "gemini-pro", server.URL)
client := NewGeminiClient("test-key", "gemini-pro", server.URL, 0)
ctx := context.Background()
_, err := client.Chat(ctx, ChatRequest{
internal/ai/providers/integration_test.go
+18 -18
View File
@@ -27,7 +27,7 @@ func getOllamaURL() string {
}
func TestIntegration_Ollama_TestConnection(t *testing.T) {
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
@@ -40,7 +40,7 @@ func TestIntegration_Ollama_TestConnection(t *testing.T) {
}
func TestIntegration_Ollama_ListModels(t *testing.T) {
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
@@ -61,7 +61,7 @@ func TestIntegration_Ollama_ListModels(t *testing.T) {
}
func TestIntegration_Ollama_SimpleChat(t *testing.T) {
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
@@ -86,7 +86,7 @@ func TestIntegration_Ollama_SimpleChat(t *testing.T) {
}
func TestIntegration_Ollama_SystemPrompt(t *testing.T) {
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
@@ -116,7 +116,7 @@ func TestIntegration_Ollama_SystemPrompt(t *testing.T) {
}
func TestIntegration_Ollama_MultiTurnConversation(t *testing.T) {
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
defer cancel()
@@ -155,7 +155,7 @@ func TestIntegration_Ollama_MultiTurnConversation(t *testing.T) {
}
func TestIntegration_Ollama_TokenCounting(t *testing.T) {
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
@@ -183,7 +183,7 @@ func TestIntegration_Ollama_TokenCounting(t *testing.T) {
}
func TestIntegration_Ollama_ErrorHandling_BadModel(t *testing.T) {
client := providers.NewOllamaClient("nonexistent-model-12345", getOllamaURL())
client := providers.NewOllamaClient("nonexistent-model-12345", getOllamaURL(, 0))
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
@@ -202,7 +202,7 @@ func TestIntegration_Ollama_ErrorHandling_BadModel(t *testing.T) {
}
func TestIntegration_Ollama_Timeout(t *testing.T) {
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
// Very short timeout - should fail
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Millisecond)
@@ -225,7 +225,7 @@ func TestIntegration_Ollama_Timeout(t *testing.T) {
// --- More useful tests below ---
func TestIntegration_Ollama_JSONOutput(t *testing.T) {
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
@@ -267,7 +267,7 @@ func TestIntegration_Ollama_JSONOutput(t *testing.T) {
}
func TestIntegration_Ollama_LongResponse(t *testing.T) {
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
defer cancel()
@@ -292,7 +292,7 @@ func TestIntegration_Ollama_LongResponse(t *testing.T) {
}
func TestIntegration_Ollama_EmptyMessage(t *testing.T) {
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
@@ -313,7 +313,7 @@ func TestIntegration_Ollama_EmptyMessage(t *testing.T) {
}
func TestIntegration_Ollama_SpecialCharacters(t *testing.T) {
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
@@ -336,7 +336,7 @@ func TestIntegration_Ollama_SpecialCharacters(t *testing.T) {
}
func TestIntegration_Ollama_ConcurrentRequests(t *testing.T) {
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
const numRequests = 3
results := make(chan error, numRequests)
@@ -372,7 +372,7 @@ func TestIntegration_Ollama_ConcurrentRequests(t *testing.T) {
func TestIntegration_Ollama_InfrastructureAnalysis(t *testing.T) {
// This simulates what Pulse actually does - send infrastructure context
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
defer cancel()
@@ -416,7 +416,7 @@ What should I investigate first?`
}
func TestIntegration_Ollama_ModelName_Preserved(t *testing.T) {
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
@@ -441,7 +441,7 @@ func TestIntegration_Ollama_ModelName_Preserved(t *testing.T) {
}
func TestIntegration_Ollama_StopReason(t *testing.T) {
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
@@ -461,7 +461,7 @@ func TestIntegration_Ollama_StopReason(t *testing.T) {
}
func TestIntegration_Ollama_VeryLongInput(t *testing.T) {
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
defer cancel()
@@ -486,7 +486,7 @@ func TestIntegration_Ollama_VeryLongInput(t *testing.T) {
}
func TestIntegration_Ollama_RapidFireRequests(t *testing.T) {
client := providers.NewOllamaClient("tinyllama", getOllamaURL())
client := providers.NewOllamaClient("tinyllama", getOllamaURL(, 0))
// Send 5 requests in rapid succession
for i := 0; i < 5; i++ {
internal/ai/providers/ollama.go
+6 -2
View File
@@ -19,7 +19,8 @@ type OllamaClient struct {
}
// NewOllamaClient creates a new Ollama API client
func NewOllamaClient(model, baseURL string) *OllamaClient {
// timeout is optional - pass 0 to use the default 5 minute timeout
func NewOllamaClient(model, baseURL string, timeout time.Duration) *OllamaClient {
if baseURL == "" {
baseURL = "http://localhost:11434"
}
@@ -28,11 +29,14 @@ func NewOllamaClient(model, baseURL string) *OllamaClient {
baseURL = strings.TrimSuffix(baseURL, "/")
baseURL = strings.TrimSuffix(baseURL, "/api")
baseURL = strings.TrimSuffix(baseURL, "/") // In case it was /api/
if timeout <= 0 {
timeout = 300 * time.Second // Default 5 minutes
}
return &OllamaClient{
model: model,
baseURL: baseURL,
client: &http.Client{
Timeout: 300 * time.Second, // Local models can be slow
Timeout: timeout,
},
}
}
internal/ai/providers/ollama_test.go
+14 -14
View File
@@ -47,7 +47,7 @@ func TestOllamaClient_Chat_Success(t *testing.T) {
}))
defer server.Close()
client := NewOllamaClient("llama2", server.URL)
client := NewOllamaClient("llama2", server.URL, 0)
ctx := context.Background()
resp, err := client.Chat(ctx, ChatRequest{
@@ -92,7 +92,7 @@ func TestOllamaClient_Chat_WithSystemPrompt(t *testing.T) {
}))
defer server.Close()
client := NewOllamaClient("llama2", server.URL)
client := NewOllamaClient("llama2", server.URL, 0)
ctx := context.Background()
_, err := client.Chat(ctx, ChatRequest{
@@ -133,7 +133,7 @@ func TestOllamaClient_Chat_WithOptions(t *testing.T) {
}))
defer server.Close()
client := NewOllamaClient("llama2", server.URL)
client := NewOllamaClient("llama2", server.URL, 0)
ctx := context.Background()
_, err := client.Chat(ctx, ChatRequest{
@@ -166,7 +166,7 @@ func TestOllamaClient_Chat_APIError(t *testing.T) {
}))
defer server.Close()
client := NewOllamaClient("nonexistent", server.URL)
client := NewOllamaClient("nonexistent", server.URL, 0)
ctx := context.Background()
_, err := client.Chat(ctx, ChatRequest{
@@ -179,7 +179,7 @@ func TestOllamaClient_Chat_APIError(t *testing.T) {
}
func TestOllamaClient_Chat_NetworkError(t *testing.T) {
client := NewOllamaClient("llama2", "http://localhost:99999")
client := NewOllamaClient("llama2", "http://localhost:99999", 0)
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
@@ -211,7 +211,7 @@ func TestOllamaClient_Chat_ModelFallback(t *testing.T) {
defer server.Close()
// Client with no default model
client := NewOllamaClient("", server.URL)
client := NewOllamaClient("", server.URL, 0)
ctx := context.Background()
_, err := client.Chat(ctx, ChatRequest{
@@ -246,7 +246,7 @@ func TestOllamaClient_Chat_StripModelPrefix(t *testing.T) {
}))
defer server.Close()
client := NewOllamaClient("default", server.URL)
client := NewOllamaClient("default", server.URL, 0)
ctx := context.Background()
_, err := client.Chat(ctx, ChatRequest{
@@ -274,7 +274,7 @@ func TestOllamaClient_TestConnection_Success(t *testing.T) {
}))
defer server.Close()
client := NewOllamaClient("llama2", server.URL)
client := NewOllamaClient("llama2", server.URL, 0)
ctx := context.Background()
err := client.TestConnection(ctx)
@@ -290,7 +290,7 @@ func TestOllamaClient_TestConnection_Failure(t *testing.T) {
}))
defer server.Close()
client := NewOllamaClient("llama2", server.URL)
client := NewOllamaClient("llama2", server.URL, 0)
ctx := context.Background()
err := client.TestConnection(ctx)
@@ -327,7 +327,7 @@ func TestOllamaClient_ListModels_Success(t *testing.T) {
}))
defer server.Close()
client := NewOllamaClient("llama2", server.URL)
client := NewOllamaClient("llama2", server.URL, 0)
ctx := context.Background()
models, err := client.ListModels(ctx)
@@ -352,7 +352,7 @@ func TestOllamaClient_ListModels_Failure(t *testing.T) {
}))
defer server.Close()
client := NewOllamaClient("llama2", server.URL)
client := NewOllamaClient("llama2", server.URL, 0)
ctx := context.Background()
_, err := client.ListModels(ctx)
@@ -376,7 +376,7 @@ func TestNewOllamaClient_NormalizesBaseURL(t *testing.T) {
for _, tc := range tests {
t.Run(tc.in, func(t *testing.T) {
client := NewOllamaClient("llama3", tc.in)
client := NewOllamaClient("llama3", tc.in, 0)
if client.baseURL != tc.expected {
t.Fatalf("baseURL = %q, want %q", client.baseURL, tc.expected)
}
@@ -408,7 +408,7 @@ func TestOllamaClient_Chat_ToolCallsResponse(t *testing.T) {
}))
defer server.Close()
client := NewOllamaClient("llama3", server.URL)
client := NewOllamaClient("llama3", server.URL, 0)
out, err := client.Chat(context.Background(), ChatRequest{
Messages: []Message{{Role: "user", Content: "What time is it?"}},
})
@@ -442,7 +442,7 @@ func TestOllamaClient_Chat_ToolCallsAndToolResultsInRequest(t *testing.T) {
}))
defer server.Close()
client := NewOllamaClient("llama3", server.URL)
client := NewOllamaClient("llama3", server.URL, 0)
_, err := client.Chat(context.Background(), ChatRequest{
System: "system prompt",
Messages: []Message{
internal/ai/providers/openai.go
+6 -3
View File
@@ -30,7 +30,8 @@ type OpenAIClient struct {
}
// NewOpenAIClient creates a new OpenAI API client
func NewOpenAIClient(apiKey, model, baseURL string) *OpenAIClient {
// timeout is optional - pass 0 to use the default 5 minute timeout
func NewOpenAIClient(apiKey, model, baseURL string, timeout time.Duration) *OpenAIClient {
if baseURL == "" {
baseURL = openaiAPIURL
}
@@ -40,13 +41,15 @@ func NewOpenAIClient(apiKey, model, baseURL string) *OpenAIClient {
} else if strings.HasPrefix(model, "deepseek:") {
model = strings.TrimPrefix(model, "deepseek:")
}
if timeout <= 0 {
timeout = 300 * time.Second // Default 5 minutes
}
return &OpenAIClient{
apiKey: apiKey,
model: model,
baseURL: baseURL,
client: &http.Client{
// 5 minutes timeout - DeepSeek reasoning models can take a long time
Timeout: 300 * time.Second,
Timeout: timeout,
},
}
}
internal/ai/providers/openai_test.go
+11 -11
View File
@@ -115,7 +115,7 @@ func TestOpenAIClient_Chat_Success(t *testing.T) {
defer server.Close()
// Create client with mock server URL
client := NewOpenAIClient("test-api-key", "gpt-4o", server.URL+"/v1/chat/completions")
client := NewOpenAIClient("test-api-key", "gpt-4o", server.URL+"/v1/chat/completions", 0)
// Execute chat request
ctx := context.Background()
@@ -158,7 +158,7 @@ func TestOpenAIClient_Chat_APIError(t *testing.T) {
}))
defer server.Close()
client := NewOpenAIClient("invalid-key", "gpt-4o", server.URL+"/v1/chat/completions")
client := NewOpenAIClient("invalid-key", "gpt-4o", server.URL+"/v1/chat/completions", 0)
ctx := context.Background()
_, err := client.Chat(ctx, ChatRequest{
@@ -172,7 +172,7 @@ func TestOpenAIClient_Chat_APIError(t *testing.T) {
func TestOpenAIClient_Chat_NetworkError(t *testing.T) {
// Create client pointing to non-existent server
client := NewOpenAIClient("test-key", "gpt-4o", "http://localhost:99999")
client := NewOpenAIClient("test-key", "gpt-4o", "http://localhost:99999", 0)
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
@@ -193,7 +193,7 @@ func TestOpenAIClient_Chat_ContextCanceled(t *testing.T) {
}))
defer server.Close()
client := NewOpenAIClient("test-key", "gpt-4o", server.URL+"/v1/chat/completions")
client := NewOpenAIClient("test-key", "gpt-4o", server.URL+"/v1/chat/completions", 0)
ctx, cancel := context.WithCancel(context.Background())
cancel() // Cancel immediately
@@ -236,7 +236,7 @@ func TestOpenAIClient_Chat_WithSystemPrompt(t *testing.T) {
}))
defer server.Close()
client := NewOpenAIClient("test-key", "gpt-4o", server.URL+"/v1/chat/completions")
client := NewOpenAIClient("test-key", "gpt-4o", server.URL+"/v1/chat/completions", 0)
ctx := context.Background()
_, err := client.Chat(ctx, ChatRequest{
@@ -284,7 +284,7 @@ func TestOpenAIClient_ListModels_UsesConfiguredHostAndFilters(t *testing.T) {
}))
defer server.Close()
client := NewOpenAIClient("test-api-key", "gpt-4o", server.URL+"/v1/chat/completions")
client := NewOpenAIClient("test-api-key", "gpt-4o", server.URL+"/v1/chat/completions", 0)
models, err := client.ListModels(context.Background())
if err != nil {
t.Fatalf("ListModels: %v", err)
@@ -312,7 +312,7 @@ func TestOpenAIClient_TestConnection_CallsListModels(t *testing.T) {
}))
defer server.Close()
client := NewOpenAIClient("test-api-key", "gpt-4o", server.URL+"/v1/chat/completions")
client := NewOpenAIClient("test-api-key", "gpt-4o", server.URL+"/v1/chat/completions", 0)
if err := client.TestConnection(context.Background()); err != nil {
t.Fatalf("TestConnection: %v", err)
}
@@ -345,7 +345,7 @@ func TestOpenAIClient_Chat_UsesMaxCompletionTokensForOpenAI(t *testing.T) {
}))
defer server.Close()
client := NewOpenAIClient("test-api-key", "gpt-4o", server.URL+"/v1/chat/completions")
client := NewOpenAIClient("test-api-key", "gpt-4o", server.URL+"/v1/chat/completions", 0)
_, err := client.Chat(context.Background(), ChatRequest{
Messages: []Message{{Role: "user", Content: "Hello"}},
MaxTokens: 123,
@@ -383,7 +383,7 @@ func TestOpenAIClient_Chat_GPT52NonChat_UsesCompletionsEndpointAndPrompt(t *test
t.Fatalf("parse server url: %v", err)
}
client := NewOpenAIClient("test-api-key", "gpt-5.2-pro", "https://api.openai.com/v1/chat/completions")
client := NewOpenAIClient("test-api-key", "gpt-5.2-pro", "https://api.openai.com/v1/chat/completions", 0)
client.client.Transport = rewriteToServerTransport{serverBase: serverURL, rt: http.DefaultTransport}
_, err = client.Chat(context.Background(), ChatRequest{
@@ -426,7 +426,7 @@ func TestOpenAIClient_ListModels_DeepSeekUsesModelsEndpoint(t *testing.T) {
t.Fatalf("parse server url: %v", err)
}
client := NewOpenAIClient("test-api-key", "deepseek-chat", "https://api.deepseek.com/v1/chat/completions")
client := NewOpenAIClient("test-api-key", "deepseek-chat", "https://api.deepseek.com/v1/chat/completions", 0)
client.client.Transport = rewriteToServerTransport{serverBase: serverURL, rt: http.DefaultTransport}
_, err = client.ListModels(context.Background())
@@ -462,7 +462,7 @@ func TestOpenAIClient_Chat_O1OmitsTemperature(t *testing.T) {
}))
defer server.Close()
client := NewOpenAIClient("test-api-key", "o1-mini", server.URL+"/v1/chat/completions")
client := NewOpenAIClient("test-api-key", "o1-mini", server.URL+"/v1/chat/completions", 0)
_, err := client.Chat(context.Background(), ChatRequest{
Messages: []Message{{Role: "user", Content: "Hello"}},
Model: "o1-mini",
internal/config/ai.go
+13
View File
@@ -53,6 +53,10 @@ type AIConfig struct {
// Alert-triggered AI analysis - analyze specific resources when alerts fire
AlertTriggeredAnalysis bool `json:"alert_triggered_analysis"` // Enable AI analysis when alerts fire (token-efficient)
// Request timeout - how long to wait for AI responses (default: 300s / 5 min)
// Increase this for slow hardware running local models (e.g., Ollama on low-power devices)
RequestTimeoutSeconds int `json:"request_timeout_seconds,omitempty"`
// AI cost controls
// Budget is expressed as an estimated USD amount over a 30-day window (pro-rated in UI for other ranges).
CostBudgetUSD30d float64 `json:"cost_budget_usd_30d,omitempty"`
@@ -454,3 +458,12 @@ func (c *AIConfig) IsPatrolEnabled() bool {
func (c *AIConfig) IsAlertTriggeredAnalysisEnabled() bool {
return c.AlertTriggeredAnalysis
}
// GetRequestTimeout returns the timeout duration for AI requests
// Default is 5 minutes (300 seconds) if not configured
func (c *AIConfig) GetRequestTimeout() time.Duration {
if c.RequestTimeoutSeconds > 0 {
return time.Duration(c.RequestTimeoutSeconds) * time.Second
}
return 300 * time.Second // 5 minutes default
}