feat: set heartbeat default to 5m and improve grok web_search parsing
This commit is contained in:
parent
8c2058fc5c
commit
9e186f7b3d
5 changed files with 161 additions and 10 deletions
|
|
@ -243,7 +243,7 @@
|
||||||
},
|
},
|
||||||
"heartbeat": {
|
"heartbeat": {
|
||||||
"enabled": true,
|
"enabled": true,
|
||||||
"interval": 30
|
"interval": 5
|
||||||
},
|
},
|
||||||
"devices": {
|
"devices": {
|
||||||
"enabled": false,
|
"enabled": false,
|
||||||
|
|
|
||||||
|
|
@ -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,
|
||||||
|
|
|
||||||
|
|
@ -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.
|
||||||
|
|
|
||||||
|
|
@ -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 {
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue