fix(tools): harden Bocha search provider per review feedback
- Cap response body read at 1MB via io.LimitReader to prevent OOM - Reuse HTTP client across searches instead of creating per request - Add mock HTTP server tests for success and error paths Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
efad854d9c
commit
295c0186f6
2 changed files with 118 additions and 8 deletions
|
|
@ -394,7 +394,7 @@ func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, cou
|
||||||
type BochaSearchProvider struct {
|
type BochaSearchProvider struct {
|
||||||
apiKey string
|
apiKey string
|
||||||
baseURL string
|
baseURL string
|
||||||
proxy string
|
client *http.Client
|
||||||
}
|
}
|
||||||
|
|
||||||
func (p *BochaSearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
|
func (p *BochaSearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
|
||||||
|
|
@ -422,17 +422,13 @@ func (p *BochaSearchProvider) Search(ctx context.Context, query string, count in
|
||||||
req.Header.Set("Content-Type", "application/json")
|
req.Header.Set("Content-Type", "application/json")
|
||||||
req.Header.Set("Authorization", "Bearer "+p.apiKey)
|
req.Header.Set("Authorization", "Bearer "+p.apiKey)
|
||||||
|
|
||||||
client, err := createHTTPClient(p.proxy, 15*time.Second)
|
resp, err := p.client.Do(req)
|
||||||
if err != nil {
|
|
||||||
return "", fmt.Errorf("failed to create HTTP client: %w", err)
|
|
||||||
}
|
|
||||||
resp, err := client.Do(req)
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return "", fmt.Errorf("request failed: %w", err)
|
return "", fmt.Errorf("request failed: %w", err)
|
||||||
}
|
}
|
||||||
defer resp.Body.Close()
|
defer resp.Body.Close()
|
||||||
|
|
||||||
body, err := io.ReadAll(resp.Body)
|
body, err := io.ReadAll(io.LimitReader(resp.Body, 1*1024*1024))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return "", fmt.Errorf("failed to read response: %w", err)
|
return "", fmt.Errorf("failed to read response: %w", err)
|
||||||
}
|
}
|
||||||
|
|
@ -566,10 +562,14 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
|
||||||
maxResults = opts.DuckDuckGoMaxResults
|
maxResults = opts.DuckDuckGoMaxResults
|
||||||
}
|
}
|
||||||
} else if opts.BochaEnabled && opts.BochaAPIKey != "" {
|
} else if opts.BochaEnabled && opts.BochaAPIKey != "" {
|
||||||
|
bochaClient, err := createHTTPClient(opts.Proxy, 15*time.Second)
|
||||||
|
if err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
provider = &BochaSearchProvider{
|
provider = &BochaSearchProvider{
|
||||||
apiKey: opts.BochaAPIKey,
|
apiKey: opts.BochaAPIKey,
|
||||||
baseURL: opts.BochaBaseURL,
|
baseURL: opts.BochaBaseURL,
|
||||||
proxy: opts.Proxy,
|
client: bochaClient,
|
||||||
}
|
}
|
||||||
if opts.BochaMaxResults > 0 {
|
if opts.BochaMaxResults > 0 {
|
||||||
maxResults = opts.BochaMaxResults
|
maxResults = opts.BochaMaxResults
|
||||||
|
|
|
||||||
|
|
@ -681,3 +681,113 @@ 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_BochaSearch_Success verifies successful Bocha search
|
||||||
|
func TestWebTool_BochaSearch_Success(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 r.Header.Get("Content-Type") != "application/json" {
|
||||||
|
t.Errorf("Expected Content-Type application/json, got %s", r.Header.Get("Content-Type"))
|
||||||
|
}
|
||||||
|
if r.Header.Get("Authorization") != "Bearer test-bocha-key" {
|
||||||
|
t.Errorf("Expected Authorization Bearer test-bocha-key, got %s", r.Header.Get("Authorization"))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify payload
|
||||||
|
var payload map[string]any
|
||||||
|
json.NewDecoder(r.Body).Decode(&payload)
|
||||||
|
if payload["query"] != "test query" {
|
||||||
|
t.Errorf("Expected query 'test query', got %v", payload["query"])
|
||||||
|
}
|
||||||
|
|
||||||
|
// Return mock Bocha response
|
||||||
|
response := map[string]any{
|
||||||
|
"code": 200,
|
||||||
|
"data": map[string]any{
|
||||||
|
"webPages": map[string]any{
|
||||||
|
"value": []map[string]any{
|
||||||
|
{
|
||||||
|
"name": "Bocha Result 1",
|
||||||
|
"url": "https://example.com/bocha/1",
|
||||||
|
"snippet": "Snippet for result 1",
|
||||||
|
"summary": "Summary for result 1",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"name": "Bocha Result 2",
|
||||||
|
"url": "https://example.com/bocha/2",
|
||||||
|
"snippet": "Snippet for result 2",
|
||||||
|
"summary": "",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"msg": "success",
|
||||||
|
}
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
json.NewEncoder(w).Encode(response)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
tool := NewWebSearchTool(WebSearchToolOptions{
|
||||||
|
BochaEnabled: true,
|
||||||
|
BochaAPIKey: "test-bocha-key",
|
||||||
|
BochaBaseURL: server.URL,
|
||||||
|
BochaMaxResults: 5,
|
||||||
|
})
|
||||||
|
|
||||||
|
ctx := context.Background()
|
||||||
|
result := tool.Execute(ctx, map[string]any{"query": "test query"})
|
||||||
|
|
||||||
|
if result.IsError {
|
||||||
|
t.Fatalf("Expected success, got IsError=true: %s", result.ForLLM)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Should contain result titles and URLs
|
||||||
|
if !strings.Contains(result.ForUser, "Bocha Result 1") ||
|
||||||
|
!strings.Contains(result.ForUser, "https://example.com/bocha/1") {
|
||||||
|
t.Errorf("Expected results in output, got: %s", result.ForUser)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Should prefer summary over snippet
|
||||||
|
if !strings.Contains(result.ForUser, "Summary for result 1") {
|
||||||
|
t.Errorf("Expected summary preferred over snippet, got: %s", result.ForUser)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Should fall back to snippet when summary is empty
|
||||||
|
if !strings.Contains(result.ForUser, "Snippet for result 2") {
|
||||||
|
t.Errorf("Expected snippet fallback when summary empty, got: %s", result.ForUser)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Should mention via Bocha
|
||||||
|
if !strings.Contains(result.ForUser, "via Bocha") {
|
||||||
|
t.Errorf("Expected 'via Bocha' in output, got: %s", result.ForUser)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestWebTool_BochaSearch_APIError verifies Bocha error handling
|
||||||
|
func TestWebTool_BochaSearch_APIError(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusUnauthorized)
|
||||||
|
w.Write([]byte("unauthorized"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
tool := NewWebSearchTool(WebSearchToolOptions{
|
||||||
|
BochaEnabled: true,
|
||||||
|
BochaAPIKey: "bad-key",
|
||||||
|
BochaBaseURL: server.URL,
|
||||||
|
})
|
||||||
|
|
||||||
|
ctx := context.Background()
|
||||||
|
result := tool.Execute(ctx, map[string]any{"query": "test"})
|
||||||
|
|
||||||
|
if !result.IsError {
|
||||||
|
t.Errorf("Expected error for unauthorized response")
|
||||||
|
}
|
||||||
|
if !strings.Contains(result.ForLLM, "bocha API error") {
|
||||||
|
t.Errorf("Expected bocha API error message, got: %s", result.ForLLM)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue