feat: set heartbeat default to 5m and improve grok web_search parsing

This commit is contained in:
root 2026-02-23 22:22:51 +08:00
parent 8c2058fc5c
commit 9e186f7b3d
5 changed files with 161 additions and 10 deletions

View file

@ -243,7 +243,7 @@
}, },
"heartbeat": { "heartbeat": {
"enabled": true, "enabled": true,
"interval": 30 "interval": 5
}, },
"devices": { "devices": {
"enabled": false, "enabled": false,

View file

@ -317,7 +317,7 @@ func DefaultConfig() *Config {
}, },
Heartbeat: HeartbeatConfig{ Heartbeat: HeartbeatConfig{
Enabled: true, Enabled: true,
Interval: 30, Interval: 5,
}, },
Devices: DevicesConfig{ Devices: DevicesConfig{
Enabled: false, Enabled: false,

View file

@ -23,7 +23,7 @@ import (
const ( const (
minIntervalMinutes = 5 minIntervalMinutes = 5
defaultIntervalMinutes = 30 defaultIntervalMinutes = 5
) )
// HeartbeatHandler is the function type for handling heartbeat. // HeartbeatHandler is the function type for handling heartbeat.

View file

@ -349,6 +349,7 @@ func (p *GrokSearchProvider) Search(ctx context.Context, query string, count int
}, },
}, },
"max_tokens": 1000, "max_tokens": 1000,
"stream": false,
} }
payloadBytes, err := json.Marshal(payload) 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) 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 { if err != nil {
return "", fmt.Errorf("failed to create request: %w", err) return "", fmt.Errorf("failed to create request: %w", err)
} }
req.Header.Set("Content-Type", "application/json") req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json")
req.Header.Set("Authorization", "Bearer "+p.apiKey) req.Header.Set("Authorization", "Bearer "+p.apiKey)
req.Header.Set("User-Agent", userAgent) 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)) 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 { var searchResp struct {
Choices []struct { Choices []struct {
Message struct { Message struct {
@ -391,16 +406,65 @@ func (p *GrokSearchProvider) Search(ctx context.Context, query string, count int
} `json:"message"` } `json:"message"`
} `json:"choices"` } `json:"choices"`
} }
if err := json.Unmarshal(body, &searchResp); err == nil {
if err := json.Unmarshal(body, &searchResp); err != nil { if len(searchResp.Choices) > 0 {
return "", fmt.Errorf("failed to parse response: %w", err) return strings.TrimSpace(searchResp.Choices[0].Message.Content), nil
}
} }
if len(searchResp.Choices) == 0 { // Some OpenAI-compatible gateways return SSE chunks even when stream=false.
return fmt.Sprintf("No results for: %s", query), nil // 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 { type WebSearchTool struct {

View file

@ -572,3 +572,90 @@ func TestWebTool_TavilySearch_Success(t *testing.T) {
t.Errorf("Expected 'via Tavily' in output, got: %s", result.ForUser) 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)
}
}