Add Search API configuration and enhance documentation
- Introduced new environment variables for TAVILY_API_KEY, SERPAPI_API_KEY, and SERPER_API_KEY in both `pr-test.yml` and `unit-test.yml` workflows to support additional search providers. - Updated `DESIGN.md` to reflect the inclusion of SerpAPI as a search provider, detailing its usage and configuration options, including support for multiple search engines. - Enhanced the `builtinSearch` function in `handler.go` to accommodate the new SerpAPI provider, ensuring proper handling of search requests. - Revised the `WebConfig` struct in `config.go` to include an `Engine` field for specifying the search engine when using SerpAPI, improving flexibility in search configurations. - Updated documentation to clarify the roles of new search providers and their integration within the search module, ensuring comprehensive guidance for developers.
This commit is contained in:
Max
parent
f8ba875cd3
commit
e4d1701dab
11 changed files with 1859 additions and 15 deletions
5
.github/workflows/pr-test.yml
vendored
5
.github/workflows/pr-test.yml
vendored
|
|
@ -49,6 +49,11 @@ env:
|
|||
DEEPSEEK_MODELS_V3: ${{ secrets.DEEPSEEK_MODELS_V3 }}
|
||||
DEEPSEEK_MODELS_V3_1: ${{ secrets.DEEPSEEK_MODELS_V3_1 }}
|
||||
|
||||
# Search API Configuration
|
||||
TAVILY_API_KEY: ${{ secrets.TAVILY_API_KEY }}
|
||||
SERPAPI_API_KEY: ${{ secrets.SERPAPI_API_KEY }}
|
||||
SERPER_API_KEY: ${{ secrets.SERPER_API_KEY }}
|
||||
|
||||
# Claude API Configuration
|
||||
CLAUDE_API_KEY: ${{ secrets.CLAUDE_API_KEY }}
|
||||
CLAUDE_PROXY: ${{ secrets.CLAUDE_PROXY }}
|
||||
|
|
|
|||
5
.github/workflows/unit-test.yml
vendored
5
.github/workflows/unit-test.yml
vendored
|
|
@ -53,6 +53,11 @@ env:
|
|||
DEEPSEEK_MODELS_V3: ${{ secrets.DEEPSEEK_MODELS_V3 }}
|
||||
DEEPSEEK_MODELS_V3_1: ${{ secrets.DEEPSEEK_MODELS_V3_1 }}
|
||||
|
||||
# Search API Configuration
|
||||
TAVILY_API_KEY: ${{ secrets.TAVILY_API_KEY }}
|
||||
SERPAPI_API_KEY: ${{ secrets.SERPAPI_API_KEY }}
|
||||
SERPER_API_KEY: ${{ secrets.SERPER_API_KEY }}
|
||||
|
||||
# Claude API Configuration
|
||||
CLAUDE_API_KEY: ${{ secrets.CLAUDE_API_KEY }}
|
||||
CLAUDE_PROXY: ${{ secrets.CLAUDE_PROXY }}
|
||||
|
|
|
|||
|
|
@ -160,7 +160,8 @@ agent/search/
|
|||
│ ├── web/ # Web search
|
||||
│ │ ├── handler.go # Web search handler (mode dispatch)
|
||||
│ │ ├── tavily.go # Tavily provider (builtin)
|
||||
│ │ ├── serper.go # Serper provider (builtin)
|
||||
│ │ ├── serper.go # Serper provider (serper.dev, builtin)
|
||||
│ │ ├── serpapi.go # SerpAPI provider (serpapi.com, multi-engine, builtin)
|
||||
│ │ ├── agent.go # Agent mode (AI Search)
|
||||
│ │ └── mcp.go # MCP mode (external service)
|
||||
│ │
|
||||
|
|
@ -663,9 +664,10 @@ type Config struct {
|
|||
// Note: uses.web determines the mode (builtin/agent/mcp)
|
||||
// Provider is only used when uses.web = "builtin"
|
||||
type WebConfig struct {
|
||||
Provider string `json:"provider,omitempty"` // "tavily" or "serper" (for builtin mode)
|
||||
Provider string `json:"provider,omitempty"` // "tavily", "serper", or "serpapi" (for builtin mode)
|
||||
APIKeyEnv string `json:"api_key_env,omitempty"` // Environment variable for API key
|
||||
MaxResults int `json:"max_results,omitempty"` // Max results (default: 10)
|
||||
Engine string `json:"engine,omitempty"` // Search engine for SerpAPI: "google", "bing", "baidu", etc. (default: "google")
|
||||
}
|
||||
|
||||
// KBConfig for knowledge base search settings
|
||||
|
|
@ -1299,9 +1301,10 @@ func (ast *Assistant) GetMergedSearchConfig() *searchTypes.Config {
|
|||
|
||||
# Web search settings
|
||||
web:
|
||||
provider: "tavily" # "tavily", "serper" (builtin providers only)
|
||||
provider: "tavily" # "tavily", "serper", or "serpapi" (builtin providers only)
|
||||
api_key_env: "TAVILY_API_KEY"
|
||||
max_results: 10
|
||||
# engine: "google" # For SerpAPI only: "google", "bing", "baidu", "yandex", etc.
|
||||
|
||||
# Knowledge base search settings
|
||||
kb:
|
||||
|
|
@ -1743,16 +1746,21 @@ func (h *Handler) Search(ctx *context.Context, req *types.Request) (*types.Resul
|
|||
}
|
||||
}
|
||||
|
||||
// builtinSearch uses Tavily/Serper directly
|
||||
// builtinSearch uses Tavily/Serper/SerpAPI directly
|
||||
func (h *Handler) builtinSearch(ctx *context.Context, req *types.Request) (*types.Result, error) {
|
||||
var provider Provider
|
||||
switch h.config.Provider {
|
||||
case "tavily":
|
||||
provider = NewTavilyProvider(h.config)
|
||||
return NewTavilyProvider(h.config).Search(req)
|
||||
case "serper":
|
||||
provider = NewSerperProvider(h.config)
|
||||
// Serper (serper.dev) - POST request with X-API-KEY header
|
||||
return NewSerperProvider(h.config).Search(req)
|
||||
case "serpapi":
|
||||
// SerpAPI (serpapi.com) - GET request with api_key parameter
|
||||
// Supports multiple engines: google, bing, baidu, yandex, etc.
|
||||
return NewSerpAPIProvider(h.config).Search(req)
|
||||
default:
|
||||
return nil, fmt.Errorf("unknown provider: %s", h.config.Provider)
|
||||
}
|
||||
return provider.Search(ctx, req)
|
||||
}
|
||||
|
||||
// agentSearch delegates to an assistant for AI-powered search
|
||||
|
|
@ -1780,10 +1788,40 @@ func (h *Handler) mcpSearch(ctx *context.Context, req *types.Request) (*types.Re
|
|||
|
||||
**Built-in Providers (when `uses.web = "builtin"`):**
|
||||
|
||||
| Provider | File | Notes |
|
||||
| -------- | ----------- | ------------------------------- |
|
||||
| Tavily | `tavily.go` | Recommended for AI applications |
|
||||
| Serper | `serper.go` | Google search API |
|
||||
| Provider | File | Notes |
|
||||
| -------- | ------------ | ----------------------------------------------- |
|
||||
| Tavily | `tavily.go` | Recommended for AI applications |
|
||||
| Serper | `serper.go` | Google search via serper.dev (POST + X-API-KEY) |
|
||||
| SerpAPI | `serpapi.go` | Multi-engine search via serpapi.com (GET + URL) |
|
||||
|
||||
**SerpAPI Engine Support:**
|
||||
|
||||
SerpAPI supports multiple search engines via the `engine` config:
|
||||
|
||||
| Engine | Description |
|
||||
| ------------ | ---------------------------- |
|
||||
| `google` | Google Search (default) |
|
||||
| `bing` | Bing Search |
|
||||
| `baidu` | Baidu (百度) |
|
||||
| `yandex` | Yandex Search |
|
||||
| `yahoo` | Yahoo Search |
|
||||
| `duckduckgo` | DuckDuckGo Search |
|
||||
| `naver` | Naver Search (Korean) |
|
||||
| `ecosia` | Ecosia Search (eco-friendly) |
|
||||
| `seznam` | Seznam Search (Czech) |
|
||||
|
||||
See [SerpAPI Documentation](https://serpapi.com/search-api) for the full list of supported engines.
|
||||
|
||||
Configuration example:
|
||||
|
||||
```yaml
|
||||
# agent/search.yml
|
||||
web:
|
||||
provider: "serpapi"
|
||||
api_key_env: "SERPAPI_API_KEY"
|
||||
engine: "bing" # Use Bing instead of Google
|
||||
max_results: 10
|
||||
```
|
||||
|
||||
**Agent Mode (AI Search):**
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
package web
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/yaoapp/yao/agent/search/types"
|
||||
)
|
||||
|
||||
|
|
@ -21,14 +24,88 @@ func (h *Handler) Type() types.SearchType {
|
|||
}
|
||||
|
||||
// Search executes web search based on uses.web mode
|
||||
// TODO: Implement actual search logic
|
||||
func (h *Handler) Search(req *types.Request) (*types.Result, error) {
|
||||
// Skeleton implementation - returns empty result
|
||||
switch {
|
||||
case h.usesWeb == "builtin" || h.usesWeb == "":
|
||||
return h.builtinSearch(req)
|
||||
case strings.HasPrefix(h.usesWeb, "mcp:"):
|
||||
return h.mcpSearch(req)
|
||||
default:
|
||||
// Agent mode: delegate to assistant for AI-powered search
|
||||
return h.agentSearch(req)
|
||||
}
|
||||
}
|
||||
|
||||
// builtinSearch uses Tavily/Serper/SerpAPI directly
|
||||
func (h *Handler) builtinSearch(req *types.Request) (*types.Result, error) {
|
||||
// Determine provider from config
|
||||
providerName := "tavily" // default
|
||||
if h.config != nil && h.config.Provider != "" {
|
||||
providerName = h.config.Provider
|
||||
}
|
||||
|
||||
switch providerName {
|
||||
case "tavily":
|
||||
return NewTavilyProvider(h.config).Search(req)
|
||||
case "serper":
|
||||
// Serper (serper.dev) - POST request with X-API-KEY header
|
||||
return NewSerperProvider(h.config).Search(req)
|
||||
case "serpapi":
|
||||
// SerpAPI (serpapi.com) - GET request with api_key parameter
|
||||
return NewSerpAPIProvider(h.config).Search(req)
|
||||
default:
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeWeb,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Error: fmt.Sprintf("Unknown provider: %s (supported: tavily, serper, serpapi)", providerName),
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
|
||||
// agentSearch delegates to an assistant for AI-powered search
|
||||
func (h *Handler) agentSearch(req *types.Request) (*types.Result, error) {
|
||||
// TODO: Implement agent mode
|
||||
// 1. Call assistant with search request
|
||||
// 2. Assistant understands intent, generates optimized queries
|
||||
// 3. Assistant executes searches (may call builtin internally)
|
||||
// 4. Assistant analyzes and returns structured results
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeWeb,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Error: "Agent mode not yet implemented",
|
||||
}, nil
|
||||
}
|
||||
|
||||
// mcpSearch calls external MCP tool
|
||||
func (h *Handler) mcpSearch(req *types.Request) (*types.Result, error) {
|
||||
// TODO: Implement MCP mode
|
||||
// Parse "mcp:server.tool"
|
||||
mcpRef := strings.TrimPrefix(h.usesWeb, "mcp:")
|
||||
parts := strings.SplitN(mcpRef, ".", 2)
|
||||
if len(parts) != 2 {
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeWeb,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Error: fmt.Sprintf("Invalid MCP format, expected 'mcp:server.tool', got '%s'", h.usesWeb),
|
||||
}, nil
|
||||
}
|
||||
// serverID, toolName := parts[0], parts[1]
|
||||
// Call MCP tool
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeWeb,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Error: "MCP mode not yet implemented",
|
||||
}, nil
|
||||
}
|
||||
|
|
|
|||
302
agent/search/handlers/web/serpapi.go
Normal file
302
agent/search/handlers/web/serpapi.go
Normal file
|
|
@ -0,0 +1,302 @@
|
|||
package web
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/yaoapp/yao/agent/search/types"
|
||||
)
|
||||
|
||||
const (
|
||||
serpAPIURL = "https://serpapi.com/search.json"
|
||||
serpAPITimeout = 30 * time.Second
|
||||
)
|
||||
|
||||
// SerpAPIProvider implements web search using SerpAPI (supports multiple search engines)
|
||||
type SerpAPIProvider struct {
|
||||
apiKey string
|
||||
maxResults int
|
||||
engine string // Search engine: "google", "bing", "baidu", "yandex", "duckduckgo", etc.
|
||||
}
|
||||
|
||||
// NewSerpAPIProvider creates a new SerpAPI provider
|
||||
func NewSerpAPIProvider(cfg *types.WebConfig) *SerpAPIProvider {
|
||||
apiKey := ""
|
||||
if cfg != nil && cfg.APIKeyEnv != "" {
|
||||
// Support both "$ENV.VAR_NAME" and "VAR_NAME" formats
|
||||
envName := cfg.APIKeyEnv
|
||||
if len(envName) > 5 && envName[:5] == "$ENV." {
|
||||
envName = envName[5:]
|
||||
}
|
||||
apiKey = os.Getenv(envName)
|
||||
}
|
||||
|
||||
maxResults := 10
|
||||
if cfg != nil && cfg.MaxResults > 0 {
|
||||
maxResults = cfg.MaxResults
|
||||
}
|
||||
|
||||
engine := "google" // Default to Google
|
||||
if cfg != nil && cfg.Engine != "" {
|
||||
engine = cfg.Engine
|
||||
}
|
||||
|
||||
return &SerpAPIProvider{
|
||||
apiKey: apiKey,
|
||||
maxResults: maxResults,
|
||||
engine: engine,
|
||||
}
|
||||
}
|
||||
|
||||
// serpAPIResponse represents the response from SerpAPI
|
||||
type serpAPIResponse struct {
|
||||
SearchMetadata serpAPIMetadata `json:"search_metadata"`
|
||||
SearchParameters serpAPIParams `json:"search_parameters"`
|
||||
SearchInformation serpAPIInfo `json:"search_information"`
|
||||
OrganicResults []serpAPIResult `json:"organic_results"`
|
||||
AnswerBox *serpAPIAnswerBox `json:"answer_box,omitempty"`
|
||||
KnowledgeGraph *serpAPIKnowledge `json:"knowledge_graph,omitempty"`
|
||||
RelatedSearches []serpAPIRelated `json:"related_searches,omitempty"`
|
||||
RelatedQuestions []serpAPIQuestion `json:"related_questions,omitempty"`
|
||||
}
|
||||
|
||||
// serpAPIMetadata contains metadata from response
|
||||
type serpAPIMetadata struct {
|
||||
ID string `json:"id"`
|
||||
Status string `json:"status"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
ProcessedAt string `json:"processed_at"`
|
||||
TotalTimeTaken float64 `json:"total_time_taken"`
|
||||
}
|
||||
|
||||
// serpAPIParams contains search parameters from response
|
||||
type serpAPIParams struct {
|
||||
Engine string `json:"engine"`
|
||||
Q string `json:"q"`
|
||||
Location string `json:"location_used"`
|
||||
GoogleDomain string `json:"google_domain"`
|
||||
HL string `json:"hl"`
|
||||
GL string `json:"gl"`
|
||||
Device string `json:"device"`
|
||||
}
|
||||
|
||||
// serpAPIInfo contains search information
|
||||
type serpAPIInfo struct {
|
||||
QueryDisplayed string `json:"query_displayed"`
|
||||
TotalResults int64 `json:"total_results"`
|
||||
TimeTakenDisplayed float64 `json:"time_taken_displayed"`
|
||||
OrganicResultsState string `json:"organic_results_state"`
|
||||
}
|
||||
|
||||
// serpAPIResult represents a single organic search result
|
||||
type serpAPIResult struct {
|
||||
Position int `json:"position"`
|
||||
Title string `json:"title"`
|
||||
Link string `json:"link"`
|
||||
RedirectLink string `json:"redirect_link,omitempty"`
|
||||
DisplayedLink string `json:"displayed_link"`
|
||||
Snippet string `json:"snippet"`
|
||||
Date string `json:"date,omitempty"`
|
||||
CachedPageLink string `json:"cached_page_link,omitempty"`
|
||||
}
|
||||
|
||||
// serpAPIAnswerBox represents the answer box (featured snippet)
|
||||
type serpAPIAnswerBox struct {
|
||||
Type string `json:"type,omitempty"`
|
||||
Title string `json:"title,omitempty"`
|
||||
Snippet string `json:"snippet,omitempty"`
|
||||
Link string `json:"link,omitempty"`
|
||||
}
|
||||
|
||||
// serpAPIKnowledge represents knowledge graph data
|
||||
type serpAPIKnowledge struct {
|
||||
Title string `json:"title,omitempty"`
|
||||
Type string `json:"type,omitempty"`
|
||||
Description string `json:"description,omitempty"`
|
||||
}
|
||||
|
||||
// serpAPIRelated represents related searches
|
||||
type serpAPIRelated struct {
|
||||
Query string `json:"query"`
|
||||
Link string `json:"link"`
|
||||
}
|
||||
|
||||
// serpAPIQuestion represents related questions (People Also Ask)
|
||||
type serpAPIQuestion struct {
|
||||
Question string `json:"question"`
|
||||
Snippet string `json:"snippet,omitempty"`
|
||||
Title string `json:"title,omitempty"`
|
||||
Link string `json:"link,omitempty"`
|
||||
}
|
||||
|
||||
// Search executes a web search using SerpAPI
|
||||
func (p *SerpAPIProvider) Search(req *types.Request) (*types.Result, error) {
|
||||
startTime := time.Now()
|
||||
|
||||
// Validate API key
|
||||
if p.apiKey == "" {
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeWeb,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Error: "SerpAPI API key not configured",
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Determine max results
|
||||
maxResults := p.maxResults
|
||||
if req.Limit > 0 {
|
||||
maxResults = req.Limit
|
||||
}
|
||||
|
||||
// Build query parameters
|
||||
params := url.Values{}
|
||||
params.Set("engine", p.engine)
|
||||
params.Set("api_key", p.apiKey)
|
||||
params.Set("num", fmt.Sprintf("%d", maxResults))
|
||||
|
||||
// Build search query with site restrictions if specified
|
||||
query := req.Query
|
||||
if len(req.Sites) > 0 {
|
||||
siteQuery := ""
|
||||
for i, site := range req.Sites {
|
||||
if i > 0 {
|
||||
siteQuery += " OR "
|
||||
}
|
||||
siteQuery += "site:" + site
|
||||
}
|
||||
query = "(" + siteQuery + ") " + req.Query
|
||||
}
|
||||
params.Set("q", query)
|
||||
|
||||
// Add time range if specified (tbs parameter)
|
||||
if req.TimeRange != "" {
|
||||
tbs := convertSerpAPITimeRange(req.TimeRange)
|
||||
if tbs != "" {
|
||||
params.Set("tbs", tbs)
|
||||
}
|
||||
}
|
||||
|
||||
// Execute API call
|
||||
serpResp, err := p.callAPI(params)
|
||||
if err != nil {
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeWeb,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Duration: time.Since(startTime).Milliseconds(),
|
||||
Error: fmt.Sprintf("SerpAPI error: %v", err),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Convert results
|
||||
items := make([]*types.ResultItem, 0, len(serpResp.OrganicResults))
|
||||
|
||||
// Add answer box as first result if available
|
||||
if serpResp.AnswerBox != nil && serpResp.AnswerBox.Snippet != "" {
|
||||
items = append(items, &types.ResultItem{
|
||||
Type: types.SearchTypeWeb,
|
||||
Title: serpResp.AnswerBox.Title,
|
||||
Content: serpResp.AnswerBox.Snippet,
|
||||
URL: serpResp.AnswerBox.Link,
|
||||
Score: 1.0, // Featured snippet gets highest score
|
||||
Source: req.Source,
|
||||
Metadata: map[string]interface{}{
|
||||
"type": "answer_box",
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// Add organic results
|
||||
for _, r := range serpResp.OrganicResults {
|
||||
// Calculate score based on position (1st = 0.95, 2nd = 0.90, etc.)
|
||||
score := 1.0 - float64(r.Position)*0.05
|
||||
if score < 0.1 {
|
||||
score = 0.1
|
||||
}
|
||||
|
||||
items = append(items, &types.ResultItem{
|
||||
Type: types.SearchTypeWeb,
|
||||
Title: r.Title,
|
||||
Content: r.Snippet,
|
||||
URL: r.Link,
|
||||
Score: score,
|
||||
Source: req.Source,
|
||||
})
|
||||
}
|
||||
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeWeb,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: items,
|
||||
Total: len(items),
|
||||
Duration: time.Since(startTime).Milliseconds(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// callAPI makes the HTTP GET request to SerpAPI
|
||||
func (p *SerpAPIProvider) callAPI(params url.Values) (*serpAPIResponse, error) {
|
||||
// Build URL with query parameters
|
||||
reqURL := serpAPIURL + "?" + params.Encode()
|
||||
|
||||
// Create HTTP request
|
||||
httpReq, err := http.NewRequest(http.MethodGet, reqURL, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
// Execute request
|
||||
client := &http.Client{Timeout: serpAPITimeout}
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("request failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// Read response body
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read response: %w", err)
|
||||
}
|
||||
|
||||
// Check status code
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("API returned status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
// Parse response
|
||||
var serpResp serpAPIResponse
|
||||
if err := json.Unmarshal(respBody, &serpResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse response: %w", err)
|
||||
}
|
||||
|
||||
return &serpResp, nil
|
||||
}
|
||||
|
||||
// convertSerpAPITimeRange converts time range to SerpAPI tbs format
|
||||
func convertSerpAPITimeRange(timeRange string) string {
|
||||
switch timeRange {
|
||||
case "hour":
|
||||
return "qdr:h"
|
||||
case "day":
|
||||
return "qdr:d"
|
||||
case "week":
|
||||
return "qdr:w"
|
||||
case "month":
|
||||
return "qdr:m"
|
||||
case "year":
|
||||
return "qdr:y"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
394
agent/search/handlers/web/serpapi_test.go
Normal file
394
agent/search/handlers/web/serpapi_test.go
Normal file
|
|
@ -0,0 +1,394 @@
|
|||
package web_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/yaoapp/yao/agent/assistant"
|
||||
"github.com/yaoapp/yao/agent/search/handlers/web"
|
||||
"github.com/yaoapp/yao/agent/search/types"
|
||||
"github.com/yaoapp/yao/agent/testutils"
|
||||
)
|
||||
|
||||
// TestSerpAPIProviderWithAssistantConfig tests SerpAPIProvider using web-serpapi assistant config
|
||||
func TestSerpAPIProviderWithAssistantConfig(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Load the web-serpapi test assistant to get its config
|
||||
ast, err := assistant.LoadPath("/assistants/tests/web-serpapi")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, ast)
|
||||
require.NotNil(t, ast.Search)
|
||||
require.NotNil(t, ast.Search.Web)
|
||||
|
||||
// Verify assistant config
|
||||
assert.Equal(t, "tests.web-serpapi", ast.ID)
|
||||
assert.Equal(t, "serpapi", ast.Search.Web.Provider)
|
||||
assert.Equal(t, "$ENV.SERPAPI_API_KEY", ast.Search.Web.APIKeyEnv)
|
||||
assert.Equal(t, 10, ast.Search.Web.MaxResults)
|
||||
|
||||
// Create SerpAPIProvider with assistant's web config
|
||||
provider := web.NewSerpAPIProvider(ast.Search.Web)
|
||||
require.NotNil(t, provider)
|
||||
|
||||
// Execute search
|
||||
req := &types.Request{
|
||||
Query: "Yao App Engine",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
Limit: 5,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// Verify result structure
|
||||
assert.Equal(t, types.SearchTypeWeb, result.Type)
|
||||
assert.Equal(t, "Yao App Engine", result.Query)
|
||||
assert.Equal(t, types.SourceAuto, result.Source)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
|
||||
|
||||
// Verify we got results
|
||||
assert.Greater(t, result.Total, 0)
|
||||
assert.NotEmpty(t, result.Items)
|
||||
assert.Greater(t, result.Duration, int64(0))
|
||||
|
||||
// Verify result item structure
|
||||
for _, item := range result.Items {
|
||||
assert.Equal(t, types.SearchTypeWeb, item.Type)
|
||||
assert.Equal(t, types.SourceAuto, item.Source)
|
||||
assert.NotEmpty(t, item.Title)
|
||||
assert.NotEmpty(t, item.URL)
|
||||
assert.Greater(t, item.Score, 0.0)
|
||||
}
|
||||
|
||||
t.Logf("Search returned %d results in %dms", result.Total, result.Duration)
|
||||
}
|
||||
|
||||
// TestSerpAPIProviderWithSiteRestriction tests SerpAPIProvider with domain restriction
|
||||
func TestSerpAPIProviderWithSiteRestriction(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Load the web-serpapi test assistant
|
||||
ast, err := assistant.LoadPath("/assistants/tests/web-serpapi")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, ast.Search)
|
||||
require.NotNil(t, ast.Search.Web)
|
||||
|
||||
// Create SerpAPIProvider
|
||||
provider := web.NewSerpAPIProvider(ast.Search.Web)
|
||||
|
||||
// Execute search with site restriction
|
||||
req := &types.Request{
|
||||
Query: "documentation",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceHook,
|
||||
Sites: []string{"github.com"},
|
||||
Limit: 3,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
assert.Equal(t, types.SearchTypeWeb, result.Type)
|
||||
assert.Equal(t, types.SourceHook, result.Source)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
|
||||
require.NotEmpty(t, result.Items, "Search should return results")
|
||||
|
||||
// All results should be from github.com
|
||||
for _, item := range result.Items {
|
||||
assert.Contains(t, item.URL, "github.com", "Result URL should be from github.com")
|
||||
}
|
||||
t.Logf("Site-restricted search returned %d results from github.com", result.Total)
|
||||
}
|
||||
|
||||
// TestSerpAPIProviderWithMultipleSites tests SerpAPIProvider with multiple domain restrictions
|
||||
func TestSerpAPIProviderWithMultipleSites(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Load the web-serpapi test assistant
|
||||
ast, err := assistant.LoadPath("/assistants/tests/web-serpapi")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, ast.Search)
|
||||
require.NotNil(t, ast.Search.Web)
|
||||
|
||||
// Create SerpAPIProvider
|
||||
provider := web.NewSerpAPIProvider(ast.Search.Web)
|
||||
|
||||
// Execute search with multiple site restrictions
|
||||
req := &types.Request{
|
||||
Query: "golang tutorial",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
Sites: []string{"github.com", "golang.org"},
|
||||
Limit: 5,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
|
||||
require.NotEmpty(t, result.Items, "Search should return results")
|
||||
|
||||
// Results should be from either github.com or golang.org
|
||||
for _, item := range result.Items {
|
||||
isValidSite := false
|
||||
for _, site := range req.Sites {
|
||||
if containsSite(item.URL, site) {
|
||||
isValidSite = true
|
||||
break
|
||||
}
|
||||
}
|
||||
assert.True(t, isValidSite, "Result URL should be from github.com or golang.org: %s", item.URL)
|
||||
}
|
||||
t.Logf("Multi-site search returned %d results", result.Total)
|
||||
}
|
||||
|
||||
// TestSerpAPIProviderWithTimeRange tests SerpAPIProvider with time range filter
|
||||
func TestSerpAPIProviderWithTimeRange(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Load the web-serpapi test assistant
|
||||
ast, err := assistant.LoadPath("/assistants/tests/web-serpapi")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, ast.Search)
|
||||
require.NotNil(t, ast.Search.Web)
|
||||
|
||||
// Create SerpAPIProvider
|
||||
provider := web.NewSerpAPIProvider(ast.Search.Web)
|
||||
|
||||
// Execute search with time range
|
||||
req := &types.Request{
|
||||
Query: "artificial intelligence news",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
TimeRange: "week", // Last week
|
||||
Limit: 5,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
|
||||
t.Logf("Time-ranged search (last week) returned %d results in %dms", result.Total, result.Duration)
|
||||
}
|
||||
|
||||
// TestSerpAPIProviderWithoutAPIKey tests graceful degradation when API key is missing
|
||||
func TestSerpAPIProviderWithoutAPIKey(t *testing.T) {
|
||||
// Create provider with nil config (no API key)
|
||||
provider := web.NewSerpAPIProvider(nil)
|
||||
require.NotNil(t, provider)
|
||||
|
||||
req := &types.Request{
|
||||
Query: "test query",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
|
||||
// Should not return error, but result should have error message
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
assert.Equal(t, types.SearchTypeWeb, result.Type)
|
||||
assert.Equal(t, "test query", result.Query)
|
||||
assert.NotEmpty(t, result.Error)
|
||||
assert.Contains(t, result.Error, "API key")
|
||||
assert.Empty(t, result.Items)
|
||||
assert.Equal(t, 0, result.Total)
|
||||
}
|
||||
|
||||
// TestSerpAPIProviderWithEmptyConfig tests provider with empty config
|
||||
func TestSerpAPIProviderWithEmptyConfig(t *testing.T) {
|
||||
// Create provider with empty config
|
||||
cfg := &types.WebConfig{}
|
||||
provider := web.NewSerpAPIProvider(cfg)
|
||||
require.NotNil(t, provider)
|
||||
|
||||
req := &types.Request{
|
||||
Query: "test query",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceUser,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
assert.NotEmpty(t, result.Error)
|
||||
assert.Contains(t, result.Error, "API key")
|
||||
}
|
||||
|
||||
// TestSerpAPIProviderMaxResults tests that max_results from config is respected
|
||||
func TestSerpAPIProviderMaxResults(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Load the web-serpapi test assistant
|
||||
ast, err := assistant.LoadPath("/assistants/tests/web-serpapi")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, ast.Search)
|
||||
require.NotNil(t, ast.Search.Web)
|
||||
|
||||
// Create SerpAPIProvider
|
||||
provider := web.NewSerpAPIProvider(ast.Search.Web)
|
||||
|
||||
// Execute search without limit (should use config's max_results)
|
||||
req := &types.Request{
|
||||
Query: "machine learning",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
// No Limit set, should use config's max_results (10)
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
|
||||
|
||||
// Should respect max_results from config (+1 for possible answer box)
|
||||
assert.LessOrEqual(t, result.Total, ast.Search.Web.MaxResults+1)
|
||||
t.Logf("Search without limit returned %d results (max: %d)", result.Total, ast.Search.Web.MaxResults)
|
||||
|
||||
// Execute search with explicit limit
|
||||
req2 := &types.Request{
|
||||
Query: "machine learning",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
Limit: 3, // Override config's max_results
|
||||
}
|
||||
|
||||
result2, err := provider.Search(req2)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result2)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result2.Error, "Search should succeed with valid API key, got error: %s", result2.Error)
|
||||
|
||||
// Should respect request's limit (+1 for possible answer box)
|
||||
assert.LessOrEqual(t, result2.Total, 4)
|
||||
t.Logf("Search with limit=3 returned %d results", result2.Total)
|
||||
}
|
||||
|
||||
// TestSerpAPIProviderWithBingEngine tests SerpAPIProvider with Bing search engine
|
||||
func TestSerpAPIProviderWithBingEngine(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Load the web-serpapi test assistant to get base config
|
||||
ast, err := assistant.LoadPath("/assistants/tests/web-serpapi")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, ast.Search)
|
||||
require.NotNil(t, ast.Search.Web)
|
||||
|
||||
// Create config with Bing engine
|
||||
bingConfig := &types.WebConfig{
|
||||
Provider: "serpapi",
|
||||
APIKeyEnv: ast.Search.Web.APIKeyEnv,
|
||||
MaxResults: 5,
|
||||
Engine: "bing",
|
||||
}
|
||||
|
||||
// Create SerpAPIProvider with Bing engine
|
||||
provider := web.NewSerpAPIProvider(bingConfig)
|
||||
require.NotNil(t, provider)
|
||||
|
||||
// Execute search
|
||||
req := &types.Request{
|
||||
Query: "Golang programming",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
Limit: 5,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// Verify result structure
|
||||
assert.Equal(t, types.SearchTypeWeb, result.Type)
|
||||
assert.Equal(t, "Golang programming", result.Query)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result.Error, "Bing search should succeed with valid API key, got error: %s", result.Error)
|
||||
|
||||
// Verify we got results
|
||||
assert.Greater(t, result.Total, 0)
|
||||
assert.NotEmpty(t, result.Items)
|
||||
|
||||
t.Logf("Bing search returned %d results in %dms", result.Total, result.Duration)
|
||||
}
|
||||
|
||||
// TestSerpAPIProviderEngineDefault tests that default engine is Google
|
||||
func TestSerpAPIProviderEngineDefault(t *testing.T) {
|
||||
// Create provider with config that has no engine specified
|
||||
cfg := &types.WebConfig{
|
||||
Provider: "serpapi",
|
||||
APIKeyEnv: "SERPAPI_API_KEY",
|
||||
MaxResults: 10,
|
||||
// Engine not set - should default to "google"
|
||||
}
|
||||
|
||||
provider := web.NewSerpAPIProvider(cfg)
|
||||
require.NotNil(t, provider)
|
||||
|
||||
// We can't directly check the engine field since it's private,
|
||||
// but we verify the provider is created successfully
|
||||
// The actual engine usage is tested in integration tests
|
||||
}
|
||||
|
||||
// containsSite checks if url contains the site domain
|
||||
func containsSite(url, site string) bool {
|
||||
return len(url) >= len(site) && containsHelper(url, site)
|
||||
}
|
||||
|
||||
// containsHelper is a helper function for string containment check
|
||||
func containsHelper(s, substr string) bool {
|
||||
for i := 0; i <= len(s)-len(substr); i++ {
|
||||
if s[i:i+len(substr)] == substr {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
280
agent/search/handlers/web/serper.go
Normal file
280
agent/search/handlers/web/serper.go
Normal file
|
|
@ -0,0 +1,280 @@
|
|||
package web
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/yaoapp/yao/agent/search/types"
|
||||
)
|
||||
|
||||
const (
|
||||
serperAPIURL = "https://google.serper.dev/search"
|
||||
serperAPITimeout = 30 * time.Second
|
||||
)
|
||||
|
||||
// SerperProvider implements web search using Serper API (serper.dev)
|
||||
type SerperProvider struct {
|
||||
apiKey string
|
||||
maxResults int
|
||||
}
|
||||
|
||||
// NewSerperProvider creates a new Serper provider
|
||||
func NewSerperProvider(cfg *types.WebConfig) *SerperProvider {
|
||||
apiKey := ""
|
||||
if cfg != nil && cfg.APIKeyEnv != "" {
|
||||
// Support both "$ENV.VAR_NAME" and "VAR_NAME" formats
|
||||
envName := cfg.APIKeyEnv
|
||||
if len(envName) > 5 && envName[:5] == "$ENV." {
|
||||
envName = envName[5:]
|
||||
}
|
||||
apiKey = os.Getenv(envName)
|
||||
}
|
||||
|
||||
maxResults := 10
|
||||
if cfg != nil && cfg.MaxResults > 0 {
|
||||
maxResults = cfg.MaxResults
|
||||
}
|
||||
|
||||
return &SerperProvider{
|
||||
apiKey: apiKey,
|
||||
maxResults: maxResults,
|
||||
}
|
||||
}
|
||||
|
||||
// serperRequest represents the request body for Serper API
|
||||
type serperRequest struct {
|
||||
Q string `json:"q"` // Search query
|
||||
Num int `json:"num,omitempty"` // Number of results (default: 10, max: 100)
|
||||
GL string `json:"gl,omitempty"` // Country code (e.g., "us", "cn")
|
||||
HL string `json:"hl,omitempty"` // Language code (e.g., "en", "zh-cn")
|
||||
TBS string `json:"tbs,omitempty"` // Time-based search (qdr:h, qdr:d, qdr:w, qdr:m, qdr:y)
|
||||
Page int `json:"page,omitempty"` // Page number (default: 1)
|
||||
AutoCor bool `json:"autocorrect"` // Auto-correct spelling
|
||||
}
|
||||
|
||||
// serperResponse represents the response from Serper API
|
||||
type serperResponse struct {
|
||||
SearchParameters serperSearchParams `json:"searchParameters"`
|
||||
Organic []serperResult `json:"organic"`
|
||||
AnswerBox *serperAnswerBox `json:"answerBox,omitempty"`
|
||||
KnowledgeGraph *serperKnowledge `json:"knowledgeGraph,omitempty"`
|
||||
RelatedSearches []serperRelated `json:"relatedSearches,omitempty"`
|
||||
}
|
||||
|
||||
// serperSearchParams contains search parameters from response
|
||||
type serperSearchParams struct {
|
||||
Q string `json:"q"`
|
||||
Type string `json:"type"`
|
||||
GL string `json:"gl"`
|
||||
HL string `json:"hl"`
|
||||
Num int `json:"num"`
|
||||
}
|
||||
|
||||
// serperResult represents a single organic search result
|
||||
type serperResult struct {
|
||||
Title string `json:"title"`
|
||||
Link string `json:"link"`
|
||||
Snippet string `json:"snippet"`
|
||||
Position int `json:"position"`
|
||||
Date string `json:"date,omitempty"`
|
||||
}
|
||||
|
||||
// serperAnswerBox represents the answer box (featured snippet)
|
||||
type serperAnswerBox struct {
|
||||
Title string `json:"title,omitempty"`
|
||||
Snippet string `json:"snippet,omitempty"`
|
||||
Link string `json:"link,omitempty"`
|
||||
}
|
||||
|
||||
// serperKnowledge represents knowledge graph data
|
||||
type serperKnowledge struct {
|
||||
Title string `json:"title,omitempty"`
|
||||
Type string `json:"type,omitempty"`
|
||||
Description string `json:"description,omitempty"`
|
||||
}
|
||||
|
||||
// serperRelated represents related searches
|
||||
type serperRelated struct {
|
||||
Query string `json:"query"`
|
||||
}
|
||||
|
||||
// Search executes a web search using Serper API
|
||||
func (p *SerperProvider) Search(req *types.Request) (*types.Result, error) {
|
||||
startTime := time.Now()
|
||||
|
||||
// Validate API key
|
||||
if p.apiKey == "" {
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeWeb,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Error: "Serper API key not configured",
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Determine max results
|
||||
maxResults := p.maxResults
|
||||
if req.Limit > 0 {
|
||||
maxResults = req.Limit
|
||||
}
|
||||
|
||||
// Build search query with site restrictions if specified
|
||||
query := req.Query
|
||||
if len(req.Sites) > 0 {
|
||||
// Serper uses "site:domain" syntax in query
|
||||
if len(req.Sites) == 1 {
|
||||
query = "site:" + req.Sites[0] + " " + req.Query
|
||||
} else {
|
||||
// Multiple sites: (site:domain1 OR site:domain2) query
|
||||
siteQuery := ""
|
||||
for i, site := range req.Sites {
|
||||
if i > 0 {
|
||||
siteQuery += " OR "
|
||||
}
|
||||
siteQuery += "site:" + site
|
||||
}
|
||||
query = "(" + siteQuery + ") " + req.Query
|
||||
}
|
||||
}
|
||||
|
||||
// Build request body
|
||||
serperReq := serperRequest{
|
||||
Q: query,
|
||||
Num: maxResults,
|
||||
AutoCor: true,
|
||||
}
|
||||
|
||||
// Add time range if specified
|
||||
if req.TimeRange != "" {
|
||||
serperReq.TBS = convertSerperTimeRange(req.TimeRange)
|
||||
}
|
||||
|
||||
// Execute API call
|
||||
serperResp, err := p.callAPI(&serperReq)
|
||||
if err != nil {
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeWeb,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Duration: time.Since(startTime).Milliseconds(),
|
||||
Error: fmt.Sprintf("Serper API error: %v", err),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Convert results
|
||||
items := make([]*types.ResultItem, 0, len(serperResp.Organic))
|
||||
|
||||
// Add answer box as first result if available
|
||||
if serperResp.AnswerBox != nil && serperResp.AnswerBox.Snippet != "" {
|
||||
items = append(items, &types.ResultItem{
|
||||
Type: types.SearchTypeWeb,
|
||||
Title: serperResp.AnswerBox.Title,
|
||||
Content: serperResp.AnswerBox.Snippet,
|
||||
URL: serperResp.AnswerBox.Link,
|
||||
Score: 1.0, // Featured snippet gets highest score
|
||||
Source: req.Source,
|
||||
Metadata: map[string]interface{}{
|
||||
"type": "answer_box",
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// Add organic results
|
||||
for _, r := range serperResp.Organic {
|
||||
// Calculate score based on position (1st = 0.95, 2nd = 0.90, etc.)
|
||||
score := 1.0 - float64(r.Position)*0.05
|
||||
if score < 0.1 {
|
||||
score = 0.1
|
||||
}
|
||||
|
||||
items = append(items, &types.ResultItem{
|
||||
Type: types.SearchTypeWeb,
|
||||
Title: r.Title,
|
||||
Content: r.Snippet,
|
||||
URL: r.Link,
|
||||
Score: score,
|
||||
Source: req.Source,
|
||||
})
|
||||
}
|
||||
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeWeb,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: items,
|
||||
Total: len(items),
|
||||
Duration: time.Since(startTime).Milliseconds(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// callAPI makes the HTTP POST request to Serper API
|
||||
func (p *SerperProvider) callAPI(req *serperRequest) (*serperResponse, error) {
|
||||
// Serialize request body
|
||||
body, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||
}
|
||||
|
||||
// Create HTTP request
|
||||
httpReq, err := http.NewRequest(http.MethodPost, serperAPIURL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("X-API-KEY", p.apiKey)
|
||||
|
||||
// Execute request
|
||||
client := &http.Client{Timeout: serperAPITimeout}
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("request failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// Read response body
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read response: %w", err)
|
||||
}
|
||||
|
||||
// Check status code
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("API returned status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
// Parse response
|
||||
var serperResp serperResponse
|
||||
if err := json.Unmarshal(respBody, &serperResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse response: %w", err)
|
||||
}
|
||||
|
||||
return &serperResp, nil
|
||||
}
|
||||
|
||||
// convertSerperTimeRange converts time range to Serper tbs format
|
||||
func convertSerperTimeRange(timeRange string) string {
|
||||
switch timeRange {
|
||||
case "hour":
|
||||
return "qdr:h"
|
||||
case "day":
|
||||
return "qdr:d"
|
||||
case "week":
|
||||
return "qdr:w"
|
||||
case "month":
|
||||
return "qdr:m"
|
||||
case "year":
|
||||
return "qdr:y"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
327
agent/search/handlers/web/serper_test.go
Normal file
327
agent/search/handlers/web/serper_test.go
Normal file
|
|
@ -0,0 +1,327 @@
|
|||
package web_test
|
||||
|
||||
import (
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/yaoapp/yao/agent/assistant"
|
||||
"github.com/yaoapp/yao/agent/search/handlers/web"
|
||||
"github.com/yaoapp/yao/agent/search/types"
|
||||
"github.com/yaoapp/yao/agent/testutils"
|
||||
)
|
||||
|
||||
// skipIfNoSerperKey skips the test if SERPER_API_KEY is not set
|
||||
// Note: Serper (serper.dev) requires registration at https://serper.dev
|
||||
func skipIfNoSerperKey(t *testing.T) {
|
||||
if os.Getenv("SERPER_API_KEY") == "" {
|
||||
t.Skip("Skipping Serper test: SERPER_API_KEY not set. Register at https://serper.dev for free 2500 queries.")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSerperProviderWithAssistantConfig tests SerperProvider using web-serper assistant config
|
||||
func TestSerperProviderWithAssistantConfig(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
skipIfNoSerperKey(t)
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Load the web-serper test assistant to get its config
|
||||
ast, err := assistant.LoadPath("/assistants/tests/web-serper")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, ast)
|
||||
require.NotNil(t, ast.Search)
|
||||
require.NotNil(t, ast.Search.Web)
|
||||
|
||||
// Verify assistant config
|
||||
assert.Equal(t, "tests.web-serper", ast.ID)
|
||||
assert.Equal(t, "serper", ast.Search.Web.Provider)
|
||||
assert.Equal(t, "$ENV.SERPER_API_KEY", ast.Search.Web.APIKeyEnv)
|
||||
assert.Equal(t, 10, ast.Search.Web.MaxResults)
|
||||
|
||||
// Create SerperProvider with assistant's web config
|
||||
provider := web.NewSerperProvider(ast.Search.Web)
|
||||
require.NotNil(t, provider)
|
||||
|
||||
// Execute search
|
||||
req := &types.Request{
|
||||
Query: "Yao App Engine",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
Limit: 5,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// Verify result structure
|
||||
assert.Equal(t, types.SearchTypeWeb, result.Type)
|
||||
assert.Equal(t, "Yao App Engine", result.Query)
|
||||
assert.Equal(t, types.SourceAuto, result.Source)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
|
||||
|
||||
// Verify we got results
|
||||
assert.Greater(t, result.Total, 0)
|
||||
assert.NotEmpty(t, result.Items)
|
||||
assert.Greater(t, result.Duration, int64(0))
|
||||
|
||||
// Verify result item structure
|
||||
for _, item := range result.Items {
|
||||
assert.Equal(t, types.SearchTypeWeb, item.Type)
|
||||
assert.Equal(t, types.SourceAuto, item.Source)
|
||||
assert.NotEmpty(t, item.Title)
|
||||
assert.NotEmpty(t, item.URL)
|
||||
assert.Greater(t, item.Score, 0.0)
|
||||
}
|
||||
|
||||
t.Logf("Search returned %d results in %dms", result.Total, result.Duration)
|
||||
}
|
||||
|
||||
// TestSerperProviderWithSiteRestriction tests SerperProvider with domain restriction
|
||||
func TestSerperProviderWithSiteRestriction(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
skipIfNoSerperKey(t)
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Load the web-serper test assistant
|
||||
ast, err := assistant.LoadPath("/assistants/tests/web-serper")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, ast.Search)
|
||||
require.NotNil(t, ast.Search.Web)
|
||||
|
||||
// Create SerperProvider
|
||||
provider := web.NewSerperProvider(ast.Search.Web)
|
||||
|
||||
// Execute search with site restriction
|
||||
req := &types.Request{
|
||||
Query: "documentation",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceHook,
|
||||
Sites: []string{"github.com"},
|
||||
Limit: 3,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
assert.Equal(t, types.SearchTypeWeb, result.Type)
|
||||
assert.Equal(t, types.SourceHook, result.Source)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
|
||||
require.NotEmpty(t, result.Items, "Search should return results")
|
||||
|
||||
// All results should be from github.com
|
||||
for _, item := range result.Items {
|
||||
assert.Contains(t, item.URL, "github.com", "Result URL should be from github.com")
|
||||
}
|
||||
t.Logf("Site-restricted search returned %d results from github.com", result.Total)
|
||||
}
|
||||
|
||||
// TestSerperProviderWithMultipleSites tests SerperProvider with multiple domain restrictions
|
||||
func TestSerperProviderWithMultipleSites(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
skipIfNoSerperKey(t)
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Load the web-serper test assistant
|
||||
ast, err := assistant.LoadPath("/assistants/tests/web-serper")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, ast.Search)
|
||||
require.NotNil(t, ast.Search.Web)
|
||||
|
||||
// Create SerperProvider
|
||||
provider := web.NewSerperProvider(ast.Search.Web)
|
||||
|
||||
// Execute search with multiple site restrictions
|
||||
req := &types.Request{
|
||||
Query: "golang tutorial",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
Sites: []string{"github.com", "golang.org"},
|
||||
Limit: 5,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
|
||||
require.NotEmpty(t, result.Items, "Search should return results")
|
||||
|
||||
// Results should be from either github.com or golang.org
|
||||
for _, item := range result.Items {
|
||||
isValidSite := false
|
||||
for _, site := range req.Sites {
|
||||
if contains(item.URL, site) {
|
||||
isValidSite = true
|
||||
break
|
||||
}
|
||||
}
|
||||
assert.True(t, isValidSite, "Result URL should be from github.com or golang.org: %s", item.URL)
|
||||
}
|
||||
t.Logf("Multi-site search returned %d results", result.Total)
|
||||
}
|
||||
|
||||
// TestSerperProviderWithTimeRange tests SerperProvider with time range filter
|
||||
func TestSerperProviderWithTimeRange(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
skipIfNoSerperKey(t)
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Load the web-serper test assistant
|
||||
ast, err := assistant.LoadPath("/assistants/tests/web-serper")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, ast.Search)
|
||||
require.NotNil(t, ast.Search.Web)
|
||||
|
||||
// Create SerperProvider
|
||||
provider := web.NewSerperProvider(ast.Search.Web)
|
||||
|
||||
// Execute search with time range
|
||||
req := &types.Request{
|
||||
Query: "artificial intelligence news",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
TimeRange: "week", // Last week
|
||||
Limit: 5,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
|
||||
t.Logf("Time-ranged search (last week) returned %d results in %dms", result.Total, result.Duration)
|
||||
}
|
||||
|
||||
// TestSerperProviderWithoutAPIKey tests graceful degradation when API key is missing
|
||||
func TestSerperProviderWithoutAPIKey(t *testing.T) {
|
||||
// Create provider with nil config (no API key)
|
||||
provider := web.NewSerperProvider(nil)
|
||||
require.NotNil(t, provider)
|
||||
|
||||
req := &types.Request{
|
||||
Query: "test query",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
|
||||
// Should not return error, but result should have error message
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
assert.Equal(t, types.SearchTypeWeb, result.Type)
|
||||
assert.Equal(t, "test query", result.Query)
|
||||
assert.NotEmpty(t, result.Error)
|
||||
assert.Contains(t, result.Error, "API key")
|
||||
assert.Empty(t, result.Items)
|
||||
assert.Equal(t, 0, result.Total)
|
||||
}
|
||||
|
||||
// TestSerperProviderWithEmptyConfig tests provider with empty config
|
||||
func TestSerperProviderWithEmptyConfig(t *testing.T) {
|
||||
// Create provider with empty config
|
||||
cfg := &types.WebConfig{}
|
||||
provider := web.NewSerperProvider(cfg)
|
||||
require.NotNil(t, provider)
|
||||
|
||||
req := &types.Request{
|
||||
Query: "test query",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceUser,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
assert.NotEmpty(t, result.Error)
|
||||
assert.Contains(t, result.Error, "API key")
|
||||
}
|
||||
|
||||
// TestSerperProviderMaxResults tests that max_results from config is respected
|
||||
func TestSerperProviderMaxResults(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
skipIfNoSerperKey(t)
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Load the web-serper test assistant
|
||||
ast, err := assistant.LoadPath("/assistants/tests/web-serper")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, ast.Search)
|
||||
require.NotNil(t, ast.Search.Web)
|
||||
|
||||
// Create SerperProvider
|
||||
provider := web.NewSerperProvider(ast.Search.Web)
|
||||
|
||||
// Execute search without limit (should use config's max_results)
|
||||
req := &types.Request{
|
||||
Query: "machine learning",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
// No Limit set, should use config's max_results (10)
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
|
||||
|
||||
// Should respect max_results from config (+1 for possible answer box)
|
||||
assert.LessOrEqual(t, result.Total, ast.Search.Web.MaxResults+1)
|
||||
t.Logf("Search without limit returned %d results (max: %d)", result.Total, ast.Search.Web.MaxResults)
|
||||
|
||||
// Execute search with explicit limit
|
||||
req2 := &types.Request{
|
||||
Query: "machine learning",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
Limit: 3, // Override config's max_results
|
||||
}
|
||||
|
||||
result2, err := provider.Search(req2)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result2)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result2.Error, "Search should succeed with valid API key, got error: %s", result2.Error)
|
||||
|
||||
// Should respect request's limit (+1 for possible answer box)
|
||||
assert.LessOrEqual(t, result2.Total, 4)
|
||||
t.Logf("Search with limit=3 returned %d results", result2.Total)
|
||||
}
|
||||
|
||||
// contains checks if s contains substr (uses containsSite from serpapi_test.go)
|
||||
func contains(s, substr string) bool {
|
||||
return containsSite(s, substr)
|
||||
}
|
||||
192
agent/search/handlers/web/tavily.go
Normal file
192
agent/search/handlers/web/tavily.go
Normal file
|
|
@ -0,0 +1,192 @@
|
|||
package web
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/yaoapp/yao/agent/search/types"
|
||||
)
|
||||
|
||||
const (
|
||||
tavilyAPIURL = "https://api.tavily.com/search"
|
||||
tavilyAPITimeout = 30 * time.Second
|
||||
)
|
||||
|
||||
// TavilyProvider implements web search using Tavily API
|
||||
type TavilyProvider struct {
|
||||
apiKey string
|
||||
maxResults int
|
||||
}
|
||||
|
||||
// NewTavilyProvider creates a new Tavily provider
|
||||
func NewTavilyProvider(cfg *types.WebConfig) *TavilyProvider {
|
||||
apiKey := ""
|
||||
if cfg != nil && cfg.APIKeyEnv != "" {
|
||||
// Support both "$ENV.VAR_NAME" and "VAR_NAME" formats
|
||||
envName := cfg.APIKeyEnv
|
||||
if len(envName) > 5 && envName[:5] == "$ENV." {
|
||||
envName = envName[5:]
|
||||
}
|
||||
apiKey = os.Getenv(envName)
|
||||
}
|
||||
|
||||
maxResults := 10
|
||||
if cfg != nil && cfg.MaxResults > 0 {
|
||||
maxResults = cfg.MaxResults
|
||||
}
|
||||
|
||||
return &TavilyProvider{
|
||||
apiKey: apiKey,
|
||||
maxResults: maxResults,
|
||||
}
|
||||
}
|
||||
|
||||
// tavilyRequest represents the request body for Tavily API
|
||||
type tavilyRequest struct {
|
||||
APIKey string `json:"api_key"`
|
||||
Query string `json:"query"`
|
||||
SearchDepth string `json:"search_depth,omitempty"` // "basic" or "advanced"
|
||||
IncludeAnswer bool `json:"include_answer,omitempty"` // Include AI-generated answer
|
||||
IncludeRawContent bool `json:"include_raw_content,omitempty"` // Include raw HTML content
|
||||
MaxResults int `json:"max_results,omitempty"` // Max number of results
|
||||
IncludeDomains []string `json:"include_domains,omitempty"` // Limit to specific domains
|
||||
ExcludeDomains []string `json:"exclude_domains,omitempty"` // Exclude specific domains
|
||||
}
|
||||
|
||||
// tavilyResponse represents the response from Tavily API
|
||||
type tavilyResponse struct {
|
||||
Query string `json:"query"`
|
||||
Answer string `json:"answer,omitempty"`
|
||||
Results []tavilyResult `json:"results"`
|
||||
}
|
||||
|
||||
// tavilyResult represents a single search result from Tavily
|
||||
type tavilyResult struct {
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
Content string `json:"content"`
|
||||
Score float64 `json:"score"`
|
||||
RawContent string `json:"raw_content,omitempty"`
|
||||
}
|
||||
|
||||
// Search executes a web search using Tavily API
|
||||
func (p *TavilyProvider) Search(req *types.Request) (*types.Result, error) {
|
||||
startTime := time.Now()
|
||||
|
||||
// Validate API key
|
||||
if p.apiKey == "" {
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeWeb,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Error: "Tavily API key not configured",
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Determine max results
|
||||
maxResults := p.maxResults
|
||||
if req.Limit > 0 {
|
||||
maxResults = req.Limit
|
||||
}
|
||||
|
||||
// Build request body
|
||||
tavilyReq := tavilyRequest{
|
||||
APIKey: p.apiKey,
|
||||
Query: req.Query,
|
||||
SearchDepth: "basic",
|
||||
IncludeAnswer: false,
|
||||
MaxResults: maxResults,
|
||||
}
|
||||
|
||||
// Add domain restrictions if specified
|
||||
if len(req.Sites) > 0 {
|
||||
tavilyReq.IncludeDomains = req.Sites
|
||||
}
|
||||
|
||||
// Execute API call
|
||||
tavilyResp, err := p.callAPI(&tavilyReq)
|
||||
if err != nil {
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeWeb,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: []*types.ResultItem{},
|
||||
Total: 0,
|
||||
Duration: time.Since(startTime).Milliseconds(),
|
||||
Error: fmt.Sprintf("Tavily API error: %v", err),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Convert results
|
||||
items := make([]*types.ResultItem, 0, len(tavilyResp.Results))
|
||||
for _, r := range tavilyResp.Results {
|
||||
items = append(items, &types.ResultItem{
|
||||
Type: types.SearchTypeWeb,
|
||||
Title: r.Title,
|
||||
Content: r.Content,
|
||||
URL: r.URL,
|
||||
Score: r.Score,
|
||||
Source: req.Source,
|
||||
})
|
||||
}
|
||||
|
||||
return &types.Result{
|
||||
Type: types.SearchTypeWeb,
|
||||
Query: req.Query,
|
||||
Source: req.Source,
|
||||
Items: items,
|
||||
Total: len(items),
|
||||
Duration: time.Since(startTime).Milliseconds(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// callAPI makes the HTTP request to Tavily API
|
||||
func (p *TavilyProvider) callAPI(req *tavilyRequest) (*tavilyResponse, error) {
|
||||
// Serialize request body
|
||||
body, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||
}
|
||||
|
||||
// Create HTTP request
|
||||
httpReq, err := http.NewRequest(http.MethodPost, tavilyAPIURL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
|
||||
// Execute request
|
||||
client := &http.Client{Timeout: tavilyAPITimeout}
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("request failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// Read response body
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read response: %w", err)
|
||||
}
|
||||
|
||||
// Check status code
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("API returned status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
// Parse response
|
||||
var tavilyResp tavilyResponse
|
||||
if err := json.Unmarshal(respBody, &tavilyResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse response: %w", err)
|
||||
}
|
||||
|
||||
return &tavilyResp, nil
|
||||
}
|
||||
223
agent/search/handlers/web/tavily_test.go
Normal file
223
agent/search/handlers/web/tavily_test.go
Normal file
|
|
@ -0,0 +1,223 @@
|
|||
package web_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/yaoapp/yao/agent/assistant"
|
||||
"github.com/yaoapp/yao/agent/search/handlers/web"
|
||||
"github.com/yaoapp/yao/agent/search/types"
|
||||
"github.com/yaoapp/yao/agent/testutils"
|
||||
)
|
||||
|
||||
// TestTavilyProviderWithAssistantConfig tests TavilyProvider using web-tavily assistant config
|
||||
func TestTavilyProviderWithAssistantConfig(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Load the web-tavily test assistant to get its config
|
||||
ast, err := assistant.LoadPath("/assistants/tests/web-tavily")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, ast)
|
||||
require.NotNil(t, ast.Search)
|
||||
require.NotNil(t, ast.Search.Web)
|
||||
|
||||
// Verify assistant config
|
||||
assert.Equal(t, "tests.web-tavily", ast.ID)
|
||||
assert.Equal(t, "tavily", ast.Search.Web.Provider)
|
||||
assert.Equal(t, "$ENV.TAVILY_API_KEY", ast.Search.Web.APIKeyEnv)
|
||||
assert.Equal(t, 10, ast.Search.Web.MaxResults)
|
||||
|
||||
// Create TavilyProvider with assistant's web config
|
||||
provider := web.NewTavilyProvider(ast.Search.Web)
|
||||
require.NotNil(t, provider)
|
||||
|
||||
// Execute search
|
||||
req := &types.Request{
|
||||
Query: "Yao App Engine",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
Limit: 5,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// Verify result structure
|
||||
assert.Equal(t, types.SearchTypeWeb, result.Type)
|
||||
assert.Equal(t, "Yao App Engine", result.Query)
|
||||
assert.Equal(t, types.SourceAuto, result.Source)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
|
||||
|
||||
// Verify we got results
|
||||
assert.Greater(t, result.Total, 0)
|
||||
assert.NotEmpty(t, result.Items)
|
||||
assert.Greater(t, result.Duration, int64(0))
|
||||
|
||||
// Verify result item structure
|
||||
for _, item := range result.Items {
|
||||
assert.Equal(t, types.SearchTypeWeb, item.Type)
|
||||
assert.Equal(t, types.SourceAuto, item.Source)
|
||||
assert.NotEmpty(t, item.Title)
|
||||
assert.NotEmpty(t, item.URL)
|
||||
// Content may be empty for some results
|
||||
}
|
||||
|
||||
t.Logf("Search returned %d results in %dms", result.Total, result.Duration)
|
||||
}
|
||||
|
||||
// TestTavilyProviderWithSiteRestriction tests TavilyProvider with domain restriction
|
||||
func TestTavilyProviderWithSiteRestriction(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Load the web-tavily test assistant
|
||||
ast, err := assistant.LoadPath("/assistants/tests/web-tavily")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, ast.Search)
|
||||
require.NotNil(t, ast.Search.Web)
|
||||
|
||||
// Create TavilyProvider
|
||||
provider := web.NewTavilyProvider(ast.Search.Web)
|
||||
|
||||
// Execute search with site restriction
|
||||
req := &types.Request{
|
||||
Query: "documentation",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceHook,
|
||||
Sites: []string{"github.com"},
|
||||
Limit: 3,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
assert.Equal(t, types.SearchTypeWeb, result.Type)
|
||||
assert.Equal(t, types.SourceHook, result.Source)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
|
||||
require.NotEmpty(t, result.Items, "Search should return results")
|
||||
|
||||
// All results should be from github.com
|
||||
for _, item := range result.Items {
|
||||
assert.Contains(t, item.URL, "github.com", "Result URL should be from github.com")
|
||||
}
|
||||
t.Logf("Site-restricted search returned %d results from github.com", result.Total)
|
||||
}
|
||||
|
||||
// TestTavilyProviderWithoutAPIKey tests graceful degradation when API key is missing
|
||||
func TestTavilyProviderWithoutAPIKey(t *testing.T) {
|
||||
// Create provider with nil config (no API key)
|
||||
provider := web.NewTavilyProvider(nil)
|
||||
require.NotNil(t, provider)
|
||||
|
||||
req := &types.Request{
|
||||
Query: "test query",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
|
||||
// Should not return error, but result should have error message
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
assert.Equal(t, types.SearchTypeWeb, result.Type)
|
||||
assert.Equal(t, "test query", result.Query)
|
||||
assert.NotEmpty(t, result.Error)
|
||||
assert.Contains(t, result.Error, "API key")
|
||||
assert.Empty(t, result.Items)
|
||||
assert.Equal(t, 0, result.Total)
|
||||
}
|
||||
|
||||
// TestTavilyProviderWithEmptyConfig tests provider with empty config
|
||||
func TestTavilyProviderWithEmptyConfig(t *testing.T) {
|
||||
// Create provider with empty config
|
||||
cfg := &types.WebConfig{}
|
||||
provider := web.NewTavilyProvider(cfg)
|
||||
require.NotNil(t, provider)
|
||||
|
||||
req := &types.Request{
|
||||
Query: "test query",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceUser,
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
assert.NotEmpty(t, result.Error)
|
||||
assert.Contains(t, result.Error, "API key")
|
||||
}
|
||||
|
||||
// TestTavilyProviderMaxResults tests that max_results from config is respected
|
||||
func TestTavilyProviderMaxResults(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping integration test")
|
||||
}
|
||||
|
||||
testutils.Prepare(t)
|
||||
defer testutils.Clean(t)
|
||||
|
||||
// Load the web-tavily test assistant
|
||||
ast, err := assistant.LoadPath("/assistants/tests/web-tavily")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, ast.Search)
|
||||
require.NotNil(t, ast.Search.Web)
|
||||
|
||||
// Create TavilyProvider
|
||||
provider := web.NewTavilyProvider(ast.Search.Web)
|
||||
|
||||
// Execute search without limit (should use config's max_results)
|
||||
req := &types.Request{
|
||||
Query: "artificial intelligence",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
// No Limit set, should use config's max_results (10)
|
||||
}
|
||||
|
||||
result, err := provider.Search(req)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result.Error, "Search should succeed with valid API key, got error: %s", result.Error)
|
||||
|
||||
// Should respect max_results from config
|
||||
assert.LessOrEqual(t, result.Total, ast.Search.Web.MaxResults)
|
||||
t.Logf("Search without limit returned %d results (max: %d)", result.Total, ast.Search.Web.MaxResults)
|
||||
|
||||
// Execute search with explicit limit
|
||||
req2 := &types.Request{
|
||||
Query: "artificial intelligence",
|
||||
Type: types.SearchTypeWeb,
|
||||
Source: types.SourceAuto,
|
||||
Limit: 3, // Override config's max_results
|
||||
}
|
||||
|
||||
result2, err := provider.Search(req2)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result2)
|
||||
|
||||
// API key must be valid - search should succeed
|
||||
require.Empty(t, result2.Error, "Search should succeed with valid API key, got error: %s", result2.Error)
|
||||
|
||||
// Should respect request's limit
|
||||
assert.LessOrEqual(t, result2.Total, 3)
|
||||
t.Logf("Search with limit=3 returned %d results", result2.Total)
|
||||
}
|
||||
|
|
@ -17,9 +17,10 @@ type Config struct {
|
|||
// Note: uses.web determines the mode (builtin/agent/mcp)
|
||||
// Provider is only used when uses.web = "builtin"
|
||||
type WebConfig struct {
|
||||
Provider string `json:"provider,omitempty" yaml:"provider,omitempty"` // "tavily" or "serper" (for builtin mode)
|
||||
Provider string `json:"provider,omitempty" yaml:"provider,omitempty"` // "tavily", "serper", or "serpapi" (for builtin mode)
|
||||
APIKeyEnv string `json:"api_key_env,omitempty" yaml:"api_key_env,omitempty"` // Environment variable for API key
|
||||
MaxResults int `json:"max_results,omitempty" yaml:"max_results,omitempty"` // Max results (default: 10)
|
||||
Engine string `json:"engine,omitempty" yaml:"engine,omitempty"` // Search engine for SerpAPI: "google", "bing", "baidu", "yandex", etc. (default: "google")
|
||||
}
|
||||
|
||||
// KBConfig for knowledge base search settings
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue