From 9e186f7b3d4eaa9667c3cc15355bc20eac3853b3 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 23 Feb 2026 22:22:51 +0800 Subject: [PATCH] feat: set heartbeat default to 5m and improve grok web_search parsing --- config/config.example.json | 2 +- pkg/config/defaults.go | 2 +- pkg/heartbeat/service.go | 2 +- pkg/tools/web.go | 78 +++++++++++++++++++++++++++++++--- pkg/tools/web_test.go | 87 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 161 insertions(+), 10 deletions(-) diff --git a/config/config.example.json b/config/config.example.json index 6db124bc7..52e993a97 100644 --- a/config/config.example.json +++ b/config/config.example.json @@ -243,7 +243,7 @@ }, "heartbeat": { "enabled": true, - "interval": 30 + "interval": 5 }, "devices": { "enabled": false, diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index c6e79d968..a1db7ab3e 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -317,7 +317,7 @@ func DefaultConfig() *Config { }, Heartbeat: HeartbeatConfig{ Enabled: true, - Interval: 30, + Interval: 5, }, Devices: DevicesConfig{ Enabled: false, diff --git a/pkg/heartbeat/service.go b/pkg/heartbeat/service.go index e05a9fdbf..3111d5145 100644 --- a/pkg/heartbeat/service.go +++ b/pkg/heartbeat/service.go @@ -23,7 +23,7 @@ import ( const ( minIntervalMinutes = 5 - defaultIntervalMinutes = 30 + defaultIntervalMinutes = 5 ) // HeartbeatHandler is the function type for handling heartbeat. diff --git a/pkg/tools/web.go b/pkg/tools/web.go index 7ad66dbb2..b7abbe44b 100644 --- a/pkg/tools/web.go +++ b/pkg/tools/web.go @@ -349,6 +349,7 @@ func (p *GrokSearchProvider) Search(ctx context.Context, query string, count int }, }, "max_tokens": 1000, + "stream": false, } payloadBytes, err := json.Marshal(payload) @@ -356,12 +357,13 @@ func (p *GrokSearchProvider) Search(ctx context.Context, query string, count int return "", fmt.Errorf("failed to marshal request: %w", err) } - req, err := http.NewRequestWithContext(ctx, "POST", searchURL, strings.NewReader(string(payloadBytes))) + req, err := http.NewRequestWithContext(ctx, "POST", searchURL, bytes.NewReader(payloadBytes)) if err != nil { return "", fmt.Errorf("failed to create request: %w", err) } req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") req.Header.Set("Authorization", "Bearer "+p.apiKey) req.Header.Set("User-Agent", userAgent) @@ -384,6 +386,19 @@ func (p *GrokSearchProvider) Search(ctx context.Context, query string, count int return "", fmt.Errorf("Grok API error: %s", string(body)) } + content, err := parseGrokResponseContent(body) + if err != nil { + return "", fmt.Errorf("failed to parse response: %w", err) + } + if strings.TrimSpace(content) == "" { + return fmt.Sprintf("No results for: %s", query), nil + } + + return fmt.Sprintf("Results for: %s (via Grok)\n%s", query, content), nil +} + +func parseGrokResponseContent(body []byte) (string, error) { + // First try regular JSON completion response. var searchResp struct { Choices []struct { Message struct { @@ -391,16 +406,65 @@ func (p *GrokSearchProvider) Search(ctx context.Context, query string, count int } `json:"message"` } `json:"choices"` } - - if err := json.Unmarshal(body, &searchResp); err != nil { - return "", fmt.Errorf("failed to parse response: %w", err) + if err := json.Unmarshal(body, &searchResp); err == nil { + if len(searchResp.Choices) > 0 { + return strings.TrimSpace(searchResp.Choices[0].Message.Content), nil + } } - if len(searchResp.Choices) == 0 { - return fmt.Sprintf("No results for: %s", query), nil + // Some OpenAI-compatible gateways return SSE chunks even when stream=false. + // Parse lines in "data: {...}" format and stitch delta content. + text := strings.TrimSpace(string(body)) + if !strings.Contains(text, "data:") { + return "", fmt.Errorf("unexpected response format") } - return fmt.Sprintf("Results for: %s (via Grok)\n%s", query, searchResp.Choices[0].Message.Content), nil + var merged strings.Builder + lines := strings.Split(text, "\n") + for _, line := range lines { + line = strings.TrimSpace(line) + if !strings.HasPrefix(line, "data:") { + continue + } + chunk := strings.TrimSpace(strings.TrimPrefix(line, "data:")) + if chunk == "" || chunk == "[DONE]" { + continue + } + + var sseChunk struct { + Choices []struct { + Delta struct { + Content string `json:"content"` + } `json:"delta"` + Message struct { + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` + } + if err := json.Unmarshal([]byte(chunk), &sseChunk); err != nil { + continue + } + if len(sseChunk.Choices) == 0 { + continue + } + + part := strings.TrimSpace(sseChunk.Choices[0].Delta.Content) + if part == "" { + part = strings.TrimSpace(sseChunk.Choices[0].Message.Content) + } + if part == "" { + continue + } + if merged.Len() > 0 { + merged.WriteByte(' ') + } + merged.WriteString(part) + } + + if merged.Len() == 0 { + return "", fmt.Errorf("empty content in SSE response") + } + return merged.String(), nil } type WebSearchTool struct { diff --git a/pkg/tools/web_test.go b/pkg/tools/web_test.go index 2cd79eb24..602131e0a 100644 --- a/pkg/tools/web_test.go +++ b/pkg/tools/web_test.go @@ -572,3 +572,90 @@ func TestWebTool_TavilySearch_Success(t *testing.T) { t.Errorf("Expected 'via Tavily' in output, got: %s", result.ForUser) } } + +// TestWebTool_GrokSearch_JSONResponse verifies Grok JSON response parsing. +func TestWebTool_GrokSearch_JSONResponse(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "POST" { + t.Errorf("Expected POST request, got %s", r.Method) + } + if got := r.Header.Get("Authorization"); got != "Bearer test-key" { + t.Errorf("Expected Authorization Bearer test-key, got %s", got) + } + + var payload map[string]any + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatalf("Failed to decode payload: %v", err) + } + if payload["model"] != "grok-4.20-beta" { + t.Errorf("Expected model grok-4.20-beta, got %v", payload["model"]) + } + if stream, ok := payload["stream"].(bool); !ok || stream { + t.Errorf("Expected stream=false, got %v", payload["stream"]) + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"choices":[{"message":{"content":"1. Test\n https://example.com\n Example snippet"}}]}`)) + })) + defer server.Close() + + tool := NewWebSearchTool(WebSearchToolOptions{ + GrokEnabled: true, + GrokAPIKey: "test-key", + GrokEndpoint: server.URL, + GrokModel: "grok-4.20-beta", + GrokMaxResults: 3, + }) + + result := tool.Execute(context.Background(), map[string]any{ + "query": "test query", + "count": 3.0, + }) + + if result.IsError { + t.Fatalf("Expected success, got error: %s", result.ForLLM) + } + if !strings.Contains(result.ForUser, "via Grok") { + t.Errorf("Expected 'via Grok' in output, got: %s", result.ForUser) + } + if !strings.Contains(result.ForUser, "https://example.com") { + t.Errorf("Expected result URL in output, got: %s", result.ForUser) + } +} + +// TestWebTool_GrokSearch_SSEResponse verifies SSE chunk response parsing. +func TestWebTool_GrokSearch_SSEResponse(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + w.Write([]byte( + "data: {\"id\":\"1\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"\"}}]}\n\n" + + "data: {\"id\":\"1\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"1. SSE Result\"}}]}\n\n" + + "data: {\"id\":\"1\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\" https://example.org\"}}]}\n\n" + + "data: [DONE]\n\n", + )) + })) + defer server.Close() + + tool := NewWebSearchTool(WebSearchToolOptions{ + GrokEnabled: true, + GrokAPIKey: "test-key", + GrokEndpoint: server.URL, + GrokModel: "grok-4.20-beta", + }) + + result := tool.Execute(context.Background(), map[string]any{ + "query": "sse query", + }) + + if result.IsError { + t.Fatalf("Expected success, got error: %s", result.ForLLM) + } + if !strings.Contains(result.ForUser, "SSE Result") { + t.Errorf("Expected SSE content in output, got: %s", result.ForUser) + } + if !strings.Contains(result.ForUser, "https://example.org") { + t.Errorf("Expected SSE URL in output, got: %s", result.ForUser) + } +}